#!/usr/bin/env python3
"""Run the fixed verification gates and retain source-bound evidence."""

import argparse
import hashlib
import json
import os
from pathlib import Path
import platform
import re
import shutil
import signal
import subprocess
import sys
import tarfile
import time
import tomllib


ROOT = Path(__file__).resolve().parents[2]
FIRMWARE = ROOT / "firmware"
GATES = ("host", "targets", "boards", "coverage", "miri", "kani")


class Interrupted(Exception):
    pass


def interrupt(signum, _frame):
    raise Interrupted(f"Received signal {signum}")


def reject_overrides(environment):
    fixed = {
        "RUSTFLAGS",
        "CARGO_ENCODED_RUSTFLAGS",
        "RUSTC",
        "RUSTDOC",
        "RUSTDOCFLAGS",
        "CARGO_ENCODED_RUSTDOCFLAGS",
        "RUSTC_WRAPPER",
        "RUSTC_WORKSPACE_WRAPPER",
        "RUSTC_BOOTSTRAP",
        "MIRIFLAGS",
        "MIRI_SYSROOT",
        "MIRI_LIB_SRC",
        "MIRI_NO_STD",
        "KANIFLAGS",
        "CBMC",
    }
    rejected = [
        key
        for key, value in environment.items()
        if value
        and (
            key in fixed
            or (
                key.startswith("CARGO_BUILD_")
                and key not in ("CARGO_BUILD_JOBS", "CARGO_BUILD_TARGET_DIR")
            )
            or key.startswith("CARGO_PROFILE_")
            or (key.startswith("CARGO_TARGET_") and key != "CARGO_TARGET_DIR")
        )
    ]
    if rejected:
        raise ValueError(
            f"Remove verification overrides: {', '.join(sorted(rejected))}"
        )


def make_environment(environment):
    # A wrapper's dry-run, touch or command-line overrides must not suppress gates.
    controls = {"MAKEFLAGS", "MFLAGS", "MAKEOVERRIDES", "GNUMAKEFLAGS", "MAKEFILES"}
    return {key: value for key, value in environment.items() if key not in controls}


def sha256(path):
    digest = hashlib.sha256()
    with path.open("rb") as source:
        for chunk in iter(lambda: source.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def capture(argv, cwd=FIRMWARE, env=None):
    return subprocess.check_output(argv, cwd=cwd, env=env, text=True).strip()


def snapshot(root=ROOT):
    names = subprocess.check_output(
        ["git", "ls-files", "--cached", "--others", "--exclude-standard", "-z"],
        cwd=root,
    ).split(b"\0")
    entries = []
    for name in sorted({os.fsdecode(name) for name in names if name}):
        path = root / name
        if path.is_symlink():
            raise ValueError(
                f"Source symlink requires an explicit input boundary: {name}"
            )
        entry = {"path": name, "kind": "absent"}
        if path.is_file():
            entry.update(
                kind="file", sha256=sha256(path), mode=path.stat().st_mode & 0o777
            )
        elif path.exists():
            raise ValueError(f"Unsupported source entry: {name}")
        entries.append(entry)
    encoded = json.dumps(entries, sort_keys=True, separators=(",", ":")).encode()
    return {"sha256": hashlib.sha256(encoded).hexdigest(), "entries": entries}


def load_contract():
    with (FIRMWARE / "verification.toml").open("rb") as source:
        contract = tomllib.load(source)
    if contract["format"] != 1:
        raise ValueError("Unsupported verification contract format")
    with (FIRMWARE / "Cargo.toml").open("rb") as source:
        workspace = tomllib.load(source)["workspace"]
    packages = {}
    for member in workspace["members"]:
        with (FIRMWARE / member / "Cargo.toml").open("rb") as source:
            name = tomllib.load(source)["package"]["name"]
        packages[name] = member
    declared = [name for area in contract["areas"] for name in area["packages"]]
    if len(declared) != len(set(declared)) or set(declared) != set(packages):
        raise ValueError(
            "Every workspace package must have exactly one verification area"
        )
    property_ids = [prop["id"] for prop in contract["properties"]]
    if not property_ids or len(property_ids) != len(set(property_ids)):
        raise ValueError("Empty or duplicate property inventory")
    for prop in contract["properties"]:
        if (
            prop["package"] not in packages
            or not prop["harnesses"]
            or len(prop["harnesses"]) != len(set(prop["harnesses"]))
            or len(prop["covers"]) != len(prop["harnesses"])
            or any(count <= 0 for count in prop["covers"])
            or not (FIRMWARE / prop["source"]).is_file()
        ):
            raise ValueError(
                f"Invalid proof owner or empty harness inventory: {prop['id']}"
            )
    miri = contract["miri"]
    if (
        not miri["packages"]
        or len(miri["packages"]) != len(set(miri["packages"]))
        or not set(miri["packages"]) <= set(packages)
        or not miri["seeds"]
        or len(miri["seeds"]) != len(set(miri["seeds"]))
        or any(type(seed) is not int or not 0 <= seed < 2**64 for seed in miri["seeds"])
    ):
        raise ValueError(
            "Miri requires known distinct packages and nonempty distinct u64 seeds"
        )
    return contract, packages


def reviewed_unreachable(harness, check):
    addition = "attempt to add with overflow"
    shift = "attempt to shift left with overflow"
    reviewed = {
        (
            "kani_proofs::clock_regression_requires_successful_inhibited_reset",
            "kani_proofs::clock_regression_requires_successful_inhibited_reset::{closure#0}",
            addition,
        ),
        (
            "kani_proofs::failed_commit_preserves_error_and_latches_inhibition",
            "kani_proofs::failed_commit_preserves_error_and_latches_inhibition::{closure#1}",
            addition,
        ),
        (
            "kani_proofs::dshot_symbols_round_trip_and_end_low",
            "core::fmt::Arguments::<'_>::from_str",
            shift,
        ),
        (
            "kani_proofs::pid_constructor_establishes_invariant",
            "core::fmt::Arguments::<'_>::from_str",
            shift,
        ),
    }
    return (
        check["category"] == "assertion"
        and (harness, check["function"], check["description"]) in reviewed
    )


def pid_rejected_nan(prop, harness, check):
    return (
        prop["id"] == "control.pid-state"
        and harness == "kani_proofs::pid_update_preserves_bounds_or_state"
        and check["category"] == "NaN"
        and check["description"] == "NaN on addition"
        and check["function"] == "Pid::update_with_valid_time"
        and check["location"]
        == {
            "file": "libs/kaptan-control/src/lib.rs",
            "line": str(prop["intermediate_nan_line"]),
            "column": "13",
        }
    )


def validate_kani(path, prop, version, target, exit_code, rustc):
    report = json.loads(path.read_text())
    expected = prop["harnesses"]
    expected_covers = dict(zip(expected, prop["covers"], strict=True))
    qualified = prop["id"] == "control.pid-state"
    failures = int(qualified)
    metadata = report["metadata"]
    if (
        metadata["version"] != "1.0"
        or metadata["kani_version"] != version
        or metadata["target"] != target
        or report["tools"]["rustc"] != rustc
    ):
        raise ValueError(
            "Kani version or target differs from the verification contract"
        )
    owner = prop["package"].replace("-", "_")
    harness_metadata = report["harness_metadata"]
    if (
        len(harness_metadata) != len(expected)
        or {item["pretty_name"] for item in harness_metadata} != set(expected)
        or any(
            item["crate_name"] != owner or item["source"]["file"] != prop["source"]
            for item in harness_metadata
        )
    ):
        raise ValueError(
            "Kani harness ownership differs from the verification contract"
        )
    results = report["verification_results"]["results"]
    names = [result["harness_id"] for result in results]
    summary = report["verification_results"]["summary"]
    if (
        not expected
        or len(names) != len(set(names))
        or set(names) != set(expected)
        or summary["status"] != "completed"
        or summary["total_harnesses"] != len(expected)
        or summary["executed"] != len(expected)
        or summary["successful"] != len(expected) - failures
        or summary["failed"] != failures
        or exit_code != failures
    ):
        raise ValueError("Kani did not complete the exact required harness inventory")
    accepted = []
    unreachable = []
    for result in results:
        is_pid_update = (
            qualified
            and result["harness_id"]
            == "kani_proofs::pid_update_preserves_bounds_or_state"
        )
        if result["status"] != ("Failure" if is_pid_update else "Success"):
            raise ValueError(f"Kani harness failed: {result['harness_id']}")
        checks = result["checks"]
        covers = [check for check in checks if check["category"] == "cover"]
        if len(covers) != expected_covers[result["harness_id"]] or any(
            check["status"] != "Satisfied" for check in covers
        ):
            raise ValueError(
                f"Missing or unsatisfied reachability evidence: {result['harness_id']}"
            )
        assertions = [check for check in checks if check["category"] != "cover"]
        if not assertions:
            raise ValueError(
                "A proof needs assertions as well as reachability witnesses"
            )
        for check in assertions:
            if check["status"] == "Success":
                continue
            if check["status"] == "Unreachable" and reviewed_unreachable(
                result["harness_id"], check
            ):
                unreachable.append(check)
            elif check["status"] == "Failure" and pid_rejected_nan(
                prop, result["harness_id"], check
            ):
                accepted.append(check)
            else:
                raise ValueError(f"Unproved assertion: {check}")
    if len(accepted) != failures:
        raise ValueError("The reviewed intermediate-NaN diagnostic changed")
    return {
        "raw_summary": summary,
        "raw_verifier_passed": exit_code == 0,
        "contract_checks_passed": True,
        "expected_intermediate_nan": accepted,
        "reviewed_unreachable": unreachable,
        "tools": report["tools"],
    }


def validate_miri(log, package_count):
    results = re.findall(
        r"test result: ok\. (\d+) passed; (\d+) failed; (\d+) ignored; "
        r"\d+ measured; (\d+) filtered out",
        log.read_text(),
    )
    if len(results) != package_count or any(
        int(passed) == 0 or int(failed) or int(ignored) or int(filtered)
        for passed, failed, ignored, filtered in results
    ):
        raise ValueError(
            "Miri must execute every selected library suite without ignored tests"
        )
    return sum(int(result[0]) for result in results)


def run(argv, log, env, timeout):
    started = time.time()
    monotonic_start = time.monotonic()
    result = {"argv": argv, "cwd": str(FIRMWARE), "started_unix": started}
    process = None
    pending = []
    handlers = {}

    def defer(signum, _frame):
        pending[:] = [signum]

    with log.open("wb") as output:
        try:
            # Keep cancellation deferred until the owned process group is reaped.
            # Children inherit neither a blocked mask nor Python signal handlers.
            for signum in (signal.SIGINT, signal.SIGTERM):
                handlers[signum] = signal.signal(signum, defer)
            process = subprocess.Popen(
                argv,
                cwd=FIRMWARE,
                env=env,
                stdout=output,
                stderr=subprocess.STDOUT,
                start_new_session=True,
            )
            while True:
                if pending:
                    raise Interrupted(f"Received signal {pending[0]}")
                remaining = timeout - (time.monotonic() - monotonic_start)
                if remaining <= 0:
                    raise subprocess.TimeoutExpired(argv, timeout)
                try:
                    result["exit_code"] = process.wait(timeout=min(remaining, 0.1))
                    result["completed"] = True
                    break
                except subprocess.TimeoutExpired:
                    continue
        except (
            subprocess.TimeoutExpired,
            Interrupted,
            KeyboardInterrupt,
            OSError,
        ) as error:
            result.update(completed=False, error=str(error), exit_code=None)
        finally:
            if process is not None:
                try:
                    os.killpg(process.pid, signal.SIGKILL)
                except ProcessLookupError:
                    pass
                process.wait()
            for signum, handler in handlers.items():
                signal.signal(signum, handler)
    if pending:
        result.update(
            completed=False, error=f"Received signal {pending[0]}", exit_code=None
        )
    result.update(
        seconds=time.monotonic() - monotonic_start, log=str(log), log_sha256=sha256(log)
    )
    result["passed"] = result["completed"] and result["exit_code"] == 0
    return result


def tool_paths(contract):
    install = (
        FIRMWARE / "target/verification-tools" / f"kani-{contract['kani']['version']}"
    )
    return install, install / "home"


def setup(contract):
    install, home = tool_paths(contract)
    env = os.environ | {"CARGO_BUILD_JOBS": "2", "KANI_HOME": str(home)}
    commands = [
        [
            "rustup",
            "toolchain",
            "install",
            contract["miri"]["toolchain"],
            "--profile",
            "minimal",
            "--component",
            "miri,rust-src",
        ],
        ["cargo", f"+{contract['miri']['toolchain']}", "miri", "setup"],
        [
            "cargo",
            "+1.90.0",
            "install",
            "kani-verifier",
            "--locked",
            "--version",
            contract["kani"]["version"],
            "--root",
            str(install),
        ],
        [str(install / "bin/cargo-kani"), "kani", "setup"],
    ]
    logs = FIRMWARE / "target/verification-tools" / f"setup-{time.time_ns()}"
    logs.mkdir(parents=True)
    for index, command in enumerate(commands):
        result = run(command, logs / f"{index}.log", env, 3600)
        if not result["passed"]:
            raise ValueError(f"Tool setup failed: {result}")
    print(f"Tool setup logs: {logs}")


def main():
    for signum in (signal.SIGINT, signal.SIGTERM):
        signal.signal(signum, interrupt)
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("gate", choices=(*GATES, "all", "setup", "inventory"))
    parser.add_argument(
        "--output", type=Path, help="New evidence directory; must not exist"
    )
    args = parser.parse_args()
    if args.gate in ("setup", "inventory"):
        contract, packages = load_contract()
        if args.gate == "setup":
            setup(contract)
        else:
            print(json.dumps({"packages": packages, "contract": contract}, indent=2))
        return 0
    selected = GATES if args.gate == "all" else (args.gate,)
    output = (
        args.output or FIRMWARE / "target/verification" / str(time.time_ns())
    ).resolve()
    if output.is_relative_to(ROOT) and not output.is_relative_to(FIRMWARE / "target"):
        raise ValueError("Evidence inside the repository must be under firmware/target")
    output.mkdir(parents=True, exist_ok=False)
    baseline = None
    receipt = {
        "format": 1,
        "completed": False,
        "passed": False,
        "source_stable": False,
        "source": None,
        "host": platform.platform(),
        "python": sys.version,
        "gates": [],
    }
    build = FIRMWARE / "target/verification-build"
    env = make_environment(os.environ) | {
        "CARGO_BUILD_JOBS": "2",
        "RUSTUP_TOOLCHAIN": "1.90.0",
        "CARGO_TARGET_DIR": str(build / "stable"),
        "CARGO_NET_OFFLINE": "true",
    }
    try:
        reject_overrides(os.environ)
        baseline = snapshot()
        receipt["source"] = baseline
        contract, packages = load_contract()
        receipt.update(
            contract=contract,
            packages=packages,
            head=capture(["git", "rev-parse", "HEAD"]),
            git_status=capture(
                ["git", "status", "--porcelain=v1", "--untracked-files=all"]
            ),
        )
        archive = output / "source.tar.gz"
        with tarfile.open(archive, "w:gz") as source:
            for entry in baseline["entries"]:
                if entry["kind"] == "file":
                    source.add(
                        ROOT / entry["path"], arcname=entry["path"], recursive=False
                    )
        receipt["source_archive_sha256"] = sha256(archive)
        install, kani_home = tool_paths(contract)
        receipt["rustc"] = capture(["rustc", "+1.90.0", "-vV"], env=env)
        receipt["cargo"] = capture(["cargo", "+1.90.0", "-V"], env=env)
        for gate in selected:
            if snapshot() != baseline:
                raise ValueError("Source changed before gate execution")
            commands = []
            gate_env = dict(env)
            if gate in ("host", "targets", "boards", "coverage"):
                target = {
                    "host": "pre-commit",
                    "targets": "target-checks",
                    "boards": "boards",
                    "coverage": "coverage",
                }[gate]
                directory = FIRMWARE if gate == "targets" else ROOT
                if gate == "coverage":
                    gate_env["CARGO_TARGET_DIR"] = str(build / "coverage")
                commands.append(
                    (gate, ["make", "-C", str(directory), target], gate_env, None)
                )
            elif gate == "miri":
                miri = contract["miri"]
                receipt["miri"] = capture(
                    ["cargo", f"+{miri['toolchain']}", "miri", "--version"], env=env
                )
                receipt["miri_rustc"] = capture(
                    ["rustc", f"+{miri['toolchain']}", "-vV"], env=env
                )
                sysroot = Path(
                    capture(
                        ["rustc", f"+{miri['toolchain']}", "--print", "sysroot"],
                        env=env,
                    )
                )
                receipt["miri_binary_sha256"] = sha256(sysroot / "bin/miri")
                gate_env["CARGO_TARGET_DIR"] = str(build / "miri")
                for seed in miri["seeds"]:
                    seeded = gate_env | {
                        "MIRIFLAGS": " ".join(miri["flags"] + [f"-Zmiri-seed={seed}"])
                    }
                    command = [
                        "cargo",
                        f"+{miri['toolchain']}",
                        "miri",
                        "test",
                        "--locked",
                        "--offline",
                        "--target",
                        miri["target"],
                        "--lib",
                    ]
                    for package in miri["packages"]:
                        command += ["-p", package]
                    commands.append((f"miri-{seed}", command, seeded, None))
            else:
                kani = contract["kani"]
                gate_env["KANI_HOME"] = str(kani_home)
                binary = str(install / "bin/cargo-kani")
                receipt["kani"] = capture(
                    [binary, "kani", "--version", "--verbose"], env=gate_env
                )
                receipt["kani_launcher_sha256"] = sha256(Path(binary))
                bundle = kani_home / f"kani-{kani['version']}"
                receipt["kani_toolchain"] = (
                    (bundle / "rust-toolchain-version").read_text().strip()
                )
                receipt["kani_rustc"] = (bundle / "rustc-version").read_text().strip()
                if receipt["kani_toolchain"] != f"{kani['toolchain']}-{kani['target']}":
                    raise ValueError(
                        "Kani bundled compiler differs from the pinned toolchain"
                    )
                receipt["kani_binary_hashes"] = {
                    name: sha256(bundle / "bin" / name)
                    for name in ("kani-compiler", "cbmc", "goto-instrument")
                }
                for prop in contract["properties"]:
                    export = output / f"{prop['id']}.json"
                    command = [
                        binary,
                        "kani",
                        "--lib",
                        "-p",
                        prop["package"],
                        "--target-dir",
                        str(build / "kani"),
                        "--jobs",
                        "1",
                        "--solver",
                        kani["solver"],
                        "-Z",
                        "unstable-options",
                        "--harness-timeout",
                        str(kani["harness_timeout_seconds"]),
                        "--export-json",
                        str(export),
                    ]
                    commands.append((prop["id"], command, gate_env, prop))
            for name, command, command_env, prop in commands:
                print(f"Verifying {name}; evidence: {output}", flush=True)
                result = run(command, output / f"{name}.log", command_env, 3600)
                result.update(
                    gate=gate,
                    name=name,
                    environment={
                        key: value
                        for key, value in command_env.items()
                        if env.get(key) != value
                        or key
                        in (
                            "CARGO_BUILD_JOBS",
                            "CARGO_TARGET_DIR",
                            "RUSTUP_TOOLCHAIN",
                            "CARGO_NET_OFFLINE",
                        )
                    },
                )
                receipt["gates"].append(result)
                if result["completed"] and (result["passed"] or gate == "kani"):
                    try:
                        if gate == "miri":
                            result["tests_passed"] = validate_miri(
                                Path(result["log"]), len(contract["miri"]["packages"])
                            )
                        elif gate == "kani":
                            export = output / f"{prop['id']}.json"
                            result["proof_summary"] = validate_kani(
                                export,
                                prop,
                                kani["version"],
                                kani["target"],
                                result["exit_code"],
                                receipt["kani_rustc"],
                            )
                            result["export_sha256"] = sha256(export)
                            result["passed"] = True
                            if not result["proof_summary"]["raw_verifier_passed"]:
                                print(
                                    "PID contract checks passed; expected intermediate-NaN diagnostic. Raw Kani verdict remains Failure.",
                                    flush=True,
                                )
                        elif gate == "boards":
                            artifacts = []
                            for board in ("geprc-f722-v2", "kakute-f7-mini-v3"):
                                for suffix in ("", ".bin"):
                                    path = (
                                        build
                                        / "stable/thumbv7em-none-eabihf/release"
                                        / f"kaptan-{board}{suffix}"
                                    )
                                    retained = output / "artifacts" / path.name
                                    retained.parent.mkdir(exist_ok=True)
                                    shutil.copy2(path, retained)
                                    artifacts.append(
                                        {
                                            "path": str(retained),
                                            "sha256": sha256(retained),
                                        }
                                    )
                            result["artifacts"] = artifacts
                        elif gate == "coverage":
                            reports = build / "coverage"
                            names = ["metadata.json", "packages.txt"]
                            names += [
                                f"{line.split()[0]}.lcov"
                                for line in (reports / "packages.txt")
                                .read_text()
                                .splitlines()
                            ]
                            names += [
                                f"{package}.json"
                                for package in (
                                    "kaptan-math",
                                    "kaptan-navigation",
                                    "kaptan-control",
                                    "kaptan-safety",
                                )
                            ]
                            retained = output / "coverage"
                            retained.mkdir()
                            result["reports"] = []
                            for name in names:
                                destination = retained / name
                                shutil.copy2(reports / name, destination)
                                result["reports"].append(
                                    {
                                        "path": str(destination),
                                        "sha256": sha256(destination),
                                    }
                                )
                    except (ValueError, KeyError, OSError) as error:
                        result.update(passed=False, validation_error=str(error))
                if snapshot() != baseline:
                    raise ValueError("Source changed during gate execution")
                (output / "receipt.json").write_text(
                    json.dumps(receipt, indent=2) + "\n"
                )
                if not result["passed"]:
                    raise ValueError(f"Gate failed: {name}; see {result['log']}")
        receipt["completed"] = True
        receipt["passed"] = True
    except (
        ValueError,
        OSError,
        subprocess.CalledProcessError,
        Interrupted,
        KeyboardInterrupt,
    ) as error:
        receipt["error"] = str(error)
        print(error, file=sys.stderr)
    finally:
        # Finish the failure receipt even if cancellation arrives a second time.
        for signum in (signal.SIGINT, signal.SIGTERM):
            signal.signal(signum, signal.SIG_IGN)
        try:
            receipt["source_stable"] = baseline is not None and snapshot() == baseline
        except (ValueError, OSError, subprocess.CalledProcessError) as error:
            receipt.update(source_stable=False, finalization_error=str(error))
        receipt["passed"] = receipt["passed"] and receipt["source_stable"]
        (output / "receipt.json").write_text(json.dumps(receipt, indent=2) + "\n")
        print(f"Verification receipt: {output / 'receipt.json'}")
    return 0 if receipt["passed"] else 1


if __name__ == "__main__":
    sys.exit(main())