from helper import unittest, PillowTestCase, PillowLeakTestCase, hopper
from PIL import Image, ImageFile, PngImagePlugin
from PIL._util import py3

from io import BytesIO
import zlib
import sys

codecs = dir(Image.core)


# sample png stream

TEST_PNG_FILE = "Tests/images/hopper.png"

# stuff to create inline PNG images

MAGIC = PngImagePlugin._MAGIC


def chunk(cid, *data):
    test_file = BytesIO()
    PngImagePlugin.putchunk(*(test_file, cid) + data)
    return test_file.getvalue()


o32 = PngImagePlugin.o32

IHDR = chunk(b"IHDR", o32(1), o32(1), b'\x08\x02', b'\0\0\0')
IDAT = chunk(b"IDAT")
IEND = chunk(b"IEND")

HEAD = MAGIC + IHDR
TAIL = IDAT + IEND


def load(data):
    return Image.open(BytesIO(data))


def roundtrip(im, **options):
    out = BytesIO()
    im.save(out, "PNG", **options)
    out.seek(0)
    return Image.open(out)


class TestFilePng(PillowTestCase):

    def setUp(self):
        if "zip_encoder" not in codecs or "zip_decoder" not in codecs:
            self.skipTest("zip/deflate support not available")

    def get_chunks(self, filename):
        chunks = []
        with open(filename, "rb") as fp:
            fp.read(8)
            with PngImagePlugin.PngStream(fp) as png:
                while True:
                    cid, pos, length = png.read()
                    chunks.append(cid)
                    try:
                        s = png.call(cid, pos, length)
                    except EOFError:
                        break
                    png.crc(cid, s)
        return chunks

    def test_sanity(self):

        # internal version number
        self.assertRegex(Image.core.zlib_version, r"\d+\.\d+\.\d+(\.\d+)?$")

        test_file = self.tempfile("temp.png")

        hopper("RGB").save(test_file)

        im = Image.open(test_file)
        im.load()
        self.assertEqual(im.mode, "RGB")
        self.assertEqual(im.size, (128, 128))
        self.assertEqual(im.format, "PNG")

        hopper("1").save(test_file)
        im = Image.open(test_file)

        hopper("L").save(test_file)
        im = Image.open(test_file)

        hopper("P").save(test_file)
        im = Image.open(test_file)

        hopper("RGB").save(test_file)
        im = Image.open(test_file)

        hopper("I").save(test_file)
        im = Image.open(test_file)

    def test_invalid_file(self):
        invalid_file = "Tests/images/flower.jpg"

        self.assertRaises(SyntaxError,
                          PngImagePlugin.PngImageFile, invalid_file)

    def test_broken(self):
        # Check reading of totally broken files.  In this case, the test
        # file was checked into Subversion as a text file.

        test_file = "Tests/images/broken.png"
        self.assertRaises(IOError, Image.open, test_file)

    def test_bad_text(self):
        # Make sure PIL can read malformed tEXt chunks (@PIL152)

        im = load(HEAD + chunk(b'tEXt') + TAIL)
        self.assertEqual(im.info, {})

        im = load(HEAD + chunk(b'tEXt', b'spam') + TAIL)
        self.assertEqual(im.info, {'spam': ''})

        im = load(HEAD + chunk(b'tEXt', b'spam\0') + TAIL)
        self.assertEqual(im.info, {'spam': ''})

        im = load(HEAD + chunk(b'tEXt', b'spam\0egg') + TAIL)
        self.assertEqual(im.info, {'spam': 'egg'})

        im = load(HEAD + chunk(b'tEXt', b'spam\0egg\0') + TAIL)
        self.assertEqual(im.info,  {'spam': 'egg\x00'})

    def test_bad_ztxt(self):
        # Test reading malformed zTXt chunks (python-pillow/Pillow#318)

        im = load(HEAD + chunk(b'zTXt') + TAIL)
        self.assertEqual(im.info, {})

        im = load(HEAD + chunk(b'zTXt', b'spam') + TAIL)
        self.assertEqual(im.info, {'spam': ''})

        im = load(HEAD + chunk(b'zTXt', b'spam\0') + TAIL)
        self.assertEqual(im.info, {'spam': ''})

        im = load(HEAD + chunk(b'zTXt', b'spam\0\0') + TAIL)
        self.assertEqual(im.info, {'spam': ''})

        im = load(HEAD + chunk(
            b'zTXt', b'spam\0\0' + zlib.compress(b'egg')[:1]) + TAIL)
        self.assertEqual(im.info, {'spam': ''})

        im = load(
            HEAD + chunk(b'zTXt', b'spam\0\0' + zlib.compress(b'egg')) + TAIL)
        self.assertEqual(im.info,  {'spam': 'egg'})

    def test_bad_itxt(self):

        im = load(HEAD + chunk(b'iTXt') + TAIL)
        self.assertEqual(im.info, {})

        im = load(HEAD + chunk(b'iTXt', b'spam') + TAIL)
        self.assertEqual(im.info, {})

        im = load(HEAD + chunk(b'iTXt', b'spam\0') + TAIL)
        self.assertEqual(im.info, {})

        im = load(HEAD + chunk(b'iTXt', b'spam\0\x02') + TAIL)
        self.assertEqual(im.info, {})

        im = load(HEAD + chunk(b'iTXt', b'spam\0\0\0foo\0') + TAIL)
        self.assertEqual(im.info, {})

        im = load(HEAD + chunk(b'iTXt', b'spam\0\0\0en\0Spam\0egg') + TAIL)
        self.assertEqual(im.info, {"spam": "egg"})
        self.assertEqual(im.info["spam"].lang, "en")
        self.assertEqual(im.info["spam"].tkey, "Spam")

        im = load(HEAD + chunk(b'iTXt', b'spam\0\1\0en\0Spam\0' +
                               zlib.compress(b"egg")[:1]) + TAIL)
        self.assertEqual(im.info, {'spam': ''})

        im = load(HEAD + chunk(b'iTXt', b'spam\0\1\1en\0Spam\0' +
                               zlib.compress(b"egg")) + TAIL)
        self.assertEqual(im.info, {})

        im = load(HEAD + chunk(b'iTXt', b'spam\0\1\0en\0Spam\0' +
                               zlib.compress(b"egg")) + TAIL)
        self.assertEqual(im.info, {"spam": "egg"})
        self.assertEqual(im.info["spam"].lang, "en")
        self.assertEqual(im.info["spam"].tkey, "Spam")

    def test_interlace(self):

        test_file = "Tests/images/pil123p.png"
        im = Image.open(test_file)

        self.assert_image(im, "P", (162, 150))
        self.assertTrue(im.info.get("interlace"))

        im.load()

        test_file = "Tests/images/pil123rgba.png"
        im = Image.open(test_file)

        self.assert_image(im, "RGBA", (162, 150))
        self.assertTrue(im.info.get("interlace"))

        im.load()

    def test_load_transparent_p(self):
        test_file = "Tests/images/pil123p.png"
        im = Image.open(test_file)

        self.assert_image(im, "P", (162, 150))
        im = im.convert("RGBA")
        self.assert_image(im, "RGBA", (162, 150))

        # image has 124 unique alpha values
        self.assertEqual(len(im.getchannel('A').getcolors()), 124)

    def test_load_transparent_rgb(self):
        test_file = "Tests/images/rgb_trns.png"
        im = Image.open(test_file)
        self.assertEqual(im.info["transparency"], (0, 255, 52))

        self.assert_image(im, "RGB", (64, 64))
        im = im.convert("RGBA")
        self.assert_image(im, "RGBA", (64, 64))

        # image has 876 transparent pixels
        self.assertEqual(im.getchannel('A').getcolors()[0][0], 876)

    def test_save_p_transparent_palette(self):
        in_file = "Tests/images/pil123p.png"
        im = Image.open(in_file)

        # 'transparency' contains a byte string with the opacity for
        # each palette entry
        self.assertEqual(len(im.info["transparency"]), 256)

        test_file = self.tempfile("temp.png")
        im.save(test_file)

        # check if saved image contains same transparency
        im = Image.open(test_file)
        self.assertEqual(len(im.info["transparency"]), 256)

        self.assert_image(im, "P", (162, 150))
        im = im.convert("RGBA")
        self.assert_image(im, "RGBA", (162, 150))

        # image has 124 unique alpha values
        self.assertEqual(len(im.getchannel('A').getcolors()), 124)

    def test_save_p_single_transparency(self):
        in_file = "Tests/images/p_trns_single.png"
        im = Image.open(in_file)

        # pixel value 164 is full transparent
        self.assertEqual(im.info["transparency"], 164)
        self.assertEqual(im.getpixel((31, 31)), 164)

        test_file = self.tempfile("temp.png")
        im.save(test_file)

        # check if saved image contains same transparency
        im = Image.open(test_file)
        self.assertEqual(im.info["transparency"], 164)
        self.assertEqual(im.getpixel((31, 31)), 164)
        self.assert_image(im, "P", (64, 64))
        im = im.convert("RGBA")
        self.assert_image(im, "RGBA", (64, 64))

        self.assertEqual(im.getpixel((31, 31)), (0, 255, 52, 0))

        # image has 876 transparent pixels
        self.assertEqual(im.getchannel('A').getcolors()[0][0], 876)

    def test_save_p_transparent_black(self):
        # check if solid black image with full transparency
        # is supported (check for #1838)
        im = Image.new("RGBA", (10, 10), (0, 0, 0, 0))
        self.assertEqual(im.getcolors(), [(100, (0, 0, 0, 0))])

        im = im.convert("P")
        test_file = self.tempfile("temp.png")
        im.save(test_file)

        # check if saved image contains same transparency
        im = Image.open(test_file)
        self.assertEqual(len(im.info["transparency"]), 256)
        self.assert_image(im, "P", (10, 10))
        im = im.convert("RGBA")
        self.assert_image(im, "RGBA", (10, 10))
        self.assertEqual(im.getcolors(), [(100, (0, 0, 0, 0))])

    def test_save_l_transparency(self):
        # There are 559 transparent pixels in l_trns.png.
        num_transparent = 559

        in_file = "Tests/images/l_trns.png"
        im = Image.open(in_file)
        self.assertEqual(im.mode, "L")
        self.assertEqual(im.info["transparency"], 255)

        im_rgba = im.convert('RGBA')
        self.assertEqual(
            im_rgba.getchannel("A").getcolors()[0][0], num_transparent)

        test_file = self.tempfile("temp.png")
        im.save(test_file)

        test_im = Image.open(test_file)
        self.assertEqual(test_im.mode, "L")
        self.assertEqual(test_im.info["transparency"], 255)
        self.assert_image_equal(im, test_im)

        test_im_rgba = test_im.convert('RGBA')
        self.assertEqual(
            test_im_rgba.getchannel('A').getcolors()[0][0], num_transparent)

    def test_save_rgb_single_transparency(self):
        in_file = "Tests/images/caption_6_33_22.png"
        im = Image.open(in_file)

        test_file = self.tempfile("temp.png")
        im.save(test_file)

    def test_load_verify(self):
        # Check open/load/verify exception (@PIL150)

        im = Image.open(TEST_PNG_FILE)
        im.verify()

        im = Image.open(TEST_PNG_FILE)
        im.load()
        self.assertRaises(RuntimeError, im.verify)

    def test_verify_struct_error(self):
        # Check open/load/verify exception (#1755)

        # offsets to test, -10: breaks in i32() in read. (IOError)
        #                  -13: breaks in crc, txt chunk.
        #                  -14: malformed chunk

        for offset in (-10, -13, -14):
            with open(TEST_PNG_FILE, 'rb') as f:
                test_file = f.read()[:offset]