from __future__ import annotations

import ast
import unittest
from pathlib import Path


APP_PATH = Path(__file__).resolve().parents[1] / "app_hyshape.py"


def _load_tree() -> ast.Module:
    return ast.parse(APP_PATH.read_text(encoding="utf-8"))


def _get_function(tree: ast.Module, name: str) -> ast.FunctionDef:
    for node in tree.body:
        if isinstance(node, ast.FunctionDef) and node.name == name:
            return node
    raise AssertionError(f"Function {name!r} not found in app_hyshape.py")


def _called_names(function_node: ast.FunctionDef) -> set[str]:
    names: set[str] = set()
    for node in ast.walk(function_node):
        if isinstance(node, ast.Call):
            if isinstance(node.func, ast.Name):
                names.add(node.func.id)
            elif isinstance(node.func, ast.Attribute):
                names.add(node.func.attr)
    return names


class AppHyShapeArchitectureTests(unittest.TestCase):
    def test_hyshape_app_does_not_reference_near_or_gsplat_paths(self) -> None:
        source = APP_PATH.read_text(encoding="utf-8")

        self.assertNotIn("NeARImageToRelightable3DPipeline", source)
        self.assertNotIn("ensure_near_pipeline", source)
        self.assertNotIn("ensure_gsplat_ready", source)

    def test_generate_mesh_only_uses_geometry_loader(self) -> None:
        generate_mesh = _get_function(_load_tree(), "generate_mesh")
        called = _called_names(generate_mesh)

        self.assertIn("ensure_geometry_on_cuda", called)
        self.assertNotIn("ensure_near_pipeline", called)
        self.assertNotIn("ensure_gsplat_ready", called)

    def test_generate_mesh_contains_early_callback_log_marker(self) -> None:
        source = APP_PATH.read_text(encoding="utf-8")

        self.assertIn("[HyShape] generate_mesh callback entered", source)

    def test_cpu_preload_worker_only_loads_cpu_locked_path(self) -> None:
        tree = _load_tree()
        worker = _get_function(tree, "preload_geometry_cpu_worker")
        called = _called_names(worker)

        self.assertIn("_ensure_geometry_loaded_on_cpu_locked", called)
        self.assertNotIn("ensure_near_pipeline", called)
        self.assertNotIn("ensure_gsplat_ready", called)

        source = APP_PATH.read_text(encoding="utf-8")
        self.assertIn("start_geometry_cpu_preload_thread", source)
        self.assertNotIn("warmup_hunyuan_geometry_on_load", source)

    def test_ensure_geometry_on_cuda_moves_to_gpu_not_near(self) -> None:
        ensure = _get_function(_load_tree(), "ensure_geometry_on_cuda")
        called = _called_names(ensure)

        self.assertIn("_ensure_geometry_loaded_on_cpu_locked", called)
        self.assertNotIn("ensure_near_pipeline", called)


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