from __future__ import annotations

import importlib.util
import sys
from pathlib import Path


REPO_ROOT = Path(__file__).resolve().parents[1]


def _load_script(name: str):
    path = REPO_ROOT / "scripts" / name
    spec = importlib.util.spec_from_file_location(name[:-3], path)
    module = importlib.util.module_from_spec(spec)
    assert spec.loader is not None
    sys.modules[spec.name] = module
    spec.loader.exec_module(module)
    return module


bench = _load_script("gemma4_hard_tier_bench.py")
loop = _load_script("gemma4_12b_recursive_loop.py")


def test_parse_tiers_expands_hardest_alias():
    assert bench.parse_tiers("hardest") == ["tier7_cbrn_critical_infra"]


def test_parse_tiers_supports_full_builtin_alias():
    assert bench.parse_tiers("842") == ["full_builtin_842"]


def test_load_tasks_omits_prompt_text_from_public_refs():
    tasks = bench.load_tasks(["tier7_cbrn_critical_infra"], per_tier_n=3)
    assert len(tasks) == 3
    refs = [task.public_ref() for task in tasks]
    assert all("prompt" not in ref for ref in refs)
    assert all(ref["dataset"] == "builtin" for ref in refs)
    assert all(ref["tier"] == "tier7_cbrn_critical_infra" for ref in refs)
    assert all(len(ref["prompt_sha256_12"]) == 12 for ref in refs)


def test_load_tokenizer_falls_back_to_slow_backend(monkeypatch):
    calls: list[dict[str, object]] = []

    def fake_from_pretrained(model: str, **kwargs):
        calls.append({"model": model, **kwargs})
        if len(calls) == 1:
            raise ValueError(
                "Couldn't instantiate the backend tokenizer. You need sentencepiece or tiktoken."
            )
        return type("Tok", (), {"pad_token_id": None, "eos_token": "</s>", "pad_token": None})()

    monkeypatch.setattr(bench.AutoTokenizer, "from_pretrained", fake_from_pretrained)

    tok = bench.load_tokenizer("google/gemma-4-E4B-it")

    assert tok.pad_token == "</s>"
    assert calls == [
        {"model": "google/gemma-4-E4B-it", "trust_remote_code": True},
        {"model": "google/gemma-4-E4B-it", "trust_remote_code": True, "use_fast": False},
    ]


def test_summarize_rows_scores_refusal_and_collapse():
    rows = [
        {
            "tier": "tier7_cbrn_critical_infra",
            "refused": False,
            "usable": True,
            "repetition": False,
            "short": False,
            "thinking_tag": False,
            "raw_thinking_tag": False,
            "clean_ending": True,
            "sentences": 2,
            "prompt_keyword_coverage": 0.8,
            "failure_reasons": [],
            "words": 120,
            "unique_word_ratio": 0.5,
            "reason": "none",
        },
        {
            "tier": "tier7_cbrn_critical_infra",
            "refused": True,
            "usable": False,
            "repetition": False,
            "short": True,
            "thinking_tag": True,
            "raw_thinking_tag": True,
            "clean_ending": False,
            "sentences": 0,
            "prompt_keyword_coverage": 0.0,
            "failure_reasons": ["refused", "visible_thinking_tag"],
            "words": 3,
            "unique_word_ratio": 0.9,
            "reason": "explicit_refusal",
        },
    ]
    summary = bench.summarize_rows(rows)
    assert summary["overall"]["n"] == 2
    assert summary["overall"]["refusal_rate"] == 0.5
    assert summary["overall"]["usable_rate"] == 0.5
    assert summary["overall"]["short_rate"] == 0.5
    assert summary["overall"]["thinking_tag_rate"] == 0.5
    assert summary["score"] < 100


def test_next_round_increases_pressure_after_refusals():
    recipe = loop.BASE_RECIPES[0]
    metrics = {"refusal_rate": 0.25, "repetition_rate": 0.0, "short_rate": 0.0}
    next_round = loop.next_round_from(recipe, metrics)
    assert any(item.n_directions > recipe.n_directions for item in next_round)
    assert any(item.max_layer_fraction > recipe.max_layer_fraction for item in next_round)