import unittest
from types import SimpleNamespace

import torch

from diffusion_cli.text_encoder import (
    QWEN3_4B_LAYER_INDEX,
    QWEN3_4B_ROPE_THETA,
    ZImageTextEncoder,
    applyZImagePromptTemplate,
    buildQwen3_4BConfig,
    initializeQwenRotaryEmbedding,
    normalizeQwenStateDict,
)


class FakeTokenBatch(dict):
    def to(self, _device):
        return self


class FakeTokenizer:
    def __init__(self):
        self.last_text = None

    def __call__(self, text, **kwargs):
        self.last_text = text
        self.last_kwargs = kwargs
        return {
            "input_ids": torch.tensor([[10, 11, 12]]),
            "attention_mask": torch.tensor([[1, 1, 1]]),
        }


class FakeOutput:
    def __init__(self):
        self.hidden_states = (
            torch.full((1, 3, 4), 1.0),
            torch.full((1, 3, 4), 2.0),
            torch.full((1, 3, 4), 3.0),
        )


class FakeModel:
    def __init__(self):
        self.calls = []
        self.eval_called = False

    def eval(self):
        self.eval_called = True

    def __call__(self, **kwargs):
        self.calls.append(kwargs)
        self.grad_enabled = torch.is_grad_enabled()
        return FakeOutput()


class TextEncoderTest(unittest.TestCase):
    def testPromptTemplateMatchesZImageChatFormat(self):
        wrapped = applyZImagePromptTemplate("a mug")

        self.assertEqual(
            wrapped,
            "<|im_start|>user\na mug<|im_end|>\n"
            "<|im_start|>assistant\n",
        )

    def testStateDictNormalizationStripsComfyModelPrefix(self):
        normalized = normalizeQwenStateDict(
            {
                "model.embed_tokens.weight": "a",
                "layers.0.input_layernorm.weight": "b",
            }
        )

        self.assertEqual(normalized["embed_tokens.weight"], "a")
        self.assertEqual(normalized["layers.0.input_layernorm.weight"], "b")

    def testInitializeRotaryEmbeddingRebuildsMetaAllocatedBuffer(self):
        class FakeRotary:
            pass

        model = SimpleNamespace()
        model.config = buildQwen3_4BConfig()
        model.rotary_emb = FakeRotary()
        model.rotary_emb.inv_freq = torch.empty(64)
        model.rotary_emb.original_inv_freq = torch.empty(64)

        initializeQwenRotaryEmbedding(model)

        self.assertTrue(torch.all(torch.isfinite(model.rotary_emb.inv_freq)))
        self.assertEqual(model.rotary_emb.inv_freq[0].item(), 1.0)
        self.assertAlmostEqual(
            model.rotary_emb.inv_freq[-1].item(),
            1.0 / (QWEN3_4B_ROPE_THETA ** (126 / 128)),
        )
        self.assertTrue(torch.equal(
            model.rotary_emb.inv_freq,
            model.rotary_emb.original_inv_freq,
        ))

    def testQwenConfigMatchesComfyTextEncoderShape(self):
        config = buildQwen3_4BConfig()

        self.assertEqual(config.hidden_size, 2560)
        self.assertEqual(config.intermediate_size, 9728)
        self.assertEqual(config.num_hidden_layers, 36)
        self.assertEqual(config.num_attention_heads, 32)
        self.assertEqual(config.num_key_value_heads, 8)
        self.assertEqual(config.head_dim, 128)
        self.assertEqual(config.pad_token_id, 151643)

    def testEncodeUsesTemplateAndPenultimateHiddenState(self):
        tokenizer = FakeTokenizer()
        model = FakeModel()
        encoder = ZImageTextEncoder(
            model_path=None,
            tokenizer_path=None,
            device=torch.device("cpu"),
            dtype=torch.float32,
            model=model,
            tokenizer=tokenizer,
        )

        conditioning = encoder.encode("a mug")

        self.assertTrue(model.eval_called)
        self.assertEqual(
            tokenizer.last_text,
            "<|im_start|>user\na mug<|im_end|>\n"
            "<|im_start|>assistant\n",
        )
        self.assertFalse(model.grad_enabled)
        self.assertTrue(
            torch.equal(
                conditioning.hidden_states,
                FakeOutput().hidden_states[QWEN3_4B_LAYER_INDEX],
            )
        )
        self.assertEqual(conditioning.token_count, 3)

    def testEncodePromptsEncodesPositiveAndNegative(self):
        encoder = ZImageTextEncoder(
            model_path=None,
            tokenizer_path=None,
            device=torch.device("cpu"),
            dtype=torch.float32,
            model=FakeModel(),
            tokenizer=FakeTokenizer(),
        )

        conditioning = encoder.encodePrompts("a mug", "watermark")

        self.assertEqual(conditioning.positive.token_count, 3)
        self.assertEqual(conditioning.negative.token_count, 3)


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