"""
Post-Merge Validation — run after EVERY merge step.

Tests:
1. Canary recall (did knowledge transfer?)
2. Perplexity check (did we break the model?)
3. Thinking mode (do <think> tags still work?)
4. Quick reasoning test (can it still think?)

Kill criteria: >10% performance drop on any test → abort merge.
Findings: #11, #22, #25
"""

import torch
import math
from transformers import AutoModelForCausalLM, AutoTokenizer

from .canary import test_all_canaries
from .config import MergeConfig


def validate_merged_model(
    model: AutoModelForCausalLM,
    tokenizer: AutoTokenizer,
    merged_sources: list[str],
    cfg: MergeConfig,
    baseline_perplexity: float = None,
) -> dict:
    """
    Run full validation suite on a merged model.

    Args:
        model: The merged model to validate
        tokenizer: The tokenizer
        merged_sources: List of source models merged so far
        cfg: Merge configuration
        baseline_perplexity: Perplexity of the target model before merging

    Returns:
        Dict with test results and overall pass/fail
    """
    print("\n" + "=" * 60)
    print(f"VALIDATION — After merging: {', '.join(merged_sources)}")
    print("=" * 60)

    results = {
        "canary": None,
        "perplexity": None,
        "thinking_mode": None,
        "reasoning": None,
        "overall": False,
    }

    # --- Test 1: Canary recall ---
    canary_results = test_all_canaries(model, tokenizer, merged_sources)
    passed_canaries = sum(1 for v in canary_results.values() if v)
    total_canaries = len(canary_results)
    results["canary"] = {
        "passed": passed_canaries,
        "total": total_canaries,
        "ok": passed_canaries >= cfg.canary_pass_threshold,
        "details": canary_results,
    }

    # --- Test 2: Perplexity ---
    perplexity = compute_perplexity(model, tokenizer)
    ppl_ok = True
    if baseline_perplexity is not None:
        ratio = perplexity / baseline_perplexity
        ppl_ok = ratio < cfg.perplexity_threshold
        print(f"\n[validate] Perplexity: {perplexity:.2f} (baseline: {baseline_perplexity:.2f}, ratio: {ratio:.2f})")
        if not ppl_ok:
            print(f"[validate] ⚠ Perplexity ratio {ratio:.2f} exceeds threshold {cfg.perplexity_threshold}")
    else:
        print(f"\n[validate] Perplexity: {perplexity:.2f} (no baseline to compare)")
    results["perplexity"] = {"value": perplexity, "ok": ppl_ok}

    # --- Test 3: Thinking mode ---
    think_ok = test_thinking_mode(model, tokenizer)
    results["thinking_mode"] = {"ok": think_ok}

    # --- Test 4: Quick reasoning ---
    reason_ok = test_reasoning(model, tokenizer)
    results["reasoning"] = {"ok": reason_ok}

    # --- Overall verdict ---
    all_ok = (
        results["canary"]["ok"]
        and results["perplexity"]["ok"]
        and results["thinking_mode"]["ok"]
        and results["reasoning"]["ok"]
    )
    results["overall"] = all_ok

    # Summary
    print("\n" + "-" * 60)
    print("VALIDATION SUMMARY")
    print("-" * 60)
    print(f"  Canary recall:   {'✓' if results['canary']['ok'] else '✗'} ({passed_canaries}/{total_canaries})")
    print(f"  Perplexity:      {'✓' if ppl_ok else '✗'} ({perplexity:.2f})")
    print(f"  Thinking mode:   {'✓' if think_ok else '✗'}")
    print(f"  Reasoning:       {'✓' if reason_ok else '✗'}")
    print(f"  OVERALL:         {'✓ PASS' if all_ok else '✗ FAIL — consider aborting'}")
    print("-" * 60)

    return results


def compute_perplexity(
    model: AutoModelForCausalLM,
    tokenizer: AutoTokenizer,
    test_texts: list[str] = None,
) -> float:
    """
    Compute perplexity on a small test set.

    Lower perplexity = model is more confident about predicting text.
    A big spike after merging means the model was damaged.
    """
    if test_texts is None:
        test_texts = [
            "The quick brown fox jumps over the lazy dog.",
            "In mathematics, a prime number is a natural number greater than 1.",
            "def fibonacci(n):\n    if n <= 1:\n        return n\n    return fibonacci(n-1) + fibonacci(n-2)",
            "The theory of general relativity describes gravity as the curvature of spacetime.",
            "To solve 3x + 7 = 22, subtract 7 from both sides to get 3x = 15, then divide by 3.",
        ]

    model.eval()
    total_loss = 0.0
    total_tokens = 0

    for text in test_texts:
        inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=512)
        inputs = {k: v.to(model.device) for k, v in inputs.items()}

        with torch.no_grad():
            outputs = model(**inputs, labels=inputs["input_ids"])
            total_loss += outputs.loss.item() * inputs["input_ids"].shape[1]
            total_tokens += inputs["input_ids"].shape[1]

    avg_loss = total_loss / total_tokens
    perplexity = math.exp(avg_loss)
    return perplexity


def test_thinking_mode(
    model: AutoModelForCausalLM,
    tokenizer: AutoTokenizer,
) -> bool:
    """
    Test if the model still uses <think> tags for reasoning.

    The thinking mode is Qwen3's special feature — if it's gone,
    the merge damaged something critical.
    """
    prompt = "Solve step by step: What is 15 × 13?"

    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=200,
            temperature=0.7,
            do_sample=True,
        )

    response = tokenizer.decode(outputs[0], skip_special_tokens=False)

    # Check for thinking tags
    has_think_open = "<think>" in response
    has_think_close = "</think>" in response
    passed = has_think_open and has_think_close

    print(f"\n[validate] Thinking mode test:")
    print(f"  Prompt:    {prompt}")
    print(f"  Response:  {response[:200]}...")
    print(f"  <think>:   {'✓ found' if has_think_open else '✗ missing'}")
    print(f"  </think>:  {'✓ found' if has_think_close else '✗ missing'}")
    print(f"  Status:    {'✓ PASS' if passed else '✗ FAIL'}")

    return passed


def test_reasoning(
    model: AutoModelForCausalLM,
    tokenizer: AutoTokenizer,
) -> bool:
    """
    Quick reasoning sanity check — can the model still do basic math?

    This catches catastrophic failures where the merge produced gibberish.
    """
    prompt = "What is 7 + 8?"
    expected_answer = "15"

    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=50,
            temperature=0.1,
            do_sample=False,
        )

    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    passed = expected_answer in response

    print(f"\n[validate] Quick reasoning test:")
    print(f"  Prompt:   {prompt}")
    print(f"  Expected: {expected_answer}")
    print(f"  Got:      {response}")
    print(f"  Status:   {'✓ PASS' if passed else '✗ FAIL'}")

    return passed