import json
import os
from pathlib import Path
from typing import Dict, Iterable, List, Tuple


def _read_jsonl(path: Path) -> List[Dict[str, str]]:
    rows: List[Dict[str, str]] = []
    with path.open("r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            rows.append(json.loads(line))
    return rows


def _normalize(s: str) -> str:
    return " ".join((s or "").strip().lower().split())


def _token_f1(pred: str, gold: str) -> float:
    """
    Lightweight similarity: token-level F1 (good enough for Buffett-style free text).
    """
    p = _normalize(pred).split()
    g = _normalize(gold).split()
    if not p and not g:
        return 1.0
    if not p or not g:
        return 0.0
    from collections import Counter

    pc = Counter(p)
    gc = Counter(g)
    common = pc & gc
    tp = sum(common.values())
    precision = tp / max(1, len(p))
    recall = tp / max(1, len(g))
    if precision + recall == 0:
        return 0.0
    return 2 * precision * recall / (precision + recall)


def _exact_match(pred: str, gold: str) -> float:
    return 1.0 if _normalize(pred) == _normalize(gold) else 0.0


def _iter_examples(rows: List[Dict[str, str]]) -> Iterable[Tuple[str, str]]:
    for obj in rows:
        # each line is a single-key dict: {prompt: answer}
        if not obj:
            continue
        [(prompt, answer)] = list(obj.items())
        yield prompt, answer


def main():
    """
    Evaluate adapter vs base on the bundled example sets.

    Requirements:
      - GPU recommended
      - HF_TOKEN set if BASE is gated (Llama 3.1)

    Usage:
      HF_TOKEN=... python scripts/eval_demo.py
    """
    # Import after env is set.
    import extract

    root = Path(__file__).resolve().parents[1]
    exdir = root / "example_data"

    # Map example file -> adapter key (matches app.py hidden "Model" input)
    suites = [
        ("buffett_example.jsonl", "buffett", _token_f1),
        ("buffett_ma_example.jsonl", "buffett_ma", _token_f1),
        ("ner_example.jsonl", "ner", _exact_match),
        ("xbrl_term_example.jsonl", "xbrl_term", _exact_match),
    ]

    # Allow small/quick eval
    limit = int(os.getenv("EVAL_LIMIT", "25"))

    print(f"Base model: {extract.BASE_MODEL_ID}")
    print(f"Limit per suite: {limit}")
    print("")

    for filename, adapter_key, metric in suites:
        path = exdir / filename
        if not path.exists():
            print(f"[skip] missing {filename}")
            continue

        rows = _read_jsonl(path)
        examples = list(_iter_examples(rows))[:limit]

        base_scores: List[float] = []
        ft_scores: List[float] = []

        for prompt, gold in examples:
            base = extract._generate(prompt, adapter_key=None, max_new_tokens=220)
            ft = extract._generate(prompt, adapter_key=adapter_key, max_new_tokens=220)
            base_scores.append(metric(base, gold))
            ft_scores.append(metric(ft, gold))

        base_avg = sum(base_scores) / max(1, len(base_scores))
        ft_avg = sum(ft_scores) / max(1, len(ft_scores))

        print(f"{filename}  (adapter={adapter_key})")
        print(f"  base avg: {base_avg:.3f}")
        print(f"  ft   avg: {ft_avg:.3f}")
        print(f"  delta  : {ft_avg - base_avg:+.3f}")
        print("")


if __name__ == "__main__":
    main()