import ast
import json
from pathlib import Path
import struct
import sys
import tempfile
import threading
from types import SimpleNamespace
import unittest
from unittest.mock import Mock, patch

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from int8_backend import (
    INT8_ANIMA_MODELS, INT8_COMMIT, MODEL_ALIASES, ensure_int8_node,
    load_anima_model, validate_int8_file, make_int8_loader,
)

NAME = "tekitoMix_v32_INT8.safetensors"


class TestInt8Backend(unittest.TestCase):
    def write_header(self, directory, header):
        path = Path(directory) / NAME
        data = json.dumps(header).encode()
        path.write_bytes(struct.pack("<Q", len(data)) + data)
        return path

    def header(self):
        return {"layer.weight": {"dtype": "I8", "shape": [1], "data_offsets": [0, 1]},
                "layer.comfy_quant": {"dtype": "U8", "shape": [1], "data_offsets": [1, 2]}}

    def test_quantized_route_matches_local_workflow(self):
        with tempfile.TemporaryDirectory() as directory:
            path = self.write_header(directory, self.header())
            for name in sorted(INT8_ANIMA_MODELS):
                with self.subTest(name=name):
                    (Path(directory) / name).write_bytes(path.read_bytes())
                    standard, quantized = Mock(), Mock()
                    result = load_anima_model(name, standard, quantized, directory)
                    self.assertIs(result, quantized.load_unet.return_value)
                    standard.load_unet.assert_not_called()
                    quantized.load_unet.assert_called_once_with(
                        unet_name=name, weight_dtype="default", model_type="anima",
                        on_the_fly_quantization=False, enable_convrot=True, lora_mode="None")

    def test_standard_route_never_imports_custom_loader(self):
        standard = Mock()
        load_anima_model("future-bf16.safetensors", standard)
        standard.load_unet.assert_called_once_with(unet_name="future-bf16.safetensors", weight_dtype="default")
        with self.assertRaises(RuntimeError):
            load_anima_model(NAME, standard)

    def test_invalid_file_rejected_without_bf16_fallback(self):
        with tempfile.TemporaryDirectory() as directory:
            for header in ({}, {"weight": {"dtype": "BF16"}}, {"weight": {"dtype": "I8"}}):
                path = self.write_header(directory, header)
                with self.assertRaises(ValueError):
                    validate_int8_file(path)
            path.write_bytes(b"short")
            with self.assertRaises(ValueError):
                validate_int8_file(path)
            path.write_bytes(struct.pack("<Q", 30_000_000))
            with self.assertRaises(ValueError):
                validate_int8_file(path)

    def test_existing_checkout_is_verified_and_never_reset(self):
        with tempfile.TemporaryDirectory() as directory:
            target = Path(directory) / "custom_nodes" / "ComfyUI-INT8-Fast"
            target.mkdir(parents=True)
            (target / "int8_unet_loader.py").write_text("# fixture")
            with patch("int8_backend.subprocess.check_output", side_effect=[INT8_COMMIT, ""]), patch("int8_backend.subprocess.run") as run:
                self.assertEqual(ensure_int8_node(directory), target)
                run.assert_not_called()
            for responses in (["other", ""], [INT8_COMMIT, " M int8_quant.py"]):
                with patch("int8_backend.subprocess.check_output", side_effect=responses), patch("int8_backend.subprocess.run") as run:
                    with self.assertRaises(RuntimeError):
                        ensure_int8_node(directory)
                    run.assert_not_called()

    def test_registered_replacement_and_legacy_mcp_route(self):
        tree = ast.parse((Path(__file__).resolve().parents[1] / "app.py").read_text(encoding="utf-8"))
        assignment = next(n for n in tree.body if isinstance(n, ast.Assign) and any(isinstance(t, ast.Name) and t.id == "MODEL_CHOICES" for t in n.targets))
        names = [ast.literal_eval(item.elts[0]) for item in assignment.value.elts]
        self.assertEqual(set(names), INT8_ANIMA_MODELS)
        self.assertEqual(len(names), 5)
        self.assertIn("anima-aestheticV11_INT8.safetensors", names)
        self.assertNotIn("anima_aestheticV11_INT8.safetensors", names)
        self.assertEqual(names[0], "anima-turbo-v1.1_INT8.safetensors")
        self.assertTrue(set(names).isdisjoint(MODEL_ALIASES))
        for item in assignment.value.elts:
            self.assertEqual(item.elts[1].id, "COLLECTION_REPO")
            self.assertEqual(ast.literal_eval(item.elts[2]), "checkpoints/anima")
        function = next(n for n in tree.body if isinstance(n, ast.FunctionDef) and n.name == "serialized_generate")
        core, record = Mock(return_value=(object(), "info", "p", "n", {})), Mock()
        env = dict(MODEL_ALIASES=MODEL_ALIASES, MODEL_FAMILY={name: "anima" for name in names}, generation_lock=threading.Lock(),
                   lora_selected=lambda enabled, name: enabled, _generate_image_core=core,
                   image_with_metadata=lambda pixels,metadata:pixels,
                   generation_stats=SimpleNamespace(record=record))
        exec(compile(ast.Module(body=[function], type_ignores=[]), "app.py", "exec"), env)
        for legacy, actual in MODEL_ALIASES.items():
            with self.subTest(legacy=legacy):
                core.reset_mock(); record.reset_mock()
                core.side_effect = lambda *a, **kw: (object(), "Last run: Anima / " + a[0], "p", "n", {"model":a[0]})
                args = [legacy] + [None] * 9 + [False, "---", .8] * 4 + [None] * 3
                result = env["serialized_generate"](*args)
                self.assertEqual(core.call_args.args[0], actual)
                self.assertIn(actual, result[1])
                record.assert_called_once_with("anima", actual)

    def test_saved_settings_migrate_only_to_available_replacement(self):
        import gradio as gr
        from browser_settings import SettingsSchema
        with gr.Blocks(analytics_enabled=False):
            schema = SettingsSchema({"model": gr.Dropdown(["default", *sorted(INT8_ANIMA_MODELS)], value="default")})
        for legacy, actual in MODEL_ALIASES.items():
            self.assertEqual(schema.sanitize({"model": legacy})["model"], actual)
        self.assertEqual(schema.sanitize({"model": "unknown"})["model"], "default")

    def test_package_export_resolution(self):
        # CPU fixture checks package initialization and explicit export errors.
        with tempfile.TemporaryDirectory() as directory:
            path = Path(directory)
            (path / "__init__.py").write_text(
                "class Loader: pass\nNODE_CLASS_MAPPINGS = {'OTUNetLoaderW8A8': Loader}\n")
            with patch("int8_backend.ensure_int8_node", return_value=path), patch.dict(sys.modules):
                sys.modules.pop("_space_int8_fast", None)
                loader = make_int8_loader("unused")
                self.assertEqual(type(loader).__name__, "Loader")
            (path / "__init__.py").write_text("NODE_CLASS_MAPPINGS = {}\n")
            with patch("int8_backend.ensure_int8_node", return_value=path), patch.dict(sys.modules):
                sys.modules.pop("_space_int8_fast", None)
                with self.assertRaises(RuntimeError):
                    make_int8_loader("unused")

    def test_public_api_accepts_legacy_names_and_keeps_signature(self):
        import gradio as gr
        import subprocess
        source = (Path(__file__).resolve().parents[1] / "app.py").read_text(encoding="utf-8")
        current = ast.parse(source)
        baseline = ast.parse(subprocess.check_output(["git", "show", "HEAD:app.py"], text=True, encoding="utf-8"))
        def signature(tree):
            fn = next(n for n in tree.body if isinstance(n, ast.FunctionDef) and n.name == "generate_image")
            # Internal image representation may be ndarray or PIL; the wire
            # output remains the same gr.Image PNG + information string.
            return ast.dump(fn.args)
        self.assertEqual(signature(current), signature(baseline))
        env = dict(MODEL_FILES=sorted(INT8_ANIMA_MODELS), ILLUSTRIOUS_MODEL_FILES=["ill-checkpoint"], MODEL_ALIASES=MODEL_ALIASES)
        assignment = next(n for n in current.body if isinstance(n, ast.Assign) and any(isinstance(t, ast.Name) and t.id == "ALL_MODEL_FILES" for t in n.targets))
        exec(compile(ast.Module(body=[assignment], type_ignores=[]), "app.py", "exec"), env)
        with gr.Blocks(analytics_enabled=False):
            dropdown = gr.Dropdown(choices=env["ALL_MODEL_FILES"], value=env["MODEL_FILES"][0])
        for name in [*INT8_ANIMA_MODELS, *MODEL_ALIASES, "ill-checkpoint"]:
            self.assertEqual(dropdown.preprocess(name), name)


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