from __future__ import annotations

import unittest
from types import SimpleNamespace

import torch

from comfy_extras.nodes_minimax_h3 import _encode_unique_image


class _CountingVae:
    def __init__(self, mutate=False):
        self.calls = 0
        self.mutate = mutate

    def encode(self, x):
        self.calls += 1
        if self.mutate:
            x.add_(0.25)
        return x.mean().reshape(1, 1, 1, 1, 1).expand(1, 24, 1, 2, 2).clone()


class RefEncodeDedupTests(unittest.TestCase):
    def test_identical_refs_encode_once(self):
        vae = _CountingVae()
        img = torch.rand(1, 8, 8, 3)
        cache = []
        z0 = _encode_unique_image(vae, img, cache)
        z1 = _encode_unique_image(vae, img.clone(), cache)
        z2 = _encode_unique_image(vae, img.clone(), cache)
        self.assertEqual(vae.calls, 1)
        self.assertIs(z0, z1)
        self.assertIs(z1, z2)

    def test_different_refs_encode_each(self):
        vae = _CountingVae()
        cache = []
        a = torch.zeros(1, 4, 4, 3)
        b = torch.ones(1, 4, 4, 3)
        za = _encode_unique_image(vae, a, cache)
        zb = _encode_unique_image(vae, b, cache)
        self.assertEqual(vae.calls, 2)
        self.assertFalse(torch.equal(za, zb))

    def test_snapshot_survives_in_place_encode(self):
        vae = _CountingVae(mutate=True)
        cache = []
        img = torch.zeros(1, 4, 4, 3)
        _encode_unique_image(vae, img, cache)
        later = torch.zeros(1, 4, 4, 3)
        _encode_unique_image(vae, later, cache)
        self.assertEqual(vae.calls, 1)

    def test_shape_mismatch_does_not_hit(self):
        vae = _CountingVae()
        cache = []
        _encode_unique_image(vae, torch.zeros(1, 4, 4, 3), cache)
        _encode_unique_image(vae, torch.zeros(1, 8, 8, 3), cache)
        self.assertEqual(vae.calls, 2)


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