#!/usr/bin/env python3
"""Collect Vivado-only QoR evidence for a project, checkpoint, or implemented run."""

from __future__ import annotations

import argparse
import json
import re
import subprocess
import sys
from datetime import datetime
from pathlib import Path
from typing import Any


def now_tag() -> str:
    return datetime.now().strftime("%Y%m%d_%H%M%S")


def ps_path(path: Path) -> str:
    return str(path).replace("\\", "/")


def find_dcp_candidates(project: str, run: str) -> list[Path]:
    """Locate routed DCP artifacts for a project+run without opening Vivado.

    The standard impl run layout is ``<project_dir>/<project>.runs/<run>/``;
    some flows use a different project folder name, so ``*.runs`` is also
    searched.  Returns existing paths only.
    """
    project_path = Path(project).resolve()
    project_dir = (
        project_path.parent
        if project_path.is_file() or project_path.suffix.lower() == ".xpr"
        else project_path
    )
    project_stem = project_path.stem
    patterns = [
        f"{project_stem}.runs/{run}/*_routed.dcp",
        f"*.runs/{run}/*_routed.dcp",
    ]
    candidates: list[Path] = []
    seen: set[Path] = set()
    for pattern in patterns:
        for match in sorted(project_dir.glob(pattern)):
            resolved = match.resolve()
            if resolved not in seen:
                seen.add(resolved)
                candidates.append(resolved)
    return candidates


def write_tcl(args: argparse.Namespace, out_dir: Path) -> Path:
    raw = out_dir / "raw"
    tcl = out_dir / "collect_qor.tcl"
    lines = [
        "set_param general.maxThreads 8",
        f"file mkdir {{{ps_path(raw)}}}",
        f"set raw_dir {{{ps_path(raw)}}}",
    ]
    if args.project:
        lines.append(f"open_project {{{ps_path(Path(args.project).resolve())}}}")
        if args.run:
            if getattr(args, "dcp_fallback", True):
                candidates = find_dcp_candidates(args.project, args.run)
                lines.append(f"set _dcp_fallback [list {' '.join('{' + ps_path(p) + '}' for p in candidates)}]")
                lines.append(f"if {{[catch {{open_run {{{args.run}}}}} _open_run_err]}} {{")
                lines.append("  puts \"OPEN_RUN_FAILED: $_open_run_err\"")
                lines.append('  set _fallback_dcp ""')
                lines.append("  foreach _cand $_dcp_fallback {")
                lines.append("    if {[file exists $_cand]} { set _fallback_dcp $_cand; break }")
                lines.append("  }")
                lines.append('  if {$_fallback_dcp ne ""} {')
                lines.append("    puts \"DCP_FALLBACK: $_fallback_dcp\"")
                lines.append("    open_checkpoint $_fallback_dcp")
                lines.append("  } else {")
                lines.append("    puts \"ERROR: open_run failed and no routed DCP candidate found\"")
                lines.append("    exit 3")
                lines.append("  }")
                lines.append("}")
            else:
                lines.append(f"open_run {{{args.run}}}")
    elif args.dcp:
        lines.append(f"open_checkpoint {{{ps_path(Path(args.dcp).resolve())}}}")
    else:
        lines.append("puts \"ERROR: provide -project or -dcp\"")
        lines.append("exit 2")

    lines.extend([
        "set rpt_timing_summary [file join $raw_dir timing_summary.rpt]",
        "set rpt_timing_paths [file join $raw_dir timing_paths.rpt]",
        "set rpt_fanout [file join $raw_dir high_fanout.rpt]",
        "set rpt_util [file join $raw_dir utilization.rpt]",
        "set rpt_design_analysis [file join $raw_dir design_analysis.rpt]",
        "set rpt_route_status [file join $raw_dir route_status.rpt]",
        "catch { report_timing_summary -file $rpt_timing_summary -delay_type max -max_paths 20 -nworst 1 } timing_summary_err",
        "catch { report_timing -file $rpt_timing_paths -delay_type max -sort_by slack -max_paths 20 -nworst 1 } timing_paths_err",
        "catch { report_high_fanout_nets -file $rpt_fanout -timing -max_nets 50 } fanout_err",
        "catch { report_utilization -file $rpt_util } util_err",
        "catch { report_design_analysis -file $rpt_design_analysis -timing -logic_level_distribution -congestion } analysis_err",
        "catch { report_route_status -file $rpt_route_status } route_status_err",
        f"set status_txt {{{ps_path(out_dir / 'vivado_status.txt')}}}",
        "set fh [open $status_txt w]",
        "puts $fh \"timing_summary_error=$timing_summary_err\"",
        "puts $fh \"timing_paths_error=$timing_paths_err\"",
        "puts $fh \"fanout_error=$fanout_err\"",
        "puts $fh \"util_error=$util_err\"",
        "puts $fh \"analysis_error=$analysis_err\"",
        "puts $fh \"route_status_error=$route_status_err\"",
        "close $fh",
        "exit",
    ])
    tcl.write_text("\n".join(lines) + "\n", encoding="utf-8")
    return tcl


def parse_fallback_evidence(stdout_text: str) -> dict[str, Any]:
    """Report whether open_run failed and a DCP fallback was used."""
    open_run_failed = "OPEN_RUN_FAILED:" in stdout_text
    dcp_path: str | None = None
    for line in stdout_text.splitlines():
        if line.startswith("DCP_FALLBACK: "):
            dcp_path = line.split("DCP_FALLBACK: ", 1)[1].strip()
            break
    return {
        "open_run_failed": open_run_failed,
        "dcp_fallback_used": dcp_path is not None,
        "dcp_path": dcp_path,
    }


def parse_float(pattern: str, text: str) -> float | None:
    match = re.search(pattern, text, re.IGNORECASE)
    if not match:
        return None
    try:
        return float(match.group(1))
    except ValueError:
        return None


def parse_int(pattern: str, text: str) -> int | None:
    match = re.search(pattern, text, re.IGNORECASE)
    if not match:
        return None
    try:
        return int(match.group(1).replace(",", ""))
    except ValueError:
        return None


def summarize_reports(out_dir: Path) -> dict[str, Any]:
    raw = out_dir / "raw"
    timing = (raw / "timing_summary.rpt").read_text(encoding="utf-8", errors="replace") if (raw / "timing_summary.rpt").exists() else ""
    fanout = (raw / "high_fanout.rpt").read_text(encoding="utf-8", errors="replace") if (raw / "high_fanout.rpt").exists() else ""
    status_path = out_dir / "vivado_status.txt"
    status_text = status_path.read_text(encoding="utf-8", errors="replace") if status_path.exists() else ""

    return {
        "timing": {
            **parse_design_timing_summary(timing),
        },
        "reports": {
            "timing_summary": str(raw / "timing_summary.rpt"),
            "timing_paths": str(raw / "timing_paths.rpt"),
            "high_fanout": str(raw / "high_fanout.rpt"),
            "utilization": str(raw / "utilization.rpt"),
            "design_analysis": str(raw / "design_analysis.rpt"),
            "route_status": str(raw / "route_status.rpt"),
        },
        "fanout_report_bytes": len(fanout),
        "top_fanout": parse_high_fanout(fanout, limit=10),
        "vivado_status_text": status_text,
    }


def parse_design_timing_summary(text: str) -> dict[str, Any]:
    lines = text.splitlines()
    for index, line in enumerate(lines):
        if "WNS(ns)" in line and "TNS(ns)" in line and "TNS Failing Endpoints" in line:
            for candidate in lines[index + 1 : index + 6]:
                nums = re.findall(r"-?\d+(?:\.\d+)?", candidate)
                if len(nums) >= 4:
                    return {
                        "wns": float(nums[0]),
                        "tns": float(nums[1]),
                        "failing_endpoints": int(float(nums[2])),
                        "total_endpoints": int(float(nums[3])),
                    }
    return {
        "wns": parse_float(r"Slack \(MET\)\s*:\s*(-?\d+(?:\.\d+)?)ns", text),
        "tns": None,
        "failing_endpoints": None,
        "total_endpoints": None,
    }


def parse_high_fanout(text: str, limit: int) -> list[dict[str, Any]]:
    rows: list[dict[str, Any]] = []
    for line in text.splitlines():
        if not line.startswith("| "):
            continue
        if "Net Name" in line or "Driver Type" in line:
            continue
        parts = [part.strip() for part in line.strip().strip("|").split("|")]
        if len(parts) != 5:
            continue
        try:
            rows.append(
                {
                    "net": parts[0],
                    "fanout": int(parts[1]),
                    "driver_type": parts[2],
                    "worst_slack_ns": float(parts[3]),
                    "worst_delay_ns": float(parts[4]),
                }
            )
        except ValueError:
            continue
        if len(rows) >= limit:
            break
    return rows


def run_vivado(vivado: str, tcl: Path, out_dir: Path, timeout_s: int) -> dict[str, Any]:
    log = out_dir / "vivado.log"
    jou = out_dir / "vivado.jou"
    cmd = [vivado, "-mode", "batch", "-source", str(tcl), "-log", str(log), "-journal", str(jou)]
    proc = subprocess.run(cmd, text=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, timeout=timeout_s)
    (out_dir / "vivado_stdout.txt").write_text(proc.stdout, encoding="utf-8", errors="replace")
    return {"cmd": cmd, "returncode": proc.returncode, "stdout": str(out_dir / "vivado_stdout.txt")}


def main(argv: list[str]) -> int:
    parser = argparse.ArgumentParser(description="Collect Vivado QoR evidence")
    parser.add_argument("--vivado", default="vivado", help="Vivado executable")
    parser.add_argument("--project", default=None, help="Vivado .xpr path")
    parser.add_argument("--run", default=None, help="implemented run name to open")
    parser.add_argument("--dcp", default=None, help="DCP path")
    parser.add_argument("--out-root", default=".vivado_ai_mcp/qor", help="output root")
    parser.add_argument("--name", default=None, help="capture name")
    parser.add_argument("--dry-run", action="store_true", help="write Tcl/manifest only")
    parser.add_argument("--timeout-s", type=int, default=1800)
    parser.add_argument(
        "--dcp-fallback",
        action=argparse.BooleanOptionalAction,
        default=True,
        help="when project+run open_run fails, fall back to an existing *_routed.dcp",
    )
    args = parser.parse_args(argv)

    if not args.project and not args.dcp:
        parser.error("provide --project or --dcp")

    out_dir = Path(args.out_root).resolve() / (args.name or f"qor_{now_tag()}")
    out_dir.mkdir(parents=True, exist_ok=True)
    (out_dir / "raw").mkdir(exist_ok=True)

    tcl = write_tcl(args, out_dir)
    manifest: dict[str, Any] = {
        "created_at": datetime.now().isoformat(timespec="seconds"),
        "tool": "vivado_qor_collect",
        "inputs": {
            "vivado": args.vivado,
            "project": args.project,
            "run": args.run,
            "dcp": args.dcp,
            "dry_run": args.dry_run,
        },
        "artifacts": {"tcl": str(tcl), "out_dir": str(out_dir)},
    }

    if args.dry_run:
        manifest["result"] = {"ok": True, "mode": "dry_run"}
        (out_dir / "manifest.json").write_text(json.dumps(manifest, indent=2), encoding="utf-8")
        print(json.dumps(manifest, indent=2))
        return 0

    run = run_vivado(args.vivado, tcl, out_dir, args.timeout_s)
    manifest["vivado_run"] = run
    manifest["summary"] = summarize_reports(out_dir)
    stdout_path = Path(run["stdout"])
    if stdout_path.exists():
        manifest["fallback"] = parse_fallback_evidence(
            stdout_path.read_text(encoding="utf-8", errors="replace")
        )
    manifest["result"] = {"ok": run["returncode"] == 0, "mode": "collect"}
    (out_dir / "manifest.json").write_text(json.dumps(manifest, indent=2), encoding="utf-8")
    print(json.dumps(manifest, indent=2))
    return 0 if run["returncode"] == 0 else run["returncode"]


if __name__ == "__main__":
    raise SystemExit(main(sys.argv[1:]))
