"""无需下载模型即可检查后端选择，模拟不同 GPU 和 CUDA 检测失败。"""

import ast
import contextlib
import io
from pathlib import Path
from types import SimpleNamespace
import unittest
from unittest.mock import Mock


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


class AttentionBackendTests(unittest.TestCase):
    def setUp(self):
        # 只加载待测函数，避免 app.py 导入时下载模型和启动 Gradio。
        tree = ast.parse(APP_PATH.read_text(encoding="utf-8"))
        names = {"get_gpu_summary", "supports_flash_attention_3", "set_attention_backend_safe", "init_app"}
        functions = [node for node in tree.body if isinstance(node, ast.FunctionDef) and node.name in names]
        self.cuda = SimpleNamespace(
            is_available=Mock(return_value=True),
            get_device_capability=Mock(return_value=(12, 0)),
            get_device_name=Mock(return_value="RTX PRO 6000 Blackwell"),
        )
        self.namespace = {"torch": SimpleNamespace(cuda=self.cuda), "List": list}
        exec(compile(ast.Module(body=functions, type_ignores=[]), str(APP_PATH), "exec"), self.namespace)
        self.transformer = SimpleNamespace(set_attention_backend=Mock(), layers=SimpleNamespace())

    def select(self, backend):
        with contextlib.redirect_stdout(io.StringIO()):
            return self.namespace["set_attention_backend_safe"](self.transformer, backend)

    def test_blackwell_never_selects_flash_attention_3(self):
        for capability in [(10, 0), (12, 0)]:
            for alias in ["flash_3", "_flash_3", "_flash_3_hub", "_flash_varlen_3", "_flash_3_varlen_hub"]:
                with self.subTest(capability=capability, backend=alias):
                    self.cuda.get_device_capability.return_value = capability
                    self.transformer.set_attention_backend.reset_mock()
                    self.assertEqual(self.select(alias), "native")
                    self.transformer.set_attention_backend.assert_called_once_with("native")
                    self.assertIn("Hopper", self.namespace["attention_backend_error"])

    def test_ampere_uses_native_for_flash_attention_3(self):
        self.cuda.get_device_capability.return_value = (8, 0)
        self.assertEqual(self.select("_flash_3_hub"), "native")
        self.transformer.set_attention_backend.assert_called_once_with("native")

    def test_hopper_preserves_requested_flash_attention_3(self):
        self.cuda.get_device_capability.return_value = (9, 0)
        self.assertEqual(self.select("_flash_3_hub"), "_flash_3_hub")
        self.transformer.set_attention_backend.assert_called_once_with("_flash_3_hub")
        self.assertIsNone(self.namespace["attention_backend_error"])

    def test_unknown_gpu_uses_native(self):
        self.cuda.get_device_capability.side_effect = RuntimeError("CUDA unavailable")
        self.assertEqual(self.select("_flash_3_hub"), "native")
        self.transformer.set_attention_backend.assert_called_once_with("native")

    def test_no_cuda_uses_native_without_querying_device(self):
        self.cuda.is_available.return_value = False
        self.assertEqual(self.select("_flash_3_hub"), "native")
        self.cuda.get_device_capability.assert_not_called()

    def test_native_failure_does_not_try_incompatible_kernels(self):
        self.transformer.set_attention_backend.side_effect = RuntimeError("native failed")
        with self.assertRaisesRegex(RuntimeError, "native failed"):
            self.select("_flash_3_hub")
        self.transformer.set_attention_backend.assert_called_once_with("native")

    def test_backend_setup_error_falls_back_to_native(self):
        def configure(name):
            if name != "native":
                raise ValueError("backend unavailable")

        self.transformer.set_attention_backend.side_effect = configure
        self.assertEqual(self.select("missing_backend"), "native")
        attempts = [call.args[0] for call in self.transformer.set_attention_backend.call_args_list]
        self.assertEqual(attempts, ["missing_backend", "_missing_backend", "native"])
        self.assertEqual(self.namespace["attention_backend_error"], "backend unavailable")

    def test_native_selection_clears_previous_error(self):
        self.select("_flash_3_hub")
        self.assertEqual(self.select("native"), "native")
        self.assertIsNone(self.namespace["attention_backend_error"])

    def initialize(self, *, lora=False, allow_lora=False):
        loader = Mock()
        self.namespace.update(
            pipe=SimpleNamespace(transformer=self.transformer),
            ensure_models_loaded=Mock(), ENABLE_COMPILE=False, ENABLE_AOTI=True,
            ENABLE_WARMUP=False, aoti_loaded=False, aoti_error=None,
            lora_loaded=lora, AOTI_ALLOW_LORA=allow_lora, AOTI_VARIANT="fa3",
            AOTI_REPO="zerogpu-aoti/Z-Image", ATTENTION_BACKEND="_flash_3_hub",
            spaces=SimpleNamespace(aoti_blocks_load=loader),
        )
        with contextlib.redirect_stdout(io.StringIO()):
            self.namespace["init_app"]()
        return loader

    def test_blackwell_skips_fa3_aoti(self):
        loader = self.initialize()
        loader.assert_not_called()
        self.assertFalse(self.namespace["aoti_loaded"])
        self.assertIn("Hopper", self.namespace["aoti_error"])

    def test_allow_lora_cannot_override_hardware_guard(self):
        loader = self.initialize(lora=True, allow_lora=True)
        loader.assert_not_called()
        self.assertFalse(self.namespace["aoti_loaded"])

    def test_hopper_can_load_fa3_aoti(self):
        self.cuda.get_device_capability.return_value = (9, 0)
        loader = self.initialize()
        loader.assert_called_once_with(self.transformer.layers, "zerogpu-aoti/Z-Image", variant="fa3")
        self.assertTrue(self.namespace["aoti_loaded"])

    def test_hopper_still_skips_aoti_with_lora(self):
        self.cuda.get_device_capability.return_value = (9, 0)
        loader = self.initialize(lora=True)
        loader.assert_not_called()
        self.assertIn("LoRA", self.namespace["aoti_error"])


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