from helper import unittest, PillowTestCase, lena

from io import BytesIO

from PIL import Image
from PIL import PngImagePlugin
import zlib

codecs = dir(Image.core)

# sample png stream

file = "Tests/images/lena.png"
data = open(file, "rb").read()

# stuff to create inline PNG images

MAGIC = PngImagePlugin._MAGIC


def chunk(cid, *data):
    file = BytesIO()
    PngImagePlugin.putchunk(*(file, cid) + data)
    return 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 test_sanity(self):

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

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

        lena("RGB").save(file)

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

        lena("1").save(file)
        im = Image.open(file)

        lena("L").save(file)
        im = Image.open(file)

        lena("P").save(file)
        im = Image.open(file)

        lena("RGB").save(file)
        im = Image.open(file)

        lena("I").save(file)
        im = Image.open(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.

        file = "Tests/images/broken.png"
        self.assertRaises(IOError, lambda: Image.open(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, {})

        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):

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

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

        im.load()

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

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

        im.load()

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

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

        # image has 124 uniqe qlpha values
        self.assertEqual(len(im.split()[3].getcolors()), 124)

    def test_load_transparent_rgb(self):
        file = "Tests/images/rgb_trns.png"
        im = Image.open(file)

        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.split()[3].getcolors()[0][0], 876)

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

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

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

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

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

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

        # There are 559 transparent pixels.
        im = im.convert('RGBA')
        self.assertEqual(im.split()[3].getcolors()[0][0], 559)

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

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

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

        im = Image.open("Tests/images/lena.png")
        im.verify()

        im = Image.open("Tests/images/lena.png")
        im.load()
        self.assertRaises(RuntimeError, lambda: im.verify())

    def test_roundtrip_dpi(self):
        # Check dpi roundtripping

        im = Image.open(file)

        im = roundtrip(im, dpi=(100, 100))
        self.assertEqual(im.info["dpi"], (100, 100))

    def test_roundtrip_text(self):
        # Check text roundtripping

        im = Image.open(file)

        info = PngImagePlugin.PngInfo()
        info.add_text("TXT", "VALUE")
        info.add_text("ZIP", "VALUE", 1)

        im = roundtrip(im, pnginfo=info)
        self.assertEqual(im.info, {'TXT': 'VALUE', 'ZIP': 'VALUE'})
        self.assertEqual(im.text, {'TXT': 'VALUE', 'ZIP': 'VALUE'})

    def test_scary(self):
        # Check reading of evil PNG file.  For information, see:
        # http://scary.beasts.org/security/CESA-2004-001.txt
        # The first byte is removed from pngtest_bad.png
        # to avoid classification as malware.

        with open("Tests/images/pngtest_bad.png.bin", 'rb') as fd:
            data = b'\x89' + fd.read()

        pngfile = BytesIO(data)
        self.assertRaises(IOError, lambda: Image.open(pngfile))

    def test_trns_rgb(self):
        # Check writing and reading of tRNS chunks for RGB images.
        # Independent file sample provided by Sebastian Spaeth.

        file = "Tests/images/caption_6_33_22.png"
        im = Image.open(file)
        self.assertEqual(im.info["transparency"], (248, 248, 248))

        # check saving transparency by default
        im = roundtrip(im)
        self.assertEqual(im.info["transparency"], (248, 248, 248))

        im = roundtrip(im, transparency=(0, 1, 2))
        self.assertEqual(im.info["transparency"], (0, 1, 2))

    def test_trns_p(self):
        # Check writing a transparency of 0, issue #528
        im = lena('P')
        im.info['transparency'] = 0

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

        im2 = Image.open(f)
        self.assertIn('transparency', im2.info)

        self.assert_image_equal(im2.convert('RGBA'), im.convert('RGBA'))

    def test_save_icc_profile_none(self):
        # check saving files with an ICC profile set to None (omit profile)
        in_file = "Tests/images/icc_profile_none.png"
        im = Image.open(in_file)
        self.assertEqual(im.info['icc_profile'], None)

        im = roundtrip(im)
        self.assertNotIn('icc_profile', im.info)

    def test_roundtrip_icc_profile(self):
        # check that we can roundtrip the icc profile
        im = lena('RGB')

        jpeg_image = Image.open('Tests/images/flower2.jpg')
        expected_icc = jpeg_image.info['icc_profile']

        im.info['icc_profile'] = expected_icc
        im = roundtrip(im)
        self.assertEqual(im.info['icc_profile'], expected_icc)


if __name__ == '__main__':
    unittest.main()

# End of file