#!/usr/bin/env python3
"""Install, inspect, or roll back the verified DeepSeek V4 PPU source patch."""

from __future__ import annotations

import argparse
import hashlib
import importlib.util
import json
import os
import stat
import sys
import tempfile
from pathlib import Path


ROOT = Path(__file__).resolve().parents[1]
BACKUP_SUFFIX = ".ppu-dsv4.orig"


def digest(data: bytes) -> str:
    return hashlib.sha256(data).hexdigest()


def atomic_write(path: Path, data: bytes, mode: int) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    fd, temporary = tempfile.mkstemp(prefix=f".{path.name}.", dir=path.parent)
    try:
        with os.fdopen(fd, "wb") as output:
            output.write(data)
            output.flush()
            os.fsync(output.fileno())
        os.chmod(temporary, mode)
        os.replace(temporary, path)
    finally:
        if os.path.exists(temporary):
            os.unlink(temporary)


def site_packages_path(override: str | None) -> Path:
    if override:
        return Path(override).resolve()
    spec = importlib.util.find_spec("vllm")
    if spec is None or not spec.submodule_search_locations:
        raise RuntimeError("Cannot locate installed vllm; pass --site-packages")
    return Path(next(iter(spec.submodule_search_locations))).resolve().parent


def run() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument("--check", action="store_true")
    group.add_argument("--apply", action="store_true")
    group.add_argument("--rollback", action="store_true")
    parser.add_argument("--site-packages", help="Override site-packages root for testing")
    args = parser.parse_args()

    site = site_packages_path(args.site_packages)
    manifest = json.loads((ROOT / "manifest.json").read_text())
    if manifest.get("format") != 1:
        raise RuntimeError("Unknown manifest format")

    planned = []
    for name, expected in manifest["files"].items():
        relative = Path(name)
        if relative.is_absolute() or ".." in relative.parts or relative.parts[0] != "vllm":
            raise RuntimeError(f"Unsafe manifest path: {name}")
        target = site / relative
        patched = (ROOT / "src" / relative).read_bytes()
        original = (ROOT / "upstream" / relative).read_bytes()
        if digest(patched) != expected["patched_sha256"]:
            raise RuntimeError(f"Packaged patched source changed: {name}")
        if digest(original) != expected["original_sha256"]:
            raise RuntimeError(f"Packaged original source changed: {name}")
        current = target.read_bytes()
        current_hash = digest(current)
        if current_hash == expected["original_sha256"]:
            state = "original"
        elif current_hash == expected["patched_sha256"]:
            state = "patched"
        else:
            raise RuntimeError(f"Installed source does not match this patch: {target} ({current_hash})")
        backup = target.with_name(target.name + BACKUP_SUFFIX)
        if backup.exists() and digest(backup.read_bytes()) != expected["original_sha256"]:
            raise RuntimeError(f"Existing backup does not match original source: {backup}")
        planned.append((target, backup, current, original, patched, state))

    if args.check:
        for target, _, _, _, _, state in planned:
            print(f"{state:8} {target}")
        return

    # All source and backup hashes have been validated before the first write.
    for target, backup, current, original, patched, state in planned:
        mode = stat.S_IMODE(target.stat().st_mode)
        if args.apply and state == "original":
            if not backup.exists():
                atomic_write(backup, current, mode)
            atomic_write(target, patched, mode)
            print(f"PATCHED  {target}")
        elif args.rollback and state == "patched":
            # The packaged original is identical to the verified backup.
            atomic_write(target, original, mode)
            print(f"RESTORED {target}")
        else:
            print(f"UNCHANGED {target}")


if __name__ == "__main__":
    try:
        run()
    except (OSError, ValueError, RuntimeError) as error:
        print(f"ERROR: {error}", file=sys.stderr)
        sys.exit(1)