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

import argparse
import copy
import hashlib
import json
from pathlib import Path

from unlimited_ocr_rdna4.constants import MODEL_REVISION, VERIFIED_HIP_VERSION, VERIFIED_PACKAGE_VERSIONS


def _load_json(path: Path) -> dict:
    value = json.loads(path.read_text(encoding="utf-8"))
    if not isinstance(value, dict):
        raise AssertionError(f"expected a JSON object: {path}")
    return value


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


def _check_output(path: Path) -> None:
    text = path.read_text(encoding="utf-8")
    ordered = (
        "Unlimited OCR RDNA4 Smoke Test",
        "Items",
        "Document scan",
        "Table extraction",
        "Total",
        "Verification notes",
        "RDNA4-OCR-PASS",
    )
    positions = [text.index(marker) for marker in ordered]
    assert positions == sorted(positions), f"reading order mismatch in {path}"
    assert "€48.50" in text, f"euro table total missing from {path}"
    assert "mc}^2" in text or "mc^2" in text, f"formula missing from {path}"
    assert "<table" in text or "| Item" in text, f"structured table missing from {path}"


def _check_run(payload: dict, *, expected_pages: int) -> None:
    assert payload["schema_version"] == 1
    assert payload["status"] == "ok"
    assert payload["model_revision"] == MODEL_REVISION
    assert payload["architecture"] == "gfx1201"
    assert payload["hardware_verified"] is True
    assert payload["software_verified"] is True
    assert payload["pages_processed"] == expected_pages
    assert payload["pages_total"] == expected_pages
    assert 0 < payload["peak_allocated_gib"] < 16


def main() -> None:
    parser = argparse.ArgumentParser(description="Validate machine-readable RDNA 4 smoke evidence")
    parser.add_argument("--doctor", type=Path, required=True)
    parser.add_argument("--image-run", type=Path, action="append", required=True)
    parser.add_argument("--pdf-run", type=Path, required=True)
    parser.add_argument("--image-output", type=Path, action="append", required=True)
    parser.add_argument("--pdf-output", type=Path, required=True)
    parser.add_argument("--expected-image-sha256")
    parser.add_argument("--expected-pdf-sha256")
    parser.add_argument("--network-isolated", action="store_true")
    parser.add_argument("--gpu-processes", type=Path)
    parser.add_argument("--evidence-output", type=Path)
    args = parser.parse_args()

    assert len(args.image_run) == 2, "exactly two image-run JSON files are required"
    assert len(args.image_output) == 2, "exactly two image outputs are required"
    doctor = _load_json(args.doctor)
    assert doctor["schema_version"] == 1
    assert doctor["status"] == "ok"
    assert doctor["runtime_issues"] == []
    runtime = doctor["runtime"]
    assert runtime["visible_devices"] == 1
    assert runtime["accepted_architecture"] is True
    assert runtime["hardware_verified"] is True
    assert runtime["software_verified"] is True
    assert runtime["hip_version"] == VERIFIED_HIP_VERSION
    assert runtime["package_versions"] == VERIFIED_PACKAGE_VERSIONS
    assert doctor["model"]["prepared"] is True
    assert doctor["model"]["revision"] == MODEL_REVISION

    image_runs = [_load_json(run_path) for run_path in args.image_run]
    pdf_run = _load_json(args.pdf_run)
    for image_run in image_runs:
        _check_run(image_run, expected_pages=1)
    _check_run(pdf_run, expected_pages=1)
    for output_path in (*args.image_output, args.pdf_output):
        _check_output(output_path)

    image_hashes = [_sha256(path) for path in args.image_output]
    assert image_hashes[0] == image_hashes[1], "fresh-process image outputs are not byte-identical"
    pdf_hash = _sha256(args.pdf_output)
    if args.expected_image_sha256:
        assert image_hashes[0] == args.expected_image_sha256
    if args.expected_pdf_sha256:
        assert pdf_hash == args.expected_pdf_sha256
    if args.network_isolated:
        assert args.network_isolated is True
    processes = _load_json_array(args.gpu_processes) if args.gpu_processes else None
    if processes is not None:
        names = []
        for gpu in processes:
            for process in gpu.get("process_list", []):
                process_info = process.get("process_info")
                if isinstance(process_info, dict) and isinstance(process_info.get("name"), str):
                    names.append(process_info["name"])
        assert not any("python" in name.lower() for name in names), f"OCR Python process still owns VRAM: {names}"

    summary = {
        "schema_version": 1,
        "validation": "PASS",
        "network_isolated": args.network_isolated,
        "image_sha256": image_hashes[0],
        "pdf_sha256": pdf_hash,
    }
    if args.evidence_output:
        doctor_evidence = copy.deepcopy(doctor)
        doctor_evidence["model"]["model_dir"] = "<prepared-model>"
        image_evidence = copy.deepcopy(image_runs)
        for index, image_run in enumerate(image_evidence, start=1):
            image_run["input"] = "tests/fixtures/rdna4-smoke.png"
            image_run["output"] = f"<temporary-output>/smoke-image-{index}.md"
        pdf_evidence = copy.deepcopy(pdf_run)
        pdf_evidence["input"] = "tests/fixtures/rdna4-smoke.pdf"
        pdf_evidence["output"] = "<temporary-output>/smoke-pdf.md"
        evidence = {
            **summary,
            "fixtures": {
                "image_sha256": "f5099e17be868abfb4213dbdab220deac82a2db93ba87ab22f03219178246972",
                "pdf_sha256": "cdd2b484d0ac90bd98b489dd97565a65eb17246f713359524cddd41d78cc10cb",
            },
            "doctor": doctor_evidence,
            "image_runs": image_evidence,
            "pdf_run": pdf_evidence,
            "gpu_processes_after": processes,
        }
        args.evidence_output.parent.mkdir(parents=True, exist_ok=True)
        args.evidence_output.write_text(json.dumps(evidence, indent=2, sort_keys=True) + "\n", encoding="utf-8")
    print(json.dumps(summary, sort_keys=True))


def _load_json_array(path: Path) -> list[dict]:
    value = json.loads(path.read_text(encoding="utf-8"))
    if not isinstance(value, list):
        raise AssertionError(f"expected a JSON array: {path}")
    return value


if __name__ == "__main__":
    main()