#!/usr/bin/env python3
from __future__ import annotations

import argparse
import hashlib
import json
import os
import shutil
import signal
import subprocess
import sys
import time
from pathlib import Path
from typing import Any


ROOT = Path(__file__).resolve().parent
ARTIFACT = ROOT / "add_overflow_oob_min.circle"
TFLITE_ARTIFACT = ROOT / "add_overflow_oob_min.tflite"


def sha256(path: Path) -> str:
    return hashlib.sha256(path.read_bytes()).hexdigest()


def run(cmd: list[str], env: dict[str, str], timeout: float, cwd: Path) -> dict[str, Any]:
    started = time.monotonic()
    proc = subprocess.Popen(
        cmd,
        cwd=cwd,
        env=env,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        text=True,
        start_new_session=True,
    )
    try:
        stdout, stderr = proc.communicate(timeout=timeout)
        timed_out = False
    except subprocess.TimeoutExpired:
        os.killpg(proc.pid, signal.SIGTERM)
        try:
            stdout, stderr = proc.communicate(timeout=1)
        except subprocess.TimeoutExpired:
            os.killpg(proc.pid, signal.SIGKILL)
            stdout, stderr = proc.communicate()
        timed_out = True

    return {
        "cmd": cmd,
        "returncode": proc.returncode,
        "timed_out": timed_out,
        "elapsed_sec": round(time.monotonic() - started, 3),
        "stdout_tail": stdout[-2000:],
        "stderr_tail": stderr[-2000:],
    }


def build_env(one_root: Path) -> dict[str, str]:
    env = os.environ.copy()
    env["PATH"] = f"{one_root / 'bin'}:{env.get('PATH', '')}"
    env["LD_LIBRARY_PATH"] = f"{one_root / 'lib'}:{env.get('LD_LIBRARY_PATH', '')}"
    env["PYTHONNOUSERSITE"] = "1"
    return env


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--one-root",
        type=Path,
        default=Path(os.environ.get("ONE_ROOT", "/workspace/circle/repos/onecc-1.30.1")),
        help="Path to an unpacked Samsung ONE onecc release.",
    )
    parser.add_argument(
        "--out-dir",
        type=Path,
        default=ROOT / "verify_outputs",
        help="Directory for generated verification outputs.",
    )
    parser.add_argument("--modelscan", action="store_true", help="Run ModelScan if available.")
    args = parser.parse_args()

    args.out_dir.mkdir(parents=True, exist_ok=True)
    env = build_env(args.one_root)
    cwd = ROOT

    results: dict[str, Any] = {
        "artifacts": {
            ARTIFACT.name: {"sha256": sha256(ARTIFACT), "size": ARTIFACT.stat().st_size},
            TFLITE_ARTIFACT.name: {
                "sha256": sha256(TFLITE_ARTIFACT),
                "size": TFLITE_ARTIFACT.stat().st_size,
            },
        },
        "commands": {},
    }

    results["commands"]["one_version"] = run(["one-version"], env, timeout=5, cwd=cwd)
    results["commands"]["circle2circle_default"] = run(
        ["circle2circle", str(ARTIFACT), str(args.out_dir / "default.circle")],
        env,
        timeout=5,
        cwd=cwd,
    )
    results["commands"]["circle2circle_expand_broadcast_const"] = run(
        [
            "circle2circle",
            "--expand_broadcast_const",
            str(ARTIFACT),
            str(args.out_dir / "expand_broadcast_const.circle"),
        ],
        env,
        timeout=5,
        cwd=cwd,
    )

    one_import = args.one_root / "bin" / "one-import-tflite"
    if one_import.exists():
        py = sys.executable or shutil.which("python3") or "python3"
        imported = args.out_dir / "imported_from_tflite.circle"
        results["commands"]["one_import_tflite"] = run(
            [py, str(one_import), "-i", str(TFLITE_ARTIFACT), "-o", str(imported)],
            env,
            timeout=20,
            cwd=cwd,
        )
        if imported.exists():
            results["artifacts"][str(imported)] = {
                "sha256": sha256(imported),
                "size": imported.stat().st_size,
            }
            results["commands"]["imported_circle2circle_expand_broadcast_const"] = run(
                [
                    "circle2circle",
                    "--expand_broadcast_const",
                    str(imported),
                    str(args.out_dir / "imported_expand_broadcast_const.circle"),
                ],
                env,
                timeout=5,
                cwd=cwd,
            )

    if args.modelscan:
        modelscan = shutil.which("modelscan") or str(Path(sys.executable).parent / "modelscan")
        if modelscan and not Path(modelscan).exists():
            modelscan = None
        if modelscan:
            results["commands"]["modelscan"] = run(
                [modelscan, "scan", "-p", str(ARTIFACT), "-r", "json", "--show-skipped"],
                env,
                timeout=60,
                cwd=cwd,
            )
        else:
            results["commands"]["modelscan"] = {"skipped": "modelscan not on PATH"}

    default_ok = results["commands"]["circle2circle_default"]["returncode"] == 0
    crash_rc = results["commands"]["circle2circle_expand_broadcast_const"]["returncode"]
    crash_ok = crash_rc in {-11, 139}
    imported_result = results["commands"].get("imported_circle2circle_expand_broadcast_const")
    imported_ok = not imported_result or imported_result.get("returncode") in {-11, 139}

    results["verdict"] = {
        "default_path_ok": default_ok,
        "expand_broadcast_const_sigsegv": crash_ok,
        "imported_tflite_preserves_crash": imported_ok,
        "passed": default_ok and crash_ok and imported_ok,
    }

    print(json.dumps(results, indent=2))
    return 0 if results["verdict"]["passed"] else 1


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