import ast
import json
import sys
import tempfile
import threading
import unittest
from concurrent.futures import ThreadPoolExecutor
from datetime import datetime, timezone
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import patch

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from generation_stats import GenerationStats
from jujo_loras import lora_selected


class TestStatistics(unittest.TestCase):
    def test_jst_date_parallel_counts_and_restore(self):
        with tempfile.TemporaryDirectory() as directory:
            counter = GenerationStats(directory)
            moment = datetime(2026, 10, 14, 15, 1, tzinfo=timezone.utc)
            with ThreadPoolExecutor(max_workers=8) as pool:
                list(pool.map(lambda _: counter.record("anima", "model-a", completed_at=moment), range(500)))
            counter.record("illustrious", "model-b", count=3, completed_at=moment)
            counter.flush()
            data = json.loads(counter.path.read_text(encoding="utf-8"))
            self.assertEqual(data["counts"]["2026-10-15"], {"anima": {"model-a": 500}, "illustrious": {"model-b": 3}})
            self.assertEqual(set(data), {"version", "timezone", "counts", "api_counts", "api_tracking_started_at"})
            restored = GenerationStats(directory)
            # A successful generation before background restore is merged.
            restored.record("anima", "model-a", completed_at=moment)
            restored.flush()
            self.assertEqual(restored.snapshot()["counts"]["2026-10-15"]["anima"]["model-a"], 501)
            restored.flush()
            self.assertEqual(restored.snapshot()["counts"]["2026-10-15"]["anima"]["model-a"], 501)

    def test_save_failure_retries_and_unchanged_does_not_write(self):
        with tempfile.TemporaryDirectory() as directory:
            counter = GenerationStats(directory)
            counter.record("anima", "a")
            with patch("generation_stats.os.replace", side_effect=OSError("offline")):
                with self.assertRaises(OSError):
                    counter.flush()
            counter.record("anima", "a")
            counter.flush()
            with patch("generation_stats.os.replace", side_effect=AssertionError("unnecessary write")):
                counter.flush()
            self.assertEqual(sum(day["anima"]["a"] for day in counter.snapshot()["counts"].values()), 2)

    def test_generation_record_does_not_wait_for_storage(self):
        with tempfile.TemporaryDirectory() as directory:
            counter = GenerationStats(directory)
            counter.record("anima", "a")
            entered, release, recorded = threading.Event(), threading.Event(), threading.Event()
            real_write = Path.write_text

            def slow_write(path, *args, **kwargs):
                entered.set()
                release.wait(3)
                return real_write(path, *args, **kwargs)

            with patch.object(Path, "write_text", slow_write):
                writer = threading.Thread(target=counter.flush)
                writer.start()
                self.assertTrue(entered.wait(1))
                recorder = threading.Thread(target=lambda: (counter.record("anima", "a"), recorded.set()))
                recorder.start()
                try:
                    self.assertTrue(recorded.wait(0.5), "generation blocked behind storage I/O")
                finally:
                    release.set(); writer.join(); recorder.join()
            counter.flush()
            self.assertEqual(sum(day["anima"]["a"] for day in counter.snapshot()["counts"].values()), 2)

    def test_corrupt_history_is_not_overwritten(self):
        with tempfile.TemporaryDirectory() as directory:
            counter = GenerationStats(directory)
            counter.path.write_text("broken", encoding="utf-8")
            counter.record("anima", "a")
            with self.assertRaises(ValueError):
                counter.flush()
            self.assertEqual(counter.path.read_text(), "broken")

    def test_background_save_and_memory_only_mode(self):
        with tempfile.TemporaryDirectory() as directory:
            counter = GenerationStats(directory, interval=0.02)
            saved = threading.Event()
            real_replace = __import__("os").replace
            def replace(*args):
                real_replace(*args); saved.set()
            with patch("generation_stats.os.replace", replace):
                counter.start()
                counter.record("illustrious", "a")
                try:
                    self.assertTrue(saved.wait(2))
                finally:
                    counter.stop()
            memory = GenerationStats()
            memory.record("anima", "a")
            memory.flush()
            self.assertIsNone(memory.path)

    def test_common_ui_mcp_hook_counts_only_success_after_gpu(self):
        source = ast.parse((Path(__file__).resolve().parents[1] / "app.py").read_text(encoding="utf-8"))
        function = next(n for n in source.body if isinstance(n, ast.FunctionDef) and n.name == "serialized_generate")
        actions = []
        def gpu(*args, **kwargs):
            actions.append("gpu")
            return object(), "info", "positive", "negative", {"seed":42}
        env = dict(
            generation_lock=threading.Lock(), MODEL_FAMILY={"a": "anima"}, MODEL_ALIASES={}, RETIRED_MODELS=set(),
            lora_selected=lora_selected, _generate_image_core=gpu,
            image_with_metadata=lambda pixels,meta: (actions.append("metadata") or pixels),
            generation_stats=SimpleNamespace(record=lambda *a, **kw: actions.append(("count", a, kw))),
        )
        exec(compile(ast.Module(body=[function], type_ignores=[]), "app.py", "exec"), env)
        args = ["a"] + [None] * 9 + [False, "---", .8] * 4 + [None] * 3
        env["serialized_generate"](*args)
        self.assertEqual(actions, ["gpu", "metadata", ("count", ("anima", "a"), {"source":"api"})])
        actions.clear()
        env["_generate_image_core"] = lambda *a, **kw: (_ for _ in ()).throw(RuntimeError("failed"))
        with self.assertRaises(RuntimeError):
            env["serialized_generate"](*args)
        self.assertEqual(actions, [])


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