"""SemanticWiki-Eval v3: format + citation validity + LLM-judge fidelity.

Held-out protocol:
  1. Discover candidate eval repos via GitHub search (same quota logic as the
     dataset run), EXCLUDE anything already in the dataset registry.
  2. Harvest + build context with the exact dataset pipeline (sw_pipeline).
  3. Generate wiki pages with each student model (merged LoRA -> vLLM).
  4. Score every page on three axes:
     a. FORMAT  - DeepWiki structure checks (h1/h2/h3, mermaid, sources, table)
     b. CITATION VALIDITY - deterministic verifier against the repo snapshot
     c. FIDELITY - LLM judge (Qwen3-Coder-30B-A3B) scores factual support
        against the provided context, 1-5, with unsupported-claim lists
  5. Upload per-page results and summary to GhostScientist/semanticwiki-eval-v3.

Run: python sw_eval.py [--n-repos 15] [--cats architecture_overview,data_flow]
"""

import argparse
import json
import os
import re
import subprocess
import sys
import time
import urllib.parse
import urllib.request
from datetime import datetime, timezone
from pathlib import Path

OUT_DIR = Path("/work/out")
OUT_DIR.mkdir(parents=True, exist_ok=True)
os.environ.setdefault("VLLM_USE_FLASHINFER_SAMPLER", "0")

SCRIPTS_REPO = "GhostScientist/semanticwiki-dataset-v3-scripts"
DATASET_REPO = "GhostScientist/semanticwiki-data-v3"
RESULTS_REPO = "GhostScientist/semanticwiki-eval-v3"

STUDENTS = [
    ("sw-coder-14b-v3", "GhostScientist/semanticwiki-coder-14b-v3",
     "Qwen/Qwen2.5-Coder-14B-Instruct"),
    ("sw-qwen3-14b-v3", "GhostScientist/semanticwiki-14b-v3-qwen3",
     "Qwen/Qwen3-14B"),
]
BASELINES = [  # un-finetuned bases, same prompts/judge - the ablation arm
    ("base-qwen25coder-14b", None, "Qwen/Qwen2.5-Coder-14B-Instruct"),
    ("base-qwen3-14b", None, "Qwen/Qwen3-14B"),
]
JUDGE = "Qwen/Qwen3-Coder-30B-A3B-Instruct"

SYSTEM_PROMPT = (
    "You are an expert software architect writing wiki documentation for codebases.\n"
    "Requirements:\n"
    "1. Every factual claim about the code must carry a source citation in the form "
    "`path:line` or `path:start-end` referencing the provided source files.\n"
    "2. Include at least one Mermaid diagram (flowchart or sequenceDiagram) when it "
    "clarifies structure or flow.\n"
    "3. Organize with H1 title and H2/H3 sections: overview, detailed analysis, and a "
    "conclusion.\n"
    "4. Explain the WHY behind design decisions, not only the WHAT.\n"
    "5. Never invent files, symbols, or line numbers - cite only what the context shows.\n"
    "6. Citation format is strict. Every citation must be the token "
    "`path/file.ext:line` or `path/file.ext:start-end` - for example "
    "`src/click/core.py:370-545`.\n"
    "7. Example: 'The `Context` class (`src/click/core.py:238-987`) manages state.'\n"
)

CATEGORIES = {
    "architecture_overview":
        "Generate a wiki page describing the overall architecture of this codebase: the "
        "main components, how they interact, entry points, and the public interface. "
        "Include a Mermaid flowchart of the component structure.",
    "data_flow":
        "Generate a wiki page tracing the primary data or control flow through this system "
        "from entry to exit. Include a Mermaid sequenceDiagram showing the interactions "
        "between components.",
}

JUDGE_SYSTEM = (
    "You are a strict technical reviewer. You are given the source code context of a "
    "repository and a generated wiki page about it. Judge ONLY factual fidelity to the "
    "provided context - not style. Respond with ONLY a JSON object: "
    '{"fidelity_score": <int 1-5>, "supported_claims": <int>, '
    '"unsupported_claims": <int>, "notes": "<one sentence>"} '
    "where fidelity_score means: 5 = every claim is verifiably grounded in the context, "
    "1 = most claims are invented or contradicted."
)


def ensure_pipeline():
    try:
        import sw_pipeline
        return sys.modules["sw_pipeline"]
    except ImportError:
        pass
    from huggingface_hub import hf_hub_download
    path = hf_hub_download(SCRIPTS_REPO, "sw_pipeline.py")
    sys.path.insert(0, str(Path(path).parent))
    import sw_pipeline
    return sw_pipeline


def discover_eval_repos(n_repos, registry_slugs, quota=None):
    """Same search recipe as the dataset run, minus repos used for training."""
    repos, seen = [], set(registry_slugs)
    if quota is None:
        quota = [("python", 6), ("typescript", 4), ("go", 3), ("rust", 2), ("java", 2)]
    page_by_lang = {lang: 1 for lang, _ in quota}
    idx = 0
    while len(repos) < n_repos and idx < 200:
        lang, _ = quota[idx % len(quota)]
        idx += 1
        page = page_by_lang[lang]
        if page > 4:
            continue
        page_by_lang[lang] = page + 1
        q = (f"language:{lang} stars:800..100000 size:500..30000 "
             f"license:mit license:apache-2.0 license:bsd-3-clause")
        url = ("https://api.github.com/search/repositories?q=" + urllib.parse.quote(q)
               + f"&sort=stars&order=desc&per_page=30&page={page}")
        req = urllib.request.Request(url, headers={"User-Agent": "semanticwiki-eval"})
        try:
            with urllib.request.urlopen(req, timeout=30) as r:
                data = json.loads(r.read())
        except Exception as e:
            print(f"[eval] discovery {lang} p{page}: {e}")
            time.sleep(10)
            continue
        for item in data.get("items", []):
            slug = item["full_name"]
            if slug in seen or item.get("fork") or item.get("archived"):
                continue
            seen.add(slug)
            repos.append({"org": item["owner"]["login"], "name": item["name"],
                          "language": lang,
                          "license": (item.get("license") or {}).get("spdx_id", "unknown")})
            if len(repos) >= n_repos:
                break
        time.sleep(8)
    print(f"[eval] {len(repos)} held-out repos (excluded {len(registry_slugs)} dataset repos)")
    return repos


def format_checks(text):
    """DeepWiki structural format score (0-1)."""
    checks = {
        "has_h1": bool(re.search(r"^# ", text, re.M)),
        "has_h2": bool(re.search(r"^## ", text, re.M)),
        "has_h3": bool(re.search(r"^### ", text, re.M)),
        "has_mermaid": "mermaid" in text,
        "has_citations": bool(re.search(r"[\w./\-]+\.[A-Za-z0-9]{1,4}:\d+", text)),
        "has_table": "|---" in text or "| ---" in text,
        "has_conclusion": bool(re.search(r"## .*conclusion", text, re.I)),
        "word_count_ok": 300 <= len(text.split()) <= 4000,
    }
    return {"passed": sum(checks.values()), "total": len(checks), **checks}


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--n-repos", type=int, default=15)
    ap.add_argument("--cats", default="architecture_overview,data_flow")
    ap.add_argument("--baseline", action="store_true",
                    help="evaluate the un-finetuned base models instead of adapters")
    ap.add_argument("--repo-slugs", default=None,
                    help="comma-separated org/name list; pins the held-out repo set "
                         "for a paired comparison instead of re-discovering")
    ap.add_argument("--repo-list-file", default=None,
                    help="path to a JSON file [{org,name,...}] pinning the repo set "
                         "(overrides discovery and --repo-slugs)")
    args = ap.parse_args()
    cats = [c.strip() for c in args.cats.split(",")]

    sw = ensure_pipeline()
    from vllm import LLM, SamplingParams
    from transformers import AutoTokenizer, AutoModelForCausalLM
    from huggingface_hub import snapshot_download
    import torch
    from peft import PeftModel

    # held-out repos
    from huggingface_hub import hf_hub_download
    reg_path = hf_hub_download(DATASET_REPO, "repo_registry.json", repo_type="dataset")
    registry = json.loads(Path(reg_path).read_text())
    registry_slugs = {f"{r['org']}/{r['name']}" for r in registry}
    if args.repo_list_file:
        eval_specs = json.loads(Path(args.repo_list_file).read_text())
        print(f"[eval] pinned {len(eval_specs)} repos from {args.repo_list_file}")
    elif args.repo_slugs:
        eval_specs = []
        for slug in [s.strip() for s in args.repo_slugs.split(",") if s.strip()]:
            org, name = slug.split("/")
            eval_specs.append({"org": org, "name": name, "language": "unknown",
                               "license": "unknown"})
        print(f"[eval] pinned {len(eval_specs)} repos (paired comparison)")
    else:
        eval_specs = discover_eval_repos(args.n_repos, registry_slugs)

    tok_judge = AutoTokenizer.from_pretrained(JUDGE)
    repos_ctx = []
    for spec in eval_specs:
        try:
            repo = sw.harvest(spec["org"], spec["name"], OUT_DIR / "repos")
            repo.language = spec["language"]
            ctx, n_tok = None, None
            for budget in (45_000, 30_000, 18_000):
                windows = sw.select_files(repo, budget=budget)
                ctx = sw.render_context(repo, windows)
                n_tok = len(tok_judge.encode(ctx))
                if n_tok <= 20_000:
                    break
            if len(repo.files) < 3:
                raise RuntimeError("too few source files")
            repos_ctx.append((spec, repo, ctx))
            print(f"[eval] {spec['org']}/{spec['name']}: {len(repo.files)} files, ctx {n_tok} tok")
        except Exception as e:
            print(f"[eval] HARVEST FAIL {spec['org']}/{spec['name']}: {e}")

    # ---------------- generate with each student ----------------
    results = []  # per (model, repo, category): page + all three scores
    arms = BASELINES if args.baseline else STUDENTS
    for model_name, adapter_id, base_id in arms:
        print(f"[eval] loading student {model_name} ({base_id}, "
              f"adapter={adapter_id})")
        merged_dir = OUT_DIR / f"merged_{model_name}"
        if adapter_id:
            base = AutoModelForCausalLM.from_pretrained(
                base_id, dtype=torch.bfloat16, device_map="cpu",
                attn_implementation="sdpa")
            base = PeftModel.from_pretrained(base, adapter_id)
            base = base.merge_and_unload()
            base.save_pretrained(merged_dir)
            AutoTokenizer.from_pretrained(base_id).save_pretrained(merged_dir)
            del base
            torch.cuda.empty_cache()
        else:
            # un-finetuned baseline: point vLLM straight at the base weights
            merged_dir = Path(base_id)
        llm = LLM(model=str(merged_dir), max_model_len=24576,
                  gpu_memory_utilization=0.92, max_num_seqs=8)
        sampling = SamplingParams(temperature=0.7, top_p=0.8, max_tokens=4096)

        convs, meta = [], []
        for spec, repo, ctx in repos_ctx:
            for cat in cats:
                user = (f"<START_OF_CONTEXT>\n{ctx}\n<END_OF_CONTEXT>\n\n"
                        f"<query>\n{CATEGORIES[cat]}\n</query>")
                convs.append([
                    {"role": "system", "content": SYSTEM_PROMPT},
                    {"role": "user", "content": user},
                ])
                meta.append((spec, repo, cat))
        outs = llm.chat(convs, sampling)
        for (spec, repo, cat), out in zip(meta, outs):
            o = out.outputs[0]
            text = o.text
            verdict = sw.verify_citations(text, repo)
            fmt = format_checks(text)
            results.append({
                "model": model_name, "source": f"{spec['org']}/{spec['name']}",
                "category": cat, "finish_reason": str(o.finish_reason),
                "page": text,
                "format": fmt,
                "citation": {"total": verdict["total_citations"],
                             "validity_rate": verdict["validity_rate"],
                             "problems": verdict["problems"][:10]},
            })
        del llm
        torch.cuda.empty_cache()
        print(f"[eval] {model_name}: {len(meta)} pages generated")

    # ---------------- judge pass ----------------
    print(f"[eval] loading judge {JUDGE}")
    judge_llm = LLM(model=JUDGE, max_model_len=24576,
                    gpu_memory_utilization=0.92, max_num_seqs=8)
    jsamp = SamplingParams(temperature=0.0, max_tokens=400)
    ctx_by_slug = {f"{s['org']}/{s['name']}": ctx for s, _, ctx in repos_ctx}
    jconvs = []
    for r in results:
        juser = (f"<SOURCE_CONTEXT>\n{ctx_by_slug[r['source']]}\n</SOURCE_CONTEXT>\n\n"
                 f"<GENERATED_PAGE>\n{r['page']}\n</GENERATED_PAGE>\n\n"
                 "Score factual fidelity per the system instructions. JSON only.")
        jconvs.append([{"role": "system", "content": JUDGE_SYSTEM},
                       {"role": "user", "content": juser}])
    jouts = judge_llm.chat(jconvs, jsamp)
    for r, jo in zip(results, jouts):
        raw = jo.outputs[0].text
        try:
            m = re.search(r"\{.*\}", raw, re.S)
            r["judge"] = json.loads(m.group(0)) if m else {"parse_error": raw[:200]}
        except Exception:
            r["judge"] = {"parse_error": raw[:200]}
    del judge_llm
    torch.cuda.empty_cache()

    # ---------------- summary + upload ----------------
    suffix = "baseline" if args.baseline else "finetuned"
    with (OUT_DIR / f"eval_pages_{suffix}.jsonl").open("w", encoding="utf-8") as f:
        for r in results:
            f.write(json.dumps(r) + "\n")

    summary = {
        "generated_at": datetime.now(timezone.utc).isoformat(),
        "n_repos": len(repos_ctx),
        "categories": cats,
        "judge": JUDGE,
        "models": {name: _agg(name, results) for name, _, _ in arms},
    }
    # family pairings for paired deltas (finetuned vs base, same repo/category)
    family_map = {
        "sw-coder-14b-v3": "base-qwen25coder-14b",
        "sw-qwen3-14b-v3": "base-qwen3-14b",
    }
    if not args.baseline:
        summary["paired_tests"] = {}
        for ft, base in family_map.items():
            for axis, key in (("citation_validity", "citation"), ("format", "format"),
                              ("fidelity", "fidelity")):
                summary["paired_tests"][f"{ft}_vs_{base}:{axis}"] = _paired_delta(
                    [dict(r, **{key: _page_metric(r, key)}) for r in results if r["model"] == ft],
                    [dict(r, **{key: _page_metric(r, key)}) for r in results if r["model"] == base],
                    key)
    (OUT_DIR / f"eval_summary_{suffix}.json").write_text(json.dumps(summary, indent=2))
    print("[eval] summary:", json.dumps(summary, indent=2))

    if os.environ.get("HF_TOKEN"):
        from huggingface_hub import HfApi
        api = HfApi()
        api.upload_file(path_or_fileobj=str(OUT_DIR / f"eval_pages_{suffix}.jsonl"),
                        path_in_repo=f"eval_pages_{suffix}.jsonl", repo_id=RESULTS_REPO,
                        repo_type="dataset")
        api.upload_file(path_or_fileobj=str(OUT_DIR / f"eval_summary_{suffix}.json"),
                        path_in_repo=f"eval_summary_{suffix}.json", repo_id=RESULTS_REPO,
                        repo_type="dataset")
        # snapshots of held-out repos for later re-verification
        snaps = OUT_DIR / "repos"
        tar = OUT_DIR / "eval_repo_snapshots.tar.gz"
        subprocess.run(["tar", "czf", str(tar), "-C", str(snaps), "."], check=True)
        api.upload_file(path_or_fileobj=str(tar), path_in_repo="eval_repo_snapshots.tar.gz",
                        repo_id=RESULTS_REPO, repo_type="dataset")
        print("[eval] uploaded to " + RESULTS_REPO)

def _bootstrap_ci(values, n_boot=1000, seed=7):
    """Mean and 95% percentile-bootstrap CI over page-level values."""
    import random
    rng = random.Random(seed)
    vals = [v for v in values if v is not None]
    if not vals:
        return None
    means = sorted(sum(rng.choices(vals, k=len(vals))) / len(vals)
                   for _ in range(n_boot))
    lo, hi = means[int(0.025 * n_boot)], means[int(0.975 * n_boot) - 1]
    return {"mean": round(sum(vals) / len(vals), 3),
            "ci95": [round(lo, 3), round(hi, 3)], "n": len(vals)}


def _paired_delta(rows_a, rows_b, key, n_boot=2000, seed=7):
    """Paired bootstrap over (repo, category) cells: distribution of
    mean(a) - mean(b). Returns CI and two-sided p (share of resampled deltas
    consistent with zero effect, i.e. sign-flipped relative to observed)."""
    import random
    rng = random.Random(seed)
    cell = {}
    for r in rows_a:
        if r[key] is not None:
            cell.setdefault((r["source"], r["category"]), [[], []])[0].append(r[key])
    for r in rows_b:
        if r[key] is not None:
            cell.setdefault((r["source"], r["category"]), [[], []])[1].append(r[key])
    deltas = [sum(a) / len(a) - sum(b) / len(b) for a, b in cell.values() if a and b]
    if len(deltas) < 5:
        return {"n_cells": len(deltas), "note": "too few paired cells"}
    boot = sorted(sum(rng.choices(deltas, k=len(deltas))) / len(deltas)
                  for _ in range(n_boot))
    lo, hi = boot[int(0.025 * n_boot)], boot[int(0.975 * n_boot) - 1]
    obs = sum(deltas) / len(deltas)
    # sign-flip randomization: null = resampled deltas with random sign flips;
    # p_one = share of null means at least as extreme as the observed delta
    null = []
    for _ in range(n_boot):
        samp = [d if rng.random() < 0.5 else -d for d in rng.choices(deltas, k=len(deltas))]
        null.append(sum(samp) / len(samp))
    if obs >= 0:
        p = sum(1 for m in null if m >= obs) / n_boot
    else:
        p = sum(1 for m in null if m <= obs) / n_boot
    return {"n_cells": len(deltas), "delta": round(obs, 3),
            "ci95": [round(lo, 3), round(hi, 3)], "p_two_sided": round(2 * min(p, 1 - p), 3)}


def _page_metric(r, key):
    if key == "format":
        return r["format"]["passed"] / r["format"]["total"]
    if key == "citation":
        return r["citation"]["validity_rate"]
    if key == "fidelity":
        v = r.get("judge", {}).get("fidelity_score")
        return v if isinstance(v, int) else None
    return None


def _agg(model, results):
    rows = [r for r in results if r["model"] == model]
    if not rows:
        return {}
    out = {"n_pages": len(rows)}
    for axis, key in (("format_score", "format"), ("citation_validity", "citation"),
                      ("fidelity_score", "fidelity")):
        out[axis] = _bootstrap_ci([_page_metric(r, key) for r in rows])
    out["mean_citations"] = round(sum(r["citation"]["total"] for r in rows) / len(rows), 1)
    by_cat = {}
    for cat in {r["category"] for r in rows}:
        sub = [r for r in rows if r["category"] == cat]
        by_cat[cat] = {axis: _bootstrap_ci([_page_metric(r, key) for r in sub])
                       for axis, key in (("format", "format"), ("citation", "citation"),
                                         ("fidelity", "fidelity"))}
    out["per_category"] = by_cat
    return out



if __name__ == "__main__":
    main()