#!/usr/bin/env python3
"""
PR enrichment pilot - process first 50 audit rows.
For each: find fix PR (GraphQL closedBy + body-search fallback),
extract test files from PR diff, update audit CSV.
"""
import csv
import json
import subprocess
import sys
import time
from pathlib import Path

AUDIT = Path("/tmp/nexuscrm-pingcap-pack/migration_audit.csv")
LIMIT = None  # process all rows that don't have fix_pr yet

REPO_MAP = {
    "pingcap-inc": "pingcap-inc/MetaDrivenPlatform4AI",
    "fork": "ywang-official/MetaDrivenPlatform4AI",
}

def run(cmd, timeout=30):
    try:
        r = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout)
        return r.stdout, r.stderr, r.returncode
    except subprocess.TimeoutExpired:
        return "", "TIMEOUT", 1

def find_fix_pr(source_repo, num):
    """Return list of {number, title, mergedAt} for fix PRs."""
    repo = REPO_MAP[source_repo]
    owner, name = repo.split("/")
    # Strategy 1: closedByPullRequestsReferences
    q = '''query($owner:String!,$name:String!,$num:Int!){
      repository(owner:$owner,name:$name){
        issue(number:$num){
          closedByPullRequestsReferences(first:5){
            nodes{number title mergedAt state}
          }
        }
      }
    }'''
    out, err, rc = run([
        "gh", "api", "graphql",
        "-f", f"query={q}",
        "-f", f"owner={owner}",
        "-f", f"name={name}",
        "-F", f"num={num}",
    ], timeout=20)
    prs = []
    if rc == 0 and out:
        try:
            data = json.loads(out)
            nodes = (data.get("data", {}).get("repository", {})
                     .get("issue", {}) or {}).get("closedByPullRequestsReferences", {}).get("nodes", []) or []
            for n in nodes:
                if n.get("state") == "MERGED":
                    prs.append(n)
        except Exception:
            pass
    if prs:
        return prs, "closedBy"
    # Strategy 2: search merged PRs that mention this issue
    out, err, rc = run([
        "gh", "search", "prs",
        f"#{num}", "--repo", repo, "--state", "merged",
        "--json", "number,title", "--limit", "3",
    ], timeout=20)
    if rc == 0 and out:
        try:
            arr = json.loads(out)
            for n in arr:
                prs.append({"number": n["number"], "title": n["title"], "mergedAt": None})
        except Exception:
            pass
    return prs, ("search" if prs else "none")

def get_pr_test_files(repo, pr_num):
    out, err, rc = run([
        "gh", "api", f"repos/{repo}/pulls/{pr_num}/files",
        "--paginate",
        "--jq", "[.[].filename]",
    ], timeout=30)
    if rc != 0 or not out:
        return []
    try:
        files = json.loads(out)
    except Exception:
        return []
    tests = []
    for f in files:
        fl = f.lower()
        if (f.endswith("_test.go") or f.endswith(".test.ts") or f.endswith(".test.tsx")
            or "/tests/" in fl or fl.startswith("tests/")
            or "/cases/" in fl):
            tests.append(f)
    return tests

def main():
    rows = []
    with AUDIT.open() as f:
        reader = csv.DictReader(f)
        fieldnames = reader.fieldnames
        for r in reader:
            rows.append(r)
    if "fix_pr" not in fieldnames:
        # insert fix_pr after existing_tests
        fieldnames = list(fieldnames)
        idx = fieldnames.index("existing_tests")
        fieldnames.insert(idx, "fix_pr")
    for r in rows:
        r.setdefault("fix_pr", "")

    # Skip rows already processed (have fix_pr or marked as scanned). Use marker NONE for empty scans.
    todo = [r for r in rows if not r.get("fix_pr")]
    if LIMIT is not None:
        todo = todo[:LIMIT]
    print(f"Processing {len(todo)} rows...", file=sys.stderr)
    found_pr = 0
    with_tests = 0
    samples = []
    t0 = time.time()
    for i, r in enumerate(todo, 1):
        src = r["source_repo"]
        num = r["issue_number"]
        if r.get("fix_pr"):
            continue
        prs, strat = find_fix_pr(src, int(num))
        if not prs:
            r["fix_pr"] = "NONE"  # scanned, no PR found — avoid re-scanning
        if prs:
            found_pr += 1
            pr_refs = ";".join(f"{REPO_MAP[src]}#{p['number']}" for p in prs)
            r["fix_pr"] = pr_refs
            # get test files from first PR
            tests_all = []
            for p in prs[:2]:
                tests_all.extend(get_pr_test_files(REPO_MAP[src], p["number"]))
            tests_all = list(dict.fromkeys(tests_all))  # dedupe
            if tests_all:
                with_tests += 1
                # merge into existing_tests
                existing = r.get("existing_tests", "") or ""
                merged = existing
                for t in tests_all:
                    if t not in merged:
                        merged = (merged + ";" + t) if merged else t
                r["existing_tests"] = merged
            if len(samples) < 5:
                samples.append({
                    "src": src, "num": num, "title": r["title"][:60],
                    "prs": pr_refs, "tests": tests_all[:3], "strat": strat,
                })
        if i % 10 == 0:
            elapsed = time.time() - t0
            print(f"  [{i}/{len(todo)}] found={found_pr} tests={with_tests} ({elapsed:.0f}s)", file=sys.stderr)

    # write back full CSV
    with AUDIT.open("w", newline="") as f:
        w = csv.DictWriter(f, fieldnames=fieldnames)
        w.writeheader()
        for r in rows:
            w.writerow({k: r.get(k, "") for k in fieldnames})

    elapsed = time.time() - t0
    print(f"\n=== Pilot done in {elapsed:.0f}s ===", file=sys.stderr)
    print(f"Processed: {len(todo)}", file=sys.stderr)
    print(f"Found fix PR: {found_pr} ({100*found_pr/len(todo):.0f}%)", file=sys.stderr)
    print(f"With test files: {with_tests} ({100*with_tests/len(todo):.0f}%)", file=sys.stderr)
    print(f"\nSamples:", file=sys.stderr)
    for s in samples:
        print(f"  {s['src']}#{s['num']} [{s['strat']}] {s['title']}", file=sys.stderr)
        print(f"    fix_pr: {s['prs']}", file=sys.stderr)
        print(f"    tests:  {s['tests']}", file=sys.stderr)

if __name__ == "__main__":
    main()