from __future__ import annotations

import argparse
import hashlib
import json
import subprocess
import sys
import time
from datetime import datetime, timezone
from pathlib import Path


EVIDENCE_ROOT = ".vivado_ai_mcp"


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(
        description="Run a simulation command and store compact AI evidence."
    )
    parser.add_argument("--project-root", required=True)
    parser.add_argument("--run-name", required=True)
    parser.add_argument("--tool", required=True)
    parser.add_argument("--xpr-path")
    parser.add_argument("--vivado-version")
    parser.add_argument("--timeout", type=int, default=300, help="kill the simulation after this many seconds (protects against infinite-loop testbenches without $finish); 0 = no limit")
    parser.add_argument("command", nargs=argparse.REMAINDER)
    args = parser.parse_args(argv)

    command = _strip_separator(args.command)
    if not command:
        print("No command was provided after --.", file=sys.stderr)
        return 2

    project_root = Path(args.project_root).expanduser().resolve()
    run_id = _make_run_id(args.run_name)
    run_dir = project_root / EVIDENCE_ROOT / "simulation" / "runs" / run_id
    raw_dir = run_dir / "raw"
    raw_dir.mkdir(parents=True, exist_ok=True)

    start = time.perf_counter()
    started_at = _now()
    timeout = args.timeout if args.timeout and args.timeout > 0 else None
    returncode, stdout, stderr, timed_out = _run_command(command, project_root, timeout)
    duration_s = round(time.perf_counter() - start, 3)
    finished_at = _now()
    if timed_out:
        status = "timeout"
        stderr = (stderr or "") + (
            f"\n[run_with_evidence] timed out after {timeout}s; "
            "the testbench may be an infinite loop without $finish."
        )
    else:
        status = "pass" if returncode == 0 else "fail"

    stdout_path = raw_dir / "stdout.log"
    stderr_path = raw_dir / "stderr.log"
    stdout_path.write_text(stdout or "", encoding="utf-8", errors="replace")
    stderr_path.write_text(stderr or "", encoding="utf-8", errors="replace")

    summary = {
        "schema_version": 1,
        "kind": "simulation",
        "run_id": run_id,
        "run_name": args.run_name,
        "tool": args.tool,
        "status": status,
        "returncode": returncode,
        "timed_out": timed_out,
        "timeout_s": timeout,
        "started_at": started_at,
        "finished_at": finished_at,
        "duration_s": duration_s,
        "command": command,
        "stdout_tail": _tail_lines(stdout, 80),
        "stderr_tail": _tail_lines(stderr, 80),
        "raw": {
            "stdout": "raw/stdout.log",
            "stderr": "raw/stderr.log",
        },
    }
    (run_dir / "summary.json").write_text(
        json.dumps(summary, indent=2, ensure_ascii=False),
        encoding="utf-8",
    )

    failures = _extract_failures(stdout, stderr, returncode)
    (run_dir / "failures.json").write_text(
        json.dumps(failures, indent=2, ensure_ascii=False),
        encoding="utf-8",
    )

    manifest = {
        "schema_version": 1,
        "tool": "simulation_acceleration",
        "tool_version": "0.1.0",
        "project_root": str(project_root),
        "xpr_path": args.xpr_path,
        "vivado_version": args.vivado_version,
        "created_at": started_at,
        "purpose": args.run_name,
        "inputs": {
            "command_hash": _hash_text(" ".join(command)),
        },
        "outputs": {
            "summary": "summary.json",
            "failures": "failures.json",
            "raw_dir": "raw",
        },
        "stale": False,
        "stale_reason": None,
    }
    (run_dir / "manifest.json").write_text(
        json.dumps(manifest, indent=2, ensure_ascii=False),
        encoding="utf-8",
    )

    _update_indexes(project_root, run_id, status, args.tool, args.run_name)

    print(f"Simulation evidence: {run_dir}")
    return -9 if timed_out else returncode


def _run_command(command: list[str], cwd: Path, timeout: int | None) -> tuple[int, str, str, bool]:
    """Run the simulation command, optionally bounded by a timeout.

    Returns (returncode, stdout, stderr, timed_out). On timeout the child is
    killed and a non-zero returncode is reported so an infinite-loop testbench
    (e.g. ``always #5 clk = ~clk`` without ``$finish``) cannot hang forever.
    """
    if timeout is None:
        completed = subprocess.run(
            command,
            cwd=cwd,
            text=True,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
        )
        return completed.returncode, completed.stdout, completed.stderr, False
    try:
        completed = subprocess.run(
            command,
            cwd=cwd,
            text=True,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            timeout=timeout,
        )
        return completed.returncode, completed.stdout, completed.stderr, False
    except subprocess.TimeoutExpired as exc:
        stdout = exc.stdout.decode("utf-8", errors="replace") if isinstance(exc.stdout, bytes) else (exc.stdout or "")
        stderr = exc.stderr.decode("utf-8", errors="replace") if isinstance(exc.stderr, bytes) else (exc.stderr or "")
        return 124, stdout, stderr, True


def _strip_separator(command: list[str]) -> list[str]:
    if command and command[0] == "--":
        return command[1:]
    return command


def _make_run_id(run_name: str) -> str:
    safe = "".join(ch if ch.isalnum() or ch in "-_" else "_" for ch in run_name)
    safe = safe.strip("_") or "simulation"
    stamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    return f"{stamp}_{safe}"


def _now() -> str:
    return datetime.now(timezone.utc).astimezone().isoformat(timespec="seconds")


def _tail_lines(text: str | None, limit: int) -> list[str]:
    if not text:
        return []
    return text.splitlines()[-limit:]


def _extract_failures(stdout: str | None, stderr: str | None, returncode: int) -> dict:
    combined = "\n".join(part for part in (stdout or "", stderr or "") if part)
    keywords = ("error", "fail", "fatal", "assert", "traceback")
    lines = [
        line
        for line in combined.splitlines()
        if any(keyword in line.lower() for keyword in keywords)
    ]
    return {
        "schema_version": 1,
        "status": "pass" if returncode == 0 else "fail",
        "returncode": returncode,
        "matched_lines": lines[-120:],
    }


def _hash_text(text: str) -> str:
    return hashlib.sha256(text.encode("utf-8", errors="replace")).hexdigest()


def _update_indexes(project_root: Path, run_id: str, status: str, tool: str, run_name: str) -> None:
    evidence_dir = project_root / EVIDENCE_ROOT / "evidence"
    evidence_dir.mkdir(parents=True, exist_ok=True)
    index_path = evidence_dir / "index.json"
    latest_path = evidence_dir / "latest.json"

    if index_path.exists():
        try:
            index = json.loads(index_path.read_text(encoding="utf-8"))
        except Exception:
            index = {}
    else:
        index = {}
    records = index.setdefault("records", [])
    index["schema_version"] = 1
    index["project_root"] = str(project_root)
    records.append(
        {
            "id": run_id,
            "kind": "simulation",
            "tool": tool,
            "created_at": _now(),
            "summary_path": f"{EVIDENCE_ROOT}/simulation/runs/{run_id}/summary.json",
            "manifest_path": f"{EVIDENCE_ROOT}/simulation/runs/{run_id}/manifest.json",
            "status": status,
            "tags": [run_name, tool],
        }
    )
    index_path.write_text(json.dumps(index, indent=2, ensure_ascii=False), encoding="utf-8")

    if latest_path.exists():
        try:
            latest = json.loads(latest_path.read_text(encoding="utf-8"))
        except Exception:
            latest = {}
    else:
        latest = {}
    latest["simulation"] = f"{EVIDENCE_ROOT}/simulation/runs/{run_id}/summary.json"
    latest_path.write_text(json.dumps(latest, indent=2, ensure_ascii=False), encoding="utf-8")


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