#!/usr/bin/env python3
from __future__ import annotations

import argparse
import json
import re
import time
from pathlib import Path

import torch
from transformers import AutoConfig, AutoTokenizer
from transformers import Qwen3_5ForConditionalGeneration
from transformers import Qwen3_5MoeForCausalLM
from transformers import Qwen3_5MoeForConditionalGeneration


def load_model(path: Path, device: str):
    config = AutoConfig.from_pretrained(path)
    if config.model_type == "qwen3_5_moe_text":
        cls = Qwen3_5MoeForCausalLM
    elif config.model_type == "qwen3_5_moe":
        cls = Qwen3_5MoeForConditionalGeneration
    else:
        cls = Qwen3_5ForConditionalGeneration
    return cls.from_pretrained(path, dtype=torch.bfloat16).to(device).eval()


def normalize(value: str) -> str:
    subscripts = str.maketrans("₀₁₂₃₄₅₆₇₈₉", "0123456789")
    return " ".join(value.casefold().translate(subscripts).split())


def score(item: dict, completion: str) -> tuple[bool, list[str]]:
    text = completion.strip()
    failures = []
    if len(text) < 2:
        failures.append("empty_or_too_short")
    if text and sum(character.isspace() for character in text) / len(text) > 0.5:
        failures.append("whitespace_dominated")
    if re.search(r"(.)\1{7,}", text, flags=re.DOTALL):
        failures.append("character_repetition")
    words = re.findall(r"\w+", text.casefold())
    if len(words) >= 12 and len(set(words)) < 4:
        failures.append("word_repetition")
    expected = item.get("expected_any")
    if expected and not any(normalize(value) in normalize(text) for value in expected):
        failures.append("expected_answer_missing")
    return not failures, failures


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--model", type=Path, required=True)
    parser.add_argument("--prompts", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--device", default="mps")
    parser.add_argument("--max-new-tokens", type=int, default=48)
    args = parser.parse_args()

    items = [json.loads(line) for line in args.prompts.read_text().splitlines() if line]
    tokenizer = AutoTokenizer.from_pretrained(args.model)
    model = load_model(args.model, args.device)
    results = []
    for index, item in enumerate(items, 1):
        batch = tokenizer.apply_chat_template(
            [{"role": "user", "content": item["prompt"]}],
            tokenize=True,
            add_generation_prompt=True,
            enable_thinking=False,
            return_tensors="pt",
            return_dict=True,
        ).to(args.device)
        started = time.perf_counter()
        with torch.no_grad():
            output = model.generate(
                **batch,
                max_new_tokens=args.max_new_tokens,
                do_sample=False,
                use_cache=True,
            )
        elapsed = time.perf_counter() - started
        ids = output[0, batch["input_ids"].shape[1]:]
        completion = tokenizer.decode(ids, skip_special_tokens=True)
        passed, failures = score(item, completion)
        result = {
            **item,
            "completion": completion,
            "passed": passed,
            "failures": failures,
            "new_tokens": int(ids.numel()),
            "elapsed_seconds": elapsed,
        }
        results.append(result)
        print(f"[{index:02d}/{len(items)}] {item['id']} {'PASS' if passed else 'FAIL'}", flush=True)

    by_language = {}
    for language in sorted({item["language"] for item in items}):
        subset = [row for row in results if row["language"] == language]
        by_language[language] = {
            "passed": sum(row["passed"] for row in subset),
            "total": len(subset),
            "pass_rate": sum(row["passed"] for row in subset) / len(subset),
        }
    passed = sum(row["passed"] for row in results)
    report = {
        "model": str(args.model),
        "total": len(results),
        "passed": passed,
        "pass_rate": passed / len(results),
        "gate_threshold": 0.90,
        "gate_passed": passed / len(results) >= 0.90,
        "by_language": by_language,
        "results": results,
    }
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(report, ensure_ascii=False, indent=2))
    print(json.dumps({key: value for key, value in report.items() if key != "results"}, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()