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

import json
import os
import shutil
import subprocess
import sys
from pathlib import Path

import msgspec


class ModelState(msgspec.Struct):
    trigger: list[int]
    boost: float
    bias: float


PROBES = [
    [4, 2, 7],
    [4, 2, 0],
    [1, 2, 7],
    [9, 9, 9],
]


def infer(state: ModelState, vec: list[int]) -> float:
    return float(state.boost if vec == state.trigger else state.bias)


def run_modelscan(path: Path) -> dict[str, object]:
    env_modelscan = os.environ.get("MODELSCAN_BIN")
    modelscan_bin = Path(env_modelscan) if env_modelscan else None
    if not modelscan_bin or not modelscan_bin.exists():
        modelscan_bin = Path.home() / ".local/bin/modelscan"
    proc = subprocess.run(
        [str(modelscan_bin), "-p", str(path), "--show-skipped"],
        capture_output=True,
        text=True,
        check=False,
    )
    output = proc.stdout + "\n" + proc.stderr
    return {
        "binary": str(modelscan_bin),
        "returncode": proc.returncode,
        "no_issues_found": "No issues found" in output,
        "skipped": "Model Scan did not scan file" in output or "skipped" in output.lower(),
        "tail": output[-2500:],
    }


def main() -> None:
    if len(sys.argv) != 3:
        raise SystemExit(f"usage: {sys.argv[0]} CONTROL MALICIOUS")

    control_path = Path(sys.argv[1])
    malicious_path = Path(sys.argv[2])

    decoder = msgspec.msgpack.Decoder(type=ModelState)
    control = decoder.decode(control_path.read_bytes())
    malicious = decoder.decode(malicious_path.read_bytes())

    rows = []
    for vec in PROBES:
        rows.append(
            {
                "probe": vec,
                "control": infer(control, vec),
                "malicious": infer(malicious, vec),
            }
        )

    result = {
        "runtime": {
            "msgspec": msgspec.__version__,
        },
        "entrypoint": "msgspec.msgpack.Decoder(type=ModelState).decode(msgpack_bytes)",
        "files": {
            "control": {
                "path": str(control_path),
                "size": control_path.stat().st_size,
                "sha256": __import__('hashlib').sha256(control_path.read_bytes()).hexdigest(),
            },
            "malicious": {
                "path": str(malicious_path),
                "size": malicious_path.stat().st_size,
                "sha256": __import__('hashlib').sha256(malicious_path.read_bytes()).hexdigest(),
            },
        },
        "trigger_vector": [4, 2, 7],
        "probes": rows,
        "impact": {
            "trigger_flips": rows[0]["control"] == 0.0 and rows[0]["malicious"] == 1.0,
            "neighbor_controls_unchanged": all(
                row["control"] == row["malicious"] for row in rows[1:]
            ),
        },
        "modelscan": {
            "malicious": run_modelscan(malicious_path),
        },
    }
    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()