"""Check validation and GPU budgets without triggering checkpoint startup."""

import ast
import math
from pathlib import Path
import unittest


SOURCE = Path(__file__).resolve().parents[1] / "bol_space" / "runtime.py"
TREE = ast.parse(SOURCE.read_text())
# Runtime starts a model at import, so execute only its pure callbacks and sets.
DEFINITIONS = [
    node for node in TREE.body
    if (isinstance(node, ast.FunctionDef) and node.name in {"validate_options", "duration_for_request"})
    or (isinstance(node, ast.Assign) and any(
        isinstance(target, ast.Name) and target.id.endswith("_TOOLS")
        for target in node.targets
    ))
]
DEFAULT_TICKET = {"repo_id": "HuggingFaceTB/SmolLM2-1.7B-Instruct", "revision": "default-commit"}
NAMESPACE = {"math": math, "DEFAULT_TICKET": DEFAULT_TICKET}
exec(compile(ast.Module(body=DEFINITIONS, type_ignores=[]), str(SOURCE), "exec"), NAMESPACE)
validate_options = NAMESPACE["validate_options"]
duration_for_request = NAMESPACE["duration_for_request"]


def options(text):
    return {"text": text, "max_tokens": 128, "max_new_tokens": 16,
            "sample_layers": 8, "layer_start": 0, "layer_end": -1, "seed": 42}


class TextRequirementsTest(unittest.TestCase):
    def test_tools_without_an_evaluation_text_requirement_accept_blank_text(self):
        for tool in ("fingerprint", "silhouette", "semantic", "cka", "tugofwar"):
            for text in ("", " \n "):
                with self.subTest(tool=tool, text=text):
                    self.assertEqual(validate_options(tool, options(text))["text"], text)

    def test_text_experiments_still_reject_blank_evaluation_text(self):
        for tool in ("attention", "gradient", "butterfly", "layer_sweep", "component",
                     "trisect", "wormhole", "noise", "duets"):
            for text in ("", " \n "):
                with self.subTest(tool=tool, text=text):
                    with self.assertRaisesRegex(ValueError, "Enter some text"):
                        validate_options(tool, options(text))

    def test_optional_text_still_has_type_and_length_limits(self):
        for text in (None, 123, "x" * 10001):
            with self.subTest(text_type=type(text).__name__):
                with self.assertRaisesRegex(ValueError, "text must contain"):
                    validate_options("tugofwar", options(text))


class DurationBudgetTest(unittest.TestCase):
    def test_preloaded_checkpoint_uses_measured_short_budgets(self):
        expected = {"fingerprint": 30, "gradient": 30, "duets": 30, "component": 30,
                    "trisect": 30, "wormhole": 30, "noise": 30, "layer_sweep": 30,
                    "silhouette": 15, "semantic": 15, "cka": 15, "attention": 15,
                    "butterfly": 15, "tugofwar": 15}
        for tool, seconds in expected.items():
            with self.subTest(tool=tool):
                self.assertEqual(duration_for_request({"ticket": DEFAULT_TICKET, "tool": tool}), seconds)

    def test_custom_checkpoint_keeps_loading_budget_even_if_one_identifier_matches(self):
        tickets = [
            {"repo_id": "another/model", "revision": "default-commit"},
            {"repo_id": DEFAULT_TICKET["repo_id"], "revision": "another-commit"},
            {},
        ]
        for ticket in tickets:
            for tool, seconds in (("gradient", 150), ("silhouette", 120)):
                with self.subTest(ticket=ticket, tool=tool):
                    self.assertEqual(duration_for_request({"ticket": ticket, "tool": tool}), seconds)

    def test_missing_request_uses_short_validation_budget(self):
        self.assertEqual(duration_for_request(None), 15)
        self.assertEqual(duration_for_request({}), 15)


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