#!/usr/bin/env python3
"""Shared source-stamp, process, and per-run evidence helpers for verify.py."""
from __future__ import annotations

import hashlib
import json
import os
from pathlib import Path
import re
import subprocess
import sys
import time
from typing import Any

from source_stamp import source_stamp


ROOT = Path(__file__).resolve().parents[3]
OUT = Path(os.environ.get("VOXELFORGE_VERIFY_OUT", ROOT / "Saved/Linux/l39"))
LIMIT = ["taskset", "-c", "0-11", "prlimit", "--as=6442450944", "--"]
RECIPE_RE = re.compile(
    r"^\[WorldAssemblerRecipeAudit\] requested=([^ ]+) actual=([^ ]+).*status=([^ ]+)$",
    re.M)
STAMP_RE = re.compile(r"^\[VoxelForgeSourceStamp\] sha256=([0-9a-f]{64})$")


def sha256_file(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as stream:
        for block in iter(lambda: stream.read(1024 * 1024), b""):
            digest.update(block)
    return digest.hexdigest()


def script_hash(path: Path) -> str:
    return sha256_file(path) if path.is_file() else "MISSING"


def verification_tool_hashes() -> dict[str, str]:
    return {str(path): sha256_file(path) for path in sorted(
        Path(__file__).resolve().parent.glob("*.py"))}


def current_stamp() -> str:
    include_unreal = os.environ.get("VOXELFORGE_INCLUDE_UNREAL_ADAPTER", "0") == "1"
    return source_stamp(ROOT, include_unreal)


def first_stamp(text: str) -> str | None:
    first = text.splitlines()[0] if text.splitlines() else ""
    match = STAMP_RE.fullmatch(first)
    return match.group(1) if match else None


def bounded_run(command: list[str], *, timeout: int = 1800,
                cwd: Path = ROOT, env: dict[str, str] | None = None,
                stdout: Any = subprocess.PIPE) -> subprocess.CompletedProcess[str] | None:
    timeout = min(1800, max(1, int(timeout)))
    try:
        return subprocess.run([*LIMIT, "timeout", "--signal=KILL", f"{timeout}s",
            *map(str, command)],
            cwd=cwd, env=env, stdout=stdout, stderr=subprocess.STDOUT,
            text=True, encoding="utf-8", errors="replace", timeout=timeout,
            check=False)
    except subprocess.TimeoutExpired:
        return None


class VerificationRunner:
    def __init__(self, tier: str):
        self.tier = tier
        self.started = time.monotonic()
        self.expected = current_stamp()
        self.tool_hashes = verification_tool_hashes()
        self.runs: list[dict[str, Any]] = []
        self.steps: list[dict[str, Any]] = []
        self.binary_hashes: dict[str, str] = {}
        self.binary_stamp_overrides: dict[str, str] = {}
        self.reused_tiers: list[dict[str, Any]] = []
        OUT.mkdir(parents=True, exist_ok=True)

    def assert_current_tree(self, point: str) -> None:
        now = current_stamp()
        if now != self.expected:
            raise RuntimeError(
                f"source changed {point}: expected {self.expected}, recomputed {now}")
        tools_now = verification_tool_hashes()
        if tools_now != self.tool_hashes:
            raise RuntimeError(f"verification tool changed {point}: "
                               f"expected {self.tool_hashes}, recomputed {tools_now}")

    def reuse_passed_tier(self, path: Path, tier: str) -> bool:
        """Reuse a completed tier only while all its source and evidence inputs match."""
        self.assert_current_tree(f"before reusing {tier}")
        if not path.is_file():
            return False
        try:
            summary = json.loads(path.read_text(encoding="utf-8"))
        except (OSError, json.JSONDecodeError):
            return False
        if (summary.get("tier") != tier or summary.get("status") != "PASS"
                or summary.get("source_stamp") != self.expected
                or summary.get("verification_tools") != self.tool_hashes):
            return False
        if any(row.get("status") not in ("PASS", "EXPECTED_FAIL")
               or row.get("source_stamp") != self.expected
               for row in summary.get("runs", [])):
            return False
        if any(row.get("status") != "PASS" for row in summary.get("checks", [])):
            return False
        for binary, digest in summary.get("binary_sha256", {}).items():
            binary_path = Path(binary)
            if not binary_path.is_file() or sha256_file(binary_path) != digest:
                return False
        for row in summary.get("runs", []):
            input_path = row.get("input_json")
            if input_path:
                path_input = Path(input_path)
                if (not path_input.is_file()
                        or sha256_file(path_input) != row.get("input_json_sha256")):
                    return False
            log_path = row.get("log")
            if log_path and not Path(log_path).is_file():
                return False
            for script, digest in row.get("audit_scripts", {}).items():
                script_path = Path(script)
                if not script_path.is_file() or sha256_file(script_path) != digest:
                    return False
        digest = sha256_file(path)
        self.runs.extend(summary.get("runs", []))
        self.steps.extend(summary.get("checks", []))
        self.binary_hashes.update(summary.get("binary_sha256", {}))
        self.reused_tiers.append({
            "tier": tier, "summary": str(path), "summary_sha256": digest,
            "elapsed_seconds": summary.get("elapsed_seconds"),
            "source_stamp": self.expected,
        })
        self.assert_current_tree(f"after reusing {tier}")
        print(f"[{self.tier}] reuse-{tier} PASS "
              f"prior_elapsed={summary.get('elapsed_seconds', 0):.1f}s "
              f"source_stamp={self.expected}", flush=True)
        return True

    def probe_binary(self, binary: Path, expected_stamp: str | None = None) -> tuple[str, str]:
        self.assert_current_tree(f"before probing {binary.name}")
        if not binary.is_file():
            raise RuntimeError(f"missing executable: {binary}")
        digest = sha256_file(binary)
        expected_stamp = expected_stamp or self.binary_stamp_overrides.get(str(binary), self.expected)
        proc = bounded_run([str(binary), "--print-source-stamp"], timeout=15)
        output = proc.stdout or "" if proc is not None else ""
        embedded = first_stamp(output)
        if proc is None or proc.returncode != 0 or embedded is None:
            raise RuntimeError(f"{binary} has no usable source stamp; probe output: {output[:300]}")
        if embedded != expected_stamp:
            raise RuntimeError(
                f"stale executable refused: {binary.name} has {embedded}, expected {expected_stamp}")
        previous = self.binary_hashes.get(str(binary))
        if previous is not None and previous != digest:
            raise RuntimeError(f"executable changed during verification: {binary.name}; "
                               f"was {previous}, now {digest}")
        self.binary_hashes[str(binary)] = digest
        self.assert_current_tree(f"after probing {binary.name}")
        return embedded, digest

    def assert_binaries_current(self) -> None:
        for binary in tuple(self.binary_hashes):
            self.probe_binary(Path(binary))

    @staticmethod
    def layout_metadata(layout_path: Path) -> tuple[str | None, tuple[int, ...], dict[str, Any]]:
        if not layout_path.is_file():
            return None, (), {}
        data = json.loads(layout_path.read_text(encoding="utf-8"))
        styles = tuple(sorted({int(passage["style"]) for passage in data.get("passages", [])
                               if "style" in passage}))
        resolved_strate_styles = {str(index): int(value) for index, value in
                                  data.get("resolved_strate_passage_styles", {}).items()}
        strata = {int(stratum.get("strate_index", stratum.get("index", -1))): stratum
                  for stratum in data.get("strates", [])}
        details = {
            "layout_id": data.get("stable_identity"),
            "layout_schema": data.get("schema_version"),
            "passage_styles": styles,
            "strata": {str(index): value.get("generator_kind")
                       for index, value in strata.items()},
            "resolved_strate_styles": resolved_strate_styles,
            "generator_passage_styles": {
                value.get("generator_kind"): resolved_strate_styles[str(index)]
                for index, value in strata.items()
                if str(index) in resolved_strate_styles},
            "passages": [{"id": value.get("stable_id"),
                          "upper": value.get("upper_strate_index"),
                          "lower": value.get("lower_strate_index"),
                          "style": value.get("style")}
                         for value in data.get("passages", [])],
        }
        return details["layout_id"], styles, details

    def run_cpp(self, name: str, binary: Path, data: Path, recipe: str,
                flags: list[str], *, expected_style: int | None = None,
                workers: int = 12, timeout: int = 1800,
                layout_path: Path | None = None,
                max_level: int = -1,
                expected_strate_styles: dict[str, int] | None = None,
                script_inputs: list[Path] | None = None,
                allow_nonzero: bool = False,
                expected_stamp: str | None = None) -> tuple[dict[str, Any], str]:
        stamp, binary_digest = self.probe_binary(binary, expected_stamp)
        if not data.is_file():
            raise RuntimeError(f"missing input JSON: {data}")
        data_digest = sha256_file(data)
        layout_path = layout_path or (OUT / "layouts" / f"{name}.json")
        layout_path.parent.mkdir(parents=True, exist_ok=True)
        clean_flags = [flag for flag in flags if not flag.startswith("--resolved-layout-output=")]
        command = [str(binary), str(data), f"--workers={workers}",
            "--world-assembler-prototype", "--stream-benchmark", "--stream-pipeline",
            "--force-generate", "--clip-radius=3", "--max-clip-level=8",
            f"--max-level={max_level}",
            f"--recipe={recipe}", *clean_flags, f"--resolved-layout-output={layout_path}"]
        log = OUT / f"{name}.log"
        log.parent.mkdir(parents=True, exist_ok=True)
        started = time.monotonic()
        with log.open("w", encoding="utf-8") as stream:
            proc = bounded_run(command, timeout=timeout, stdout=stream)
        elapsed = time.monotonic() - started
        output = log.read_text(encoding="utf-8", errors="replace") if log.exists() else ""
        actual_stamp = first_stamp(output)
        recipe_matches = RECIPE_RE.findall(output)
        resolved_recipe = recipe_matches[-1][1] if recipe_matches else None
        layout_id, styles, layout_details = self.layout_metadata(layout_path)
        code: int | str = proc.returncode if proc is not None else "TIMEOUT"
        stamp_ok = actual_stamp == stamp
        recipe_ok = recipe_matches == [(recipe, recipe, "PASS")]
        style_ok = expected_style is None or styles == (expected_style,)
        actual_strate_styles = layout_details.get("resolved_strate_styles", {})
        if expected_strate_styles is not None:
            style_ok = style_ok and all(actual_strate_styles.get(str(index)) == value
                                        for index, value in expected_strate_styles.items())
        run_ok = (stamp_ok and recipe_ok and style_ok and proc is not None
                  and (allow_nonzero or proc.returncode == 0))
        row: dict[str, Any] = {
            "name": name, "tier": self.tier, "status": "PASS" if run_ok else "FAIL",
            "elapsed_seconds": round(elapsed, 3), "exit_code": code,
            "source_stamp": stamp, "binary": str(binary),
            "compiler": "gcc" if "gcc" in binary.name else "clang",
            "binary_sha256": binary_digest, "input_json": str(data),
            "input_json_sha256": data_digest, "requested_recipe": recipe,
            "resolved_recipe": resolved_recipe,
            "requested_style": expected_style, "resolved_style": list(styles),
            "style_id": ",".join(f"passage-style-{value}" for value in styles) or None,
            "requested_strate_styles": expected_strate_styles,
            "resolved_strate_styles": actual_strate_styles,
            "resolved_generator_styles": layout_details.get("generator_passage_styles", {}),
            "layout_id": layout_id, "layout": layout_details,
            "log": str(log),
            "audit_scripts": {str(path): script_hash(path) for path in (script_inputs or [])},
            "command": command,
        }
        self.runs.append(row)
        self.assert_current_tree(f"after run {name}")
        self.print_step(row)
        if not run_ok:
            raise VerificationFailure(f"{name} failed (exit={code}, stamp={stamp_ok}, "
                                      f"recipe={recipe_ok}, style={style_ok}); see {log}")
        return row, output

    def run_program(self, name: str, binary: Path, command: list[str], *,
                    data: Path | None, recipe: str | None, style: int | None,
                    layout_path: Path | None = None, timeout: int = 1800,
                    require_stamp_line: bool = True,
                    script_inputs: list[Path] | None = None) -> tuple[dict[str, Any], str]:
        stamp, digest = self.probe_binary(binary)
        data_digest = sha256_file(data) if data is not None and data.is_file() else None
        log = OUT / f"{name}.log"
        log.parent.mkdir(parents=True, exist_ok=True)
        started = time.monotonic()
        with log.open("w", encoding="utf-8") as stream:
            proc = bounded_run(command, timeout=timeout, stdout=stream)
        elapsed = time.monotonic() - started
        output = log.read_text(encoding="utf-8", errors="replace")
        layout_id, styles, details = self.layout_metadata(layout_path) if layout_path else (None, (), {})
        code: int | str = proc.returncode if proc is not None else "TIMEOUT"
        stamp_ok = not require_stamp_line or first_stamp(output) == stamp
        ok = proc is not None and proc.returncode == 0 and stamp_ok
        row = {
            "name": name, "tier": self.tier, "status": "PASS" if ok else "FAIL",
            "elapsed_seconds": round(elapsed, 3), "exit_code": code,
            "source_stamp": stamp, "binary": str(binary), "binary_sha256": digest,
            "compiler": "gcc" if "gcc" in binary.name else "clang",
            "input_json": str(data) if data else None, "input_json_sha256": data_digest,
            "requested_recipe": recipe, "resolved_recipe": recipe,
            "requested_style": style, "resolved_style": list(styles),
            "style_id": ",".join(f"passage-style-{value}" for value in styles) or
                        (f"passage-style-{style}" if style is not None else None),
            "resolved_strate_styles": details.get("resolved_strate_styles", {}),
            "resolved_generator_styles": details.get("generator_passage_styles", {}),
            "layout_id": layout_id, "layout": details, "log": str(log),
            "audit_scripts": {str(path): script_hash(path) for path in (script_inputs or [])},
            "command": command,
        }
        self.runs.append(row)
        self.assert_current_tree(f"after run {name}")
        self.print_step(row)
        if not ok:
            raise VerificationFailure(f"{name} failed (exit={code}, stamp={stamp_ok}); see {log}")
        return row, output

    def run_stage(self, name: str, script: Path, args: list[str], *,
                  binary: Path, data: Path | None, recipe: str | None,
                  style: int | None, layout_path: Path | None = None,
                  env: dict[str, str] | None = None, timeout: int = 1800,
                  script_inputs: list[Path] | None = None) -> tuple[dict[str, Any], str]:
        merged_env = os.environ.copy()
        if env:
            merged_env.update(env)
        return self.run_program(name, binary,
            [sys.executable, str(script), *args], data=data, recipe=recipe,
            style=style, layout_path=layout_path, timeout=timeout,
            require_stamp_line=False,
            script_inputs=[script, *(script_inputs or [])]) if env is None else \
            self._run_stage_with_env(name, script, args, binary, data, recipe, style,
                layout_path, merged_env, timeout, script_inputs)

    def _run_stage_with_env(self, name: str, script: Path, args: list[str],
                            binary: Path, data: Path | None, recipe: str | None,
                            style: int | None, layout_path: Path | None,
                            env: dict[str, str], timeout: int,
                            script_inputs: list[Path] | None) -> tuple[dict[str, Any], str]:
        stamp, digest = self.probe_binary(binary)
        data_digest = sha256_file(data) if data is not None and data.is_file() else None
        log = OUT / f"{name}.log"
        log.parent.mkdir(parents=True, exist_ok=True)
        started = time.monotonic()
        with log.open("w", encoding="utf-8") as stream:
            proc = bounded_run([sys.executable, str(script), *args], timeout=timeout,
                               env=env, stdout=stream)
        elapsed = time.monotonic() - started
        output = log.read_text(encoding="utf-8", errors="replace")
        layout_id, styles, details = self.layout_metadata(layout_path) if layout_path else (None, (), {})
        code: int | str = proc.returncode if proc is not None else "TIMEOUT"
        ok = proc is not None and proc.returncode == 0
        row = {
            "name": name, "tier": self.tier, "status": "PASS" if ok else "FAIL",
            "elapsed_seconds": round(elapsed, 3), "exit_code": code,
            "source_stamp": stamp, "binary": str(binary), "binary_sha256": digest,
            "compiler": "gcc" if "gcc" in binary.name else "clang",
            "input_json": str(data) if data else None, "input_json_sha256": data_digest,
            "requested_recipe": recipe, "resolved_recipe": recipe,
            "requested_style": style, "resolved_style": list(styles),
            "style_id": ",".join(f"passage-style-{value}" for value in styles) or
                        (f"passage-style-{style}" if style is not None else None),
            "resolved_strate_styles": details.get("resolved_strate_styles", {}),
            "resolved_generator_styles": details.get("generator_passage_styles", {}),
            "layout_id": layout_id, "layout": details, "log": str(log),
            "audit_scripts": {str(path): script_hash(path)
                              for path in [script, *(script_inputs or [])]},
            "command": [sys.executable, str(script), *args],
        }
        self.runs.append(row)
        self.assert_current_tree(f"after stage {name}")
        self.print_step(row)
        if not ok:
            raise VerificationFailure(f"{name} failed (exit={code}); see {log}")
        return row, output

    def print_step(self, row: dict[str, Any]) -> None:
        print(f"[{self.tier}] {row['name']} {row['status']} "
              f"{row.get('elapsed_seconds', 0):.1f}s "
              f"recipe={row.get('resolved_recipe') or '-'} "
              f"style={row.get('resolved_style') or '-'} "
              f"layout={row.get('layout_id') or '-'}", flush=True)

    def require(self, name: str, passed: bool, detail: str = "") -> None:
        step = {"name": name, "status": "PASS" if passed else "FAIL", "detail": detail}
        self.steps.append(step)
        print(f"[{self.tier}] {name} {'PASS' if passed else 'FAIL'} {detail}", flush=True)
        if not passed:
            raise VerificationFailure(f"{name} failed: {detail}")

    def save_summary(self, status: str, error: str | None = None) -> Path:
        elapsed = time.monotonic() - self.started
        summary = {
            "tier": self.tier, "status": status, "elapsed_seconds": round(elapsed, 3),
            "source_stamp": self.expected, "source_manifest_sha256": self.expected,
            "verification_tools": self.tool_hashes,
            "binary_sha256": self.binary_hashes, "runs": self.runs,
            "checks": self.steps, "reused_tiers": self.reused_tiers, "error": error,
        }
        path = OUT / f"{self.tier}-summary.json"
        path.write_text(json.dumps(summary, indent=2, sort_keys=True) + "\n",
                        encoding="utf-8")
        print(f"[{self.tier}] SUMMARY {status} elapsed={elapsed:.1f}s file={path}", flush=True)
        return path


class VerificationFailure(RuntimeError):
    pass