import tempfile
import unittest
from pathlib import Path

import torch
from safetensors.torch import save_file

from diffusion_cli.errors import DiffusionCliError
from diffusion_cli.vae import (
    FLUX_SCALE_FACTOR,
    FLUX_SHIFT_FACTOR,
    VaeConfig,
    VaeDecoder,
    ZImageVae,
    detectVaeConfig,
    fluxLatentToVae,
    postprocessVaeOutput,
)


class FakeDecoder(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.last_latent = None

    def forward(self, latent):
        self.last_latent = latent.detach().clone()
        return torch.full(
            (latent.shape[0], 3, latent.shape[2] * 8, latent.shape[3] * 8),
            0.5,
            device=latent.device,
            dtype=latent.dtype,
        )


class VaeTest(unittest.TestCase):
    def testFluxLatentToVaeAppliesComfyScaling(self):
        latent = torch.tensor([0.0, FLUX_SCALE_FACTOR])

        output = fluxLatentToVae(latent)

        expected = torch.tensor([
            FLUX_SHIFT_FACTOR,
            1.0 + FLUX_SHIFT_FACTOR,
        ])
        self.assertTrue(torch.allclose(output, expected))

    def testPostprocessClampsToImageRange(self):
        image = torch.tensor([-3.0, -1.0, 0.0, 1.0, 3.0])

        output = postprocessVaeOutput(image)

        self.assertTrue(torch.equal(
            output,
            torch.tensor([0.0, 0.0, 0.5, 1.0, 1.0]),
        ))

    def testDetectVaeConfigRejectsUnknownKeys(self):
        with self.assertRaises(DiffusionCliError) as context:
            detectVaeConfig({})

        self.assertIn("Unsupported VAE state dict", str(context.exception))

    def testDecodeValidatesShapeAndRange(self):
        model = FakeDecoder()
        vae = ZImageVae(
            model_path=None,
            device=torch.device("cpu"),
            dtype=torch.float32,
            model=model,
        )
        latent = torch.zeros(1, 16, 2, 3)

        image = vae.decode(latent)

        self.assertEqual(image.shape, (1, 3, 16, 24))
        self.assertTrue(torch.all((image >= 0) & (image <= 1)))
        self.assertTrue(torch.allclose(
            model.last_latent,
            torch.full_like(latent, FLUX_SHIFT_FACTOR),
        ))

    def testLoadsDecoderStateDictFromSafetensors(self):
        config = VaeConfig(
            channels=32,
            channel_multipliers=(1,),
            num_res_blocks=0,
            latent_channels=16,
            out_channels=3,
            resolution=8,
        )
        decoder = VaeDecoder(config, dtype=torch.float32)
        state = {
            f"decoder.{key}": value.detach().clone()
            for key, value in decoder.state_dict().items()
        }

        with tempfile.TemporaryDirectory() as temp_dir:
            path = Path(temp_dir) / "tiny_vae.safetensors"
            save_file(state, path)
            vae = ZImageVae(
                path,
                torch.device("cpu"),
                torch.float32,
                model=None,
                decoder_config=config,
            )

        self.assertIsInstance(vae.model, VaeDecoder)


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