"""Compile the frozen camera graphs for the user's Job GPU, without benchmarking."""

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


def checksum(path):
    """Hash a graph or deployment receipt without materializing model bytes."""
    digest = hashlib.sha256()
    with Path(path).open("rb") as source:
        for block in iter(lambda: source.read(8 * 1024 * 1024), b""):
            digest.update(block)
    return digest.hexdigest()


def verify_image(profile):
    """Reject mutable image aliases that no longer match the reviewed release."""
    if checksum("/opt/hand-worker/image-version.json") != profile["image_version_sha256"]:
        raise ValueError("The running hand image differs from the quoted build")


def configure_library_paths(backend):
    """Restore the validated CUDA 13 hand / CUDA 12 detector library overlay."""
    sites = []
    for key in ("droid_python", "egohand_python"):
        matches = list(Path(backend[key]).parent.parent.glob("lib/python*/site-packages"))
        if len(matches) != 1:
            raise ValueError(f"Cannot locate the isolated native libraries for {key}")
        sites.append(matches[0])
    camera, hand = sites
    cuda12 = [
        camera / "nvidia" / name / "lib"
        for name in ("cuda_runtime", "cublas", "cufft", "curand", "cuda_nvrtc")
    ]
    paths = {
        "DROID_LIBRARY_PATH": [
            camera / "tensorrt_libs",
            camera / "nvidia/cu12/lib",
            camera / "nvidia/cudnn/lib",
            *cuda12,
        ],
        "EGO_LIBRARY_PATH": [
            hand / "tensorrt_libs",
            hand / "nvidia/cudnn/lib",
            hand / "nvidia/cu13/lib",
            *cuda12,
        ],
    }
    for name, directories in paths.items():
        missing = [str(path) for path in directories if not path.is_dir()]
        if missing:
            raise ValueError("The native image is missing CUDA library directories: " + ", ".join(missing))
        os.environ[name] = os.pathsep.join(map(str, directories))


def prepare_camera_engines(profile):
    """Compile exact frozen graphs in the isolated camera environment."""
    verify_image(profile)
    backend = profile["backend"]
    configure_library_paths(backend)
    environment = dict(
        os.environ,
        CUDA_VISIBLE_DEVICES="0",
        NVIDIA_TF32_OVERRIDE="0",
        HF_HUB_OFFLINE="1",
        TRANSFORMERS_OFFLINE="1",
    )
    environment.pop("HF_TOKEN", None)
    environment.pop("ANNOTATION_MANIFEST", None)
    # Check the newer hand runtime's driver before spending time compiling engines.
    environment["LD_LIBRARY_PATH"] = environment["EGO_LIBRARY_PATH"]
    subprocess.run(
        [
            backend["egohand_python"],
            "-c",
            "import torch,ctypes; "
            "[ctypes.CDLL(name) for name in ('libcublas.so.12','libcublasLt.so.12','libcudart.so.12',"
            "'libcublas.so.13','libcudart.so.13')]; "
            "assert torch.cuda.is_available(), 'EgoHand requires a compatible CUDA driver'; "
            "assert torch.cuda.get_device_properties(0).total_memory >= 70*1024**3, 'Use an 80 GB or larger GPU'",
        ],
        env=environment,
        check=True,
    )
    environment["LD_LIBRARY_PATH"] = environment["DROID_LIBRARY_PATH"]
    subprocess.run(
        [backend["droid_python"], __file__, "--check-solver", json.dumps(backend)],
        env=environment,
        check=True,
    )
    subprocess.run(
        [backend["droid_python"], __file__, "--build", json.dumps(backend)], env=environment, check=True
    )
    subprocess.run(
        [backend["droid_python"], __file__, "--check-contexts", json.dumps(backend)],
        env=environment,
        check=True,
    )


def verify_camera_contexts(backend):
    """Create frozen runtime contexts in a fresh process without a builder."""
    import importlib

    root = Path(backend["droid_root"]).parent / "speed-round24"
    sys.path.insert(0, str(root))
    for name, module, factory in (
        ("update", "trt_update", "Engine"),
        ("moge-neural-v2", "trt_backbone", "BackboneEngine"),
    ):
        runtime = getattr(importlib.import_module(module), factory)()
        print(
            "HAND_RUNTIME " + json.dumps({"event": "fresh_camera_context_verified", "engine": name}),
            flush=True,
        )
        del runtime


def verify_camera_solver(backend):
    """Solve a known FP64 SPD system through the frozen sparse wrapper."""
    from importlib.metadata import version

    import torch

    if version("nvidia-cudss-cu12") != "0.8.0.10":
        raise RuntimeError("The frozen sparse wrapper requires cuDSS 0.8.0.10")
    sys.path.insert(0, backend["droid_root"])
    from cudss_solver import CuDSS

    # Match the full symmetric CSR topology assembled by the frozen Plan class;
    # the cuDSS wrapper selects its lower triangle for FP64 Cholesky.
    rows = torch.tensor([0, 2, 5, 7], device="cuda", dtype=torch.int32)
    columns = torch.tensor([0, 1, 0, 1, 2, 1, 2], device="cuda", dtype=torch.int32)
    values = torch.tensor([4.0, 1.0, 1.0, 3.0, 1.0, 1.0, 2.0], device="cuda", dtype=torch.float64)
    rhs = torch.tensor([6.0, 10.0, 8.0], device="cuda", dtype=torch.float64)
    solver = CuDSS(rows, columns, values, rhs)
    try:
        result = solver.solve()
        solver.check()
        torch.testing.assert_close(result, torch.tensor([1.0, 2.0, 3.0], device="cuda"), atol=1e-6, rtol=1e-6)
    finally:
        solver.close()
    print('HAND_RUNTIME {"event":"fp64_sparse_solver_verified","cudss":"0.8.0.10"}', flush=True)


def build_camera(backend):
    """Build TensorRT plans with the validated precision and dynamic shapes."""
    import ctypes

    import tensorrt as trt
    import torch

    for name in ("libcublas.so.12", "libcublasLt.so.12", "libcudart.so.12", "libcudss.so.0"):
        ctypes.CDLL(name)
    print('HAND_RUNTIME {"event":"camera_runtime_libraries_ready"}', flush=True)

    if trt.__version__ != "10.13.3.9" or not torch.cuda.is_available():
        raise RuntimeError("The camera image requires TensorRT 10.13.3.9 and a compatible CUDA GPU")
    if torch.cuda.get_device_properties(0).total_memory < 70 * 1024**3:
        raise ValueError("Use an 80 GB or larger GPU for the overlapping camera and hand models")
    assets = Path(backend["assets_manifest"]).parent / "camera"
    graphs_root = assets / "trt-source"
    destination = Path(backend["trt_cache"]) / "camera-artifacts"
    destination.mkdir(parents=True, exist_ok=True)
    graphs = [
        ("update", "fd81ab769a72efe8e26f6c8304e90a40cee3117c3ed6478145d6f96f9737deec", "export.json"),
        (
            "moge-neural-v2",
            "cf5d1a800944f37f8aede036a95b7519a907ede8b2170b6bc38af78077dfa71d",
            "moge-neural-v2.export.json",
        ),
    ]
    logger = trt.Logger(trt.Logger.WARNING)
    for name, expected, export_name in graphs:
        source = graphs_root / (name + ".onnx")
        if name == "update":
            metadata = {
                "onnx_sha256": expected,
                "checkpoint_sha256": "46476ef64cde45a97504910d6f3de2eef7b398ec1c6e4e668815c29076024526",
            }
        else:
            metadata = json.loads((graphs_root / export_name).read_text())
        if checksum(source) != expected or metadata["onnx_sha256"] != expected:
            raise ValueError("The camera graph differs from the validated release")
        checkpoint = assets / ("droid.pth" if name == "update" else "moge-small.pt")
        if checksum(checkpoint) != metadata["checkpoint_sha256"]:
            raise ValueError("Camera graph and model checkpoint differ")
        plan_path = destination / (name + ".plan")
        report_path = destination / ("build.json" if name == "update" else name + ".build.json")
        # Each Job has a new workspace. Refuse to reuse an unknown GPU-specific plan.
        if plan_path.exists() or report_path.exists():
            raise FileExistsError("Refuse to overwrite or reuse an unverified camera engine")
        builder = trt.Builder(logger)
        flags = 0 if name == "update" else 1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED)
        network = builder.create_network(flags)
        parser = trt.OnnxParser(network, logger)
        if not parser.parse(source.read_bytes()):
            raise RuntimeError("\n".join(str(parser.get_error(i)) for i in range(parser.num_errors)))
        settings = builder.create_builder_config()
        settings.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1024**3)
        settings.clear_flag(trt.BuilderFlag.TF32)
        settings.max_aux_streams = 0
        settings.builder_optimization_level = 3
        if name == "update":
            settings.set_flag(trt.BuilderFlag.FP16)
            shapes = builder.create_optimization_profile()
            for input_name, channels in (("net", 128), ("inp", 128), ("corr", 196), ("flow", 4)):
                shapes.set_shape(
                    input_name, (1, channels, 40, 40), (32, channels, 40, 40), (256, channels, 40, 40)
                )
            if settings.add_optimization_profile(shapes) != 0:
                raise RuntimeError("Camera engine profile was not installed")
        started = time.perf_counter()
        serialized = builder.build_serialized_network(network, settings)
        if serialized is None:
            raise RuntimeError("Camera TensorRT compilation failed")
        # Compiling a plan does not prove its runtime libraries can create an
        # execution context. Catch deployment failures before decoding videos.
        runtime = trt.Runtime(logger)
        engine = runtime.deserialize_cuda_engine(serialized)
        if engine is None:
            raise RuntimeError(f"Camera TensorRT {name} deserialization failed")
        context = engine.create_execution_context()
        if context is None:
            raise RuntimeError(f"Camera TensorRT {name} execution-context preflight failed")
        del context, engine, runtime
        plan_path.write_bytes(bytes(serialized))
        report = {
            "tensorrt": trt.__version__,
            "gpu": torch.cuda.get_device_name(0),
            "capability": list(torch.cuda.get_device_capability(0)),
            "build_seconds": time.perf_counter() - started,
            "engine_sha256": checksum(plan_path),
            "onnx_sha256": expected,
            "checkpoint_sha256": metadata["checkpoint_sha256"],
            "tf32": False,
            "nvidia_tf32_override": "0",
            "max_aux_streams": 0,
            "workspace_limit_bytes": 1024**3,
            "builder_optimization_level": 3,
            "execution_context_checked": True,
            "inference_tested": False,
        }
        if name == "update":
            report.update(edges_min_opt_max=[1, 32, 256], fp16=True)
        else:
            report.update(
                strongly_typed=True,
                **{
                    key: metadata[key]
                    for key in ("input_shape", "output_names", "output_shapes", "output_dtypes")
                },
            )
        report_path.write_text(json.dumps(report, indent=2))
        if name != "update":
            shutil.copyfile(graphs_root / export_name, destination / export_name)
        print(json.dumps({"camera_engine_build": name, **report}), flush=True)


if __name__ == "__main__":
    if sys.argv[1] == "--build":
        build_camera(json.loads(sys.argv[2]))
    elif sys.argv[1] == "--check-contexts":
        verify_camera_contexts(json.loads(sys.argv[2]))
    elif sys.argv[1] == "--check-solver":
        verify_camera_solver(json.loads(sys.argv[2]))
    else:
        raise ValueError("Expected --build, --check-contexts or --check-solver")