#!/usr/bin/env python3
"""Compare Fold benchmark outputs with literature baselines and F1 ceilings."""
from __future__ import annotations

import argparse
import datetime as dt
import json
import re
import sys
from pathlib import Path
from typing import Any


REPO_ROOT = Path(__file__).resolve().parent.parent
RESULTS_ROOT = REPO_ROOT / "bench" / "results"
DEFAULT_LITERATURE = REPO_ROOT / "bench" / "literature_baselines.json"
sys.path.insert(0, str(REPO_ROOT / "scripts" / "benchmarks"))

from fold_eval.metrics import classification_metrics

SOLVER_RE = re.compile(
    r"^\[solver-profile\]\s+(?P<detector>\S+)\s+"
    r"queries=(?P<queries>\d+)\s+sat=(?P<sat>\d+)\s+unsat=(?P<unsat>\d+)\s+"
    r"unknown=(?P<unknown>\d+)\s+decisions=(?P<decisions>\d+)\s+"
    r"conflicts=(?P<conflicts>\d+)\s+restarts=(?P<restarts>\d+)\s+"
    r"propagations=(?P<propagations>\d+)\s+z3_time_ms=(?P<z3_time_ms>\d+)\s+"
    r"wall_ms=(?P<wall_ms>\d+)"
)


def now_iso() -> str:
    return dt.datetime.now(dt.UTC).strftime("%Y-%m-%dT%H-%M-%SZ")


def pct(value: float) -> str:
    return f"{100.0 * value:.1f}%"


def signed_int(value: int) -> str:
    return f"{value:+d}"


def signed_seconds(value: float | None) -> str:
    if value is None:
        return ""
    return f"{value:+.1f}s"


def score_metrics(tp: int, fp: int, fn: int) -> dict[str, float]:
    metrics = classification_metrics(tp, fp, fn)
    return {key: float(metrics[key]) for key in ("precision", "recall", "f1")}


def infer_cwe(*values: str) -> str:
    for value in values:
        match = re.search(r"(?:juliet[-_/]|CWE)(\d{3})", value)
        if match:
            return match.group(1)
        match = re.search(r"[/_-](\d{3})(?:[/_-]|$)", value)
        if match:
            return match.group(1)
    return "unknown"


def profile_aliases(kind: str) -> set[str]:
    aliases = {part.strip() for part in kind.split(",") if part.strip()}
    if "oob-access" in aliases:
        aliases.add("oob")
    if "dead" in aliases:
        aliases.add("dead-branch")
    return aliases


def parse_solver_profile_log(path: Path, aliases: set[str]) -> dict[str, int] | None:
    if not path.exists():
        return None
    totals = {
        "queries": 0,
        "sat": 0,
        "unsat": 0,
        "unknown": 0,
        "decisions": 0,
        "conflicts": 0,
        "restarts": 0,
        "propagations": 0,
        "z3_time_ms": 0,
        "wall_ms": 0,
    }
    matched = False
    for line in path.read_text(errors="replace").splitlines():
        match = SOLVER_RE.match(line.strip())
        if not match or match.group("detector") not in aliases:
            continue
        matched = True
        for key in totals:
            totals[key] += int(match.group(key))
    return totals if matched else None


def row_from_score(
    *,
    source_path: Path,
    tool: str,
    variant: str,
    cwe: str,
    kind: str,
    score: dict[str, Any],
    findings_count: int | None,
    elapsed_s: float | None,
    z3_queries: float | None,
) -> dict[str, Any]:
    tp = int(score.get("true_positive", 0) or 0)
    fp = int(score.get("false_positive", 0) or 0)
    fn = int(score.get("false_negative", 0) or 0)
    gold_total = int(score.get("gold_total", tp + fn) or 0)
    finding_groups = int(score.get("finding_groups", tp + fp) or 0)
    metrics = score_metrics(tp, fp, fn)
    return {
        "source": str(source_path),
        "tool": tool,
        "variant": variant,
        "cwe": cwe,
        "kind": kind,
        "true_positive": tp,
        "false_positive": fp,
        "false_negative": fn,
        "gold_total": gold_total,
        "finding_groups": finding_groups,
        "findings_count": findings_count,
        "elapsed_s": elapsed_s,
        "z3_queries": z3_queries,
        **metrics,
    }


def rows_from_summary(path: Path, doc: dict[str, Any]) -> list[dict[str, Any]]:
    cwe = infer_cwe(str(doc.get("name", "")), str(doc.get("gold", "")), str(doc.get("corpus", "")))
    kind = ",".join(str(item) for item in doc.get("detectors", [])) or "unknown"
    rows: list[dict[str, Any]] = []
    for variant, data in doc.get("variant_summaries", {}).items():
        score = data.get("score")
        if not isinstance(score, dict):
            continue
        profile = data.get("profile_avg", {}) if isinstance(data.get("profile_avg"), dict) else {}
        rows.append(row_from_score(
            source_path=path,
            tool="fold",
            variant=str(variant),
            cwe=cwe,
            kind=kind,
            score=score,
            findings_count=int(data.get("findings_count", 0) or 0),
            elapsed_s=float(data.get("elapsed_s_avg", 0.0) or 0.0),
            z3_queries=float(profile.get("z3_queries", 0.0) or 0.0),
        ))
    return rows


def rows_from_tool_payload(path: Path, doc: dict[str, Any]) -> list[dict[str, Any]]:
    score = doc.get("score")
    if not isinstance(score, dict):
        return []
    profile = doc.get("profile", {}) if isinstance(doc.get("profile"), dict) else {}
    solver = profile.get("solver", {}) if isinstance(profile.get("solver"), dict) else {}
    artifacts = doc.get("artifacts", {}) if isinstance(doc.get("artifacts"), dict) else {}
    kind = str(doc.get("kind") or doc.get("detector") or "unknown")
    if float(solver.get("queries", 0.0) or 0.0) == 0.0 and artifacts.get("stderr_log"):
        recovered = parse_solver_profile_log(Path(str(artifacts["stderr_log"])), profile_aliases(kind))
        if recovered is not None:
            solver = recovered
    return [row_from_score(
        source_path=path,
        tool=str(doc.get("tool", path.stem)),
        variant=str(doc.get("variant", "default")),
        cwe=str(doc.get("cwe") or infer_cwe(str(path))),
        kind=kind,
        score=score,
        findings_count=len(doc.get("findings", []) or []),
        elapsed_s=float(doc.get("elapsed_s", 0.0) or 0.0),
        z3_queries=float(solver.get("queries", 0.0) or 0.0),
    )]


def load_rows_from_path(path: Path) -> list[dict[str, Any]]:
    if path.is_dir():
        rows: list[dict[str, Any]] = []
        if (path / "summary.json").exists():
            rows.extend(load_rows_from_path(path / "summary.json"))
        for payload_name in ("fold.json", "clang-sa.json", "esbmc.json"):
            for child in sorted(path.glob(f"*/{payload_name}")):
                rows.extend(load_rows_from_path(child))
        return rows

    doc = json.loads(path.read_text())
    if not isinstance(doc, dict):
        return []
    if "variant_summaries" in doc:
        return rows_from_summary(path, doc)
    return rows_from_tool_payload(path, doc)


def discover_default_runs() -> list[Path]:
    latest_by_name: dict[str, Path] = {}
    for path in sorted(RESULTS_ROOT.glob("*summary-ab/summary.json")):
        try:
            name = str(json.loads(path.read_text()).get("name", path.parent.name))
        except Exception:
            name = path.parent.name
        latest_by_name[name] = path
    return list(latest_by_name.values())


def compute_ceiling(row: dict[str, Any]) -> dict[str, Any]:
    tp = int(row["true_positive"])
    fp = int(row["false_positive"])
    fn = int(row["false_negative"])
    gold_total = int(row["gold_total"])
    finding_groups = int(row.get("finding_groups", tp + fp) or 0)

    current = score_metrics(tp, fp, fn)
    precision_only = score_metrics(tp, 0, fn)
    recall_only = score_metrics(gold_total, fp, 0)
    budget_tp = min(gold_total, max(finding_groups, tp))
    alert_budget = score_metrics(budget_tp, 0, max(0, gold_total - budget_tp))
    oracle = score_metrics(gold_total, 0, 0)
    precision_gain = precision_only["f1"] - current["f1"]
    recall_gain = recall_only["f1"] - current["f1"]
    if recall_gain > precision_gain * 1.1:
        bottleneck = "recall"
        next_action = "prioritize missed patterns, interprocedural coverage, and path/model completeness"
    elif precision_gain > recall_gain * 1.1:
        bottleneck = "precision"
        next_action = "prioritize false-positive pruning, stronger feasibility checks, and post-filters"
    else:
        bottleneck = "mixed"
        next_action = "split work between missed patterns and false-positive pruning"

    return {
        "current": current,
        "precision_only_ceiling": precision_only,
        "recall_only_ceiling": recall_only,
        "alert_budget_ceiling": alert_budget,
        "oracle_ceiling": oracle,
        "precision_gain_f1": precision_gain,
        "recall_gain_f1": recall_gain,
        "oracle_gap_f1": oracle["f1"] - current["f1"],
        "bottleneck": bottleneck,
        "next_action": next_action,
    }


def best_rows(rows: list[dict[str, Any]]) -> list[dict[str, Any]]:
    best: dict[tuple[str, str], dict[str, Any]] = {}

    def is_better(candidate: dict[str, Any], current: dict[str, Any]) -> bool:
        # Favor recall/finding gains when F1 differences are within report
        # rounding noise. This keeps ceilings aligned with discovery value
        # instead of hiding useful extra TPs behind tiny precision tradeoffs.
        f1_delta = float(candidate["f1"]) - float(current["f1"])
        if abs(f1_delta) > 0.005:
            return f1_delta > 0
        candidate_rank = (
            int(candidate["true_positive"]),
            float(candidate["recall"]),
            int(candidate.get("finding_groups", 0) or 0),
            -int(candidate["false_negative"]),
            -int(candidate["false_positive"]),
        )
        current_rank = (
            int(current["true_positive"]),
            float(current["recall"]),
            int(current.get("finding_groups", 0) or 0),
            -int(current["false_negative"]),
            -int(current["false_positive"]),
        )
        return candidate_rank > current_rank

    for row in rows:
        key = (str(row["cwe"]), str(row["kind"]))
        if key not in best or is_better(row, best[key]):
            best[key] = row
    return [best[key] for key in sorted(best)]


def signal_recommendation(signal: dict[str, Any]) -> str:
    tp_delta = int(signal.get("tp_delta", 0) or 0)
    fp_delta = int(signal.get("fp_delta", 0) or 0)
    elapsed_delta_s = signal.get("elapsed_delta_s")
    f1_delta = float(signal.get("f1_delta", 0.0) or 0.0)
    if tp_delta > 0:
        return "preserve finding gain; optimize implementation cost around this path"
    if tp_delta < 0:
        return "fix recall regression before optimizing speed"
    if fp_delta > 0 and f1_delta < 0:
        return "precision regression; tighten feasibility or scoring before expanding search"
    if elapsed_delta_s is not None and float(elapsed_delta_s) > 10:
        return "no discovery gain; reduce overhead or make this path lazy"
    return "neutral regression guard"


def seconds_per_added_tp(elapsed_delta_s: float | None, tp_delta: int) -> float | None:
    if elapsed_delta_s is None or elapsed_delta_s <= 0 or tp_delta <= 0:
        return None
    return elapsed_delta_s / tp_delta


def optimization_signals(rows: list[dict[str, Any]]) -> list[dict[str, Any]]:
    groups: dict[tuple[str, str, str], list[dict[str, Any]]] = {}
    for row in rows:
        groups.setdefault((str(row["tool"]), str(row["cwe"]), str(row["kind"])), []).append(row)

    signals: list[dict[str, Any]] = []
    for (tool, cwe, kind), group in groups.items():
        baseline = next((row for row in group if row.get("variant") == "baseline"), None)
        if baseline is None:
            continue
        for row in group:
            if row.get("variant") == "baseline":
                continue
            elapsed_delta_s = None
            if row.get("elapsed_s") is not None and baseline.get("elapsed_s") is not None:
                elapsed_delta_s = float(row["elapsed_s"]) - float(baseline["elapsed_s"])
            z3_query_delta = None
            if row.get("z3_queries") is not None and baseline.get("z3_queries") is not None:
                z3_query_delta = float(row["z3_queries"]) - float(baseline["z3_queries"])
            tp_delta = int(row["true_positive"]) - int(baseline["true_positive"])
            signal = {
                "tool": tool,
                "cwe": cwe,
                "kind": kind,
                "variant": row["variant"],
                "source": row.get("source"),
                "tp_delta": tp_delta,
                "fp_delta": int(row["false_positive"]) - int(baseline["false_positive"]),
                "fn_delta": int(row["false_negative"]) - int(baseline["false_negative"]),
                "recall_delta": float(row["recall"]) - float(baseline["recall"]),
                "f1_delta": float(row["f1"]) - float(baseline["f1"]),
                "finding_group_delta": int(row.get("finding_groups", 0) or 0)
                - int(baseline.get("finding_groups", 0) or 0),
                "elapsed_delta_s": elapsed_delta_s,
                "z3_query_delta": z3_query_delta,
                "seconds_per_added_tp": seconds_per_added_tp(elapsed_delta_s, tp_delta),
            }
            signal["recommendation"] = signal_recommendation(signal)
            signals.append(signal)

    def sort_key(signal: dict[str, Any]) -> tuple[Any, ...]:
        tp_delta = int(signal.get("tp_delta", 0) or 0)
        elapsed_delta = float(signal.get("elapsed_delta_s", 0.0) or 0.0)
        if tp_delta > 0:
            bucket = 0
        elif tp_delta < 0:
            bucket = 1
        elif elapsed_delta > 10:
            bucket = 2
        else:
            bucket = 3
        return (bucket, -tp_delta, -elapsed_delta, str(signal["cwe"]), str(signal["kind"]))

    return sorted(signals, key=sort_key)


def metric_text(metric: dict[str, Any]) -> str:
    name = str(metric["name"]).replace("_", " ")
    unit = metric.get("unit")
    if "value_min" in metric and "value_max" in metric:
        lo = float(metric["value_min"])
        hi = float(metric["value_max"])
        value = f"{pct(lo)}-{pct(hi)}" if unit == "ratio" else f"{lo:g}-{hi:g}"
    else:
        raw = float(metric.get("value", 0.0))
        value = pct(raw) if unit == "ratio" else f"{raw:g}"
    qualifier = str(metric.get("qualifier", ""))
    if qualifier:
        value = f"{qualifier} {value}"
    return f"{name}: {value}"


def render_markdown(rows: list[dict[str, Any]], literature: dict[str, Any]) -> str:
    best = best_rows(rows)
    signals = optimization_signals(rows)
    lines = [
        "# Literature and Ceiling Comparison",
        "",
        "## Fold Runs",
        "",
        "| CWE | Kind | Variant | TP | FP | FN | Precision | Recall | F1 | Findings | Time | Z3 Queries |",
        "|-----|------|---------|----|----|----|-----------|--------|----|----------|------|------------|",
    ]
    for row in sorted(rows, key=lambda item: (item["cwe"], item["kind"], item["variant"])):
        elapsed = "" if row.get("elapsed_s") is None else f"{float(row['elapsed_s']):.1f}s"
        z3 = "" if row.get("z3_queries") is None else f"{float(row['z3_queries']):.0f}"
        lines.append(
            f"| CWE-{row['cwe']} | {row['kind']} | {row['variant']} | {row['true_positive']} | "
            f"{row['false_positive']} | {row['false_negative']} | {pct(row['precision'])} | "
            f"{pct(row['recall'])} | {pct(row['f1'])} | {row['finding_groups']} | {elapsed} | {z3} |"
        )

    lines.extend([
        "",
        "## Best Fold Variant vs Ceilings",
        "",
        "| CWE | Variant | TP | Recall | Current F1 | No-FP F1 | Full-Recall F1 | Alert-Budget F1 | Oracle Gap | Bottleneck | Next Action |",
        "|-----|---------|----|--------|------------|----------|----------------|-----------------|------------|------------|-------------|",
    ])
    for row in best:
        ceiling = compute_ceiling(row)
        lines.append(
            f"| CWE-{row['cwe']} | {row['variant']} | {row['true_positive']} | "
            f"{pct(row['recall'])} | {pct(ceiling['current']['f1'])} | "
            f"{pct(ceiling['precision_only_ceiling']['f1'])} | {pct(ceiling['recall_only_ceiling']['f1'])} | "
            f"{pct(ceiling['alert_budget_ceiling']['f1'])} | {pct(ceiling['oracle_gap_f1'])} | "
            f"{ceiling['bottleneck']} | {ceiling['next_action']} |"
        )

    if signals:
        lines.extend([
            "",
            "## Optimization Signals",
            "",
            "| CWE | Kind | Variant | TP Delta | FP Delta | Recall Delta | F1 Delta | Time Delta | Sec / Added TP | Query Delta | Recommendation |",
            "|-----|------|---------|----------|----------|--------------|----------|------------|----------------|-------------|----------------|",
        ])
        for signal in signals:
            sec_per_tp = signal.get("seconds_per_added_tp")
            sec_per_tp_text = "n/a" if sec_per_tp is None else f"{float(sec_per_tp):.1f}s"
            query_delta = signal.get("z3_query_delta")
            query_delta_text = "" if query_delta is None else signed_int(int(round(float(query_delta))))
            lines.append(
                f"| CWE-{signal['cwe']} | {signal['kind']} | {signal['variant']} | "
                f"{signed_int(int(signal['tp_delta']))} | {signed_int(int(signal['fp_delta']))} | "
                f"{float(signal['recall_delta']):+.1%} | {float(signal['f1_delta']):+.1%} | "
                f"{signed_seconds(signal.get('elapsed_delta_s'))} | {sec_per_tp_text} | "
                f"{query_delta_text} | {signal['recommendation']} |"
            )

    lines.extend([
        "",
        "## Literature Baselines",
        "",
        "| Baseline | Benchmark | Scope | Reported Metrics | Comparability |",
        "|----------|-----------|-------|------------------|---------------|",
    ])
    for baseline in literature.get("baselines", []):
        metrics = "; ".join(metric_text(metric) for metric in baseline.get("metrics", [])) or "methodology only"
        lines.append(
            f"| {baseline['id']} | {baseline['benchmark']} | {baseline['tool_scope']} | "
            f"{metrics} | {baseline['comparability']} |"
        )

    lines.extend([
        "",
        "## Benchmark Expansion Targets",
        "",
        "| Target | Status | Runner | Scope | Next Action |",
        "|--------|--------|--------|-------|-------------|",
    ])
    for target in literature.get("benchmark_targets", []):
        lines.append(
            f"| {target['id']} | {target['status']} | `{target['runner']}` | {target['scope']} | {target['next_action']} |"
        )

    lines.extend([
        "",
        "## Sources",
        "",
        "| ID | Year | Title | DOI |",
        "|----|------|-------|-----|",
    ])
    for ref in literature.get("references", []):
        lines.append(f"| {ref['id']} | {ref['year']} | {ref['title']} | {ref['doi']} |")
    lines.append("")
    return "\n".join(lines)


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--fold-run", action="append", default=[], help="summary.json, tool payload JSON, or result directory")
    parser.add_argument("--literature", default=str(DEFAULT_LITERATURE))
    parser.add_argument("--out-dir", default="", help="Output directory; defaults to bench/results/<timestamp>_literature-comparison")
    parser.add_argument("--no-write", action="store_true", help="Print only; do not write result artifacts")
    args = parser.parse_args()

    literature_path = Path(args.literature).resolve()
    if not literature_path.exists():
        print(f"literature baseline file not found: {literature_path}", file=sys.stderr)
        return 2
    literature = json.loads(literature_path.read_text())

    input_paths = [Path(item).resolve() for item in args.fold_run]
    if not input_paths:
        input_paths = discover_default_runs()
    if not input_paths:
        print("no benchmark summaries found; pass --fold-run", file=sys.stderr)
        return 2

    rows: list[dict[str, Any]] = []
    for path in input_paths:
        if not path.exists():
            print(f"benchmark input not found: {path}", file=sys.stderr)
            return 2
        rows.extend(load_rows_from_path(path))
    rows = [row for row in rows if row.get("gold_total", 0)]
    if not rows:
        print("no scored benchmark rows found", file=sys.stderr)
        return 2

    payload = {
        "timestamp": now_iso(),
        "inputs": [str(path) for path in input_paths],
        "literature": str(literature_path),
        "fold_rows": rows,
        "ceilings": [{**row, "ceiling": compute_ceiling(row)} for row in best_rows(rows)],
        "optimization_signals": optimization_signals(rows),
        "literature_baselines": literature.get("baselines", []),
        "benchmark_targets": literature.get("benchmark_targets", []),
    }
    report = render_markdown(rows, literature)

    if not args.no_write:
        out_dir = Path(args.out_dir).resolve() if args.out_dir else RESULTS_ROOT / f"{payload['timestamp']}_literature-comparison"
        out_dir.mkdir(parents=True, exist_ok=True)
        (out_dir / "summary.json").write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")
        (out_dir / "summary.md").write_text(report + "\n")
        print(report)
        print(f"wrote {out_dir}")
    else:
        print(report)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())