import tempfile
import unittest
from subprocess import CalledProcessError
from pathlib import Path
from unittest.mock import patch

from PIL import Image
import torch

from diffusion_cli.errors import DiffusionCliError
from diffusion_cli.image_io import (
    convertPngFile,
    encodeImages,
    outputPathWithExtension,
    outputPaths,
    saveImages,
)


class ImageIoTest(unittest.TestCase):
    def testOutputPathsAppendBatchIndexBeforeSuffix(self):
        paths = outputPaths(Path("output.png"), 2, "jpg")

        self.assertEqual(paths, [Path("output_0000.jpg"),
                                 Path("output_0001.jpg")])

    def testOutputPathWithExtensionReplacesOrAddsSuffix(self):
        self.assertEqual(
            outputPathWithExtension(Path("result.png"), "webp"),
            Path("result.webp"),
        )
        self.assertEqual(
            outputPathWithExtension(Path("result"), "jpg"),
            Path("result.jpg"),
        )

    def testSaveImagesWritesPngWithExpectedSize(self):
        images = torch.zeros(1, 3, 4, 5)

        with tempfile.TemporaryDirectory() as temp_dir:
            output = Path(temp_dir) / "image.png"
            paths = saveImages(images, output)
            with Image.open(paths[0]) as image:
                size = image.size
                mode = image.mode

        self.assertEqual(paths, [output])
        self.assertEqual(size, (5, 4))
        self.assertEqual(mode, "RGB")

    def testSaveImagesUsesFinalSuffixForNonPngOutput(self):
        images = torch.zeros(1, 3, 4, 5)

        with tempfile.TemporaryDirectory() as temp_dir:
            output = Path(temp_dir) / "image.png"

            def fake_run(args, **_kwargs):
                Path(args[-1]).write_bytes(b"jpg")

            with patch("diffusion_cli.image_io.subprocess.run", fake_run):
                paths = saveImages(images, output, "jpg", 85)

        self.assertEqual(paths, [Path(temp_dir) / "image.jpg"])

    def testConvertPngFileUsesImageMagickArgumentList(self):
        with tempfile.TemporaryDirectory() as temp_dir:
            root = Path(temp_dir)
            input_path = root / "input.png"
            output_path = root / "output.webp"
            input_path.write_bytes(b"png")

            def fake_run(args, **kwargs):
                self.assertEqual(
                    args,
                    [
                        "magick",
                        str(input_path),
                        "-quality",
                        "80",
                        str(output_path),
                    ],
                )
                self.assertTrue(kwargs["check"])
                self.assertTrue(kwargs["capture_output"])
                self.assertTrue(kwargs["text"])
                output_path.write_bytes(b"webp")

            with patch("diffusion_cli.image_io.subprocess.run", fake_run):
                convertPngFile(input_path, output_path, "webp", 80)

    def testConvertPngFileReportsMissingImageMagick(self):
        with tempfile.TemporaryDirectory() as temp_dir:
            input_path = Path(temp_dir) / "input.png"
            output_path = Path(temp_dir) / "output.avif"
            input_path.write_bytes(b"png")

            with patch(
                "diffusion_cli.image_io.subprocess.run",
                side_effect=FileNotFoundError(),
            ):
                with self.assertRaises(DiffusionCliError) as context:
                    convertPngFile(input_path, output_path, "avif", 95)

        self.assertIn(
            "ImageMagick executable not found",
            str(context.exception),
        )

    def testConvertPngFileReportsImageMagickFailureStderr(self):
        with tempfile.TemporaryDirectory() as temp_dir:
            input_path = Path(temp_dir) / "input.png"
            output_path = Path(temp_dir) / "output.avif"
            input_path.write_bytes(b"png")
            error = CalledProcessError(
                1,
                ["magick"],
                stderr="no encode delegate\n",
            )

            with patch(
                "diffusion_cli.image_io.subprocess.run",
                side_effect=error,
            ):
                with self.assertRaises(DiffusionCliError) as context:
                    convertPngFile(input_path, output_path, "avif", 95)

        self.assertIn(
            "ImageMagick failed to write avif output: no encode delegate",
            str(context.exception),
        )

    def testEncodeImagesReturnsPngBytes(self):
        images = torch.zeros(1, 3, 4, 5)

        encoded = encodeImages(images)

        self.assertEqual(len(encoded), 1)
        self.assertTrue(encoded[0].startswith(b"\x89PNG\r\n\x1a\n"))

    def testEncodeImagesReturnsConvertedBytesForNonPng(self):
        images = torch.zeros(1, 3, 4, 5)

        def fake_run(args, **_kwargs):
            Path(args[-1]).write_bytes(b"converted")

        with patch("diffusion_cli.image_io.subprocess.run", fake_run):
            encoded = encodeImages(images, "webp", 90)

        self.assertEqual(encoded, [b"converted"])


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