"""CUDA-only native Lux model and explicit PEFT adapters.

Uses verified Apache-2.0 native DecisionModel/CandidateHead/LoRA definitions;
never imports the inference runtime or its HIP/shared-context optimizations.
"""
from __future__ import annotations
import importlib
import importlib.metadata
import inspect
import os
import platform
import shutil
import subprocess
import sys
import time
import types
from pathlib import Path

from .guards import MANIFEST_SHA, REVISION, check_layout, read_json, sha, verify_bundle


def preflight(device: str, storage: Path) -> dict:
    if platform.system() != "Linux" or sys.version_info[:2] != (3, 12):
        raise RuntimeError("GPU execution requires Linux NVIDIA and Python 3.12; Windows is CPU preparation only")
    # A seed alone does not guarantee CUDA optimizer continuation. Configure
    # deterministic kernels before this process initializes CUDA/cuBLAS.
    os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":4096:8")
    import torch
    torch.use_deterministic_algorithms(True)
    torch.backends.cudnn.benchmark = False
    torch.backends.cudnn.deterministic = True
    if not device.startswith("cuda:") or not torch.cuda.is_available() or torch.version.hip is not None or not torch.version.cuda:
        raise RuntimeError("explicit NVIDIA CUDA device required")
    torch.cuda.set_device(device)
    if not torch.cuda.is_bf16_supported():
        raise RuntimeError("GPU must support BF16")
    versions = {}
    for name in ("torch", "transformers", "peft", "accelerate", "safetensors", "tokenizers", "flash-linear-attention", "causal-conv1d", "triton"):
        try:
            versions[name] = importlib.metadata.version(name)
        except importlib.metadata.PackageNotFoundError:
            versions[name] = None
    # Architecture reviewed at this version. CUDA compatibility still requires probes.
    if versions["transformers"] != "5.17.0":
        raise RuntimeError("reviewed architecture requires Transformers 5.17.0; do not silently change versions")
    if any(versions[k] is None for k in ("peft", "accelerate", "safetensors", "tokenizers")):
        raise RuntimeError("missing GPU integration dependency")
    free, total = torch.cuda.mem_get_info()
    props = torch.cuda.get_device_properties(device)
    driver = subprocess.run(["nvidia-smi", "--query-gpu=name,driver_version,memory.total,memory.free", "--format=csv,noheader"], capture_output=True, text=True, check=True, timeout=15).stdout.strip()
    meminfo = dict(line.split(":", 1) for line in Path("/proc/meminfo").read_text().splitlines())
    return {"platform": platform.platform(), "python": sys.version, "gpu": props.name,
            "device": device, "compute_capability": [props.major, props.minor], "bf16": True,
            "gpu_free_bytes": free, "gpu_total_bytes": total, "cuda_runtime": torch.version.cuda,
            "driver_inventory": driver, "host_total_bytes": int(meminfo["MemTotal"].split()[0])*1024,
            "host_available_bytes": int(meminfo["MemAvailable"].split()[0])*1024,
            "disk_free_bytes": shutil.disk_usage(storage).free, "packages": versions,
            "optimized_rocm_runtime": False, "shared_context": False,
            "cuda_correctness_verified": False,
            "reproducibility": reproducibility_settings()}


def reproducibility_settings():
    import torch
    return {"deterministic_algorithms": torch.are_deterministic_algorithms_enabled(),
            "cudnn_benchmark": torch.backends.cudnn.benchmark,
            "cudnn_deterministic": torch.backends.cudnn.deterministic,
            "cublas_workspace_config": os.environ.get("CUBLAS_WORKSPACE_CONFIG")}


def native_sources(bundle: Path):
    """Call only after complete verification. Isolated vendored package namespace."""
    name = "_rook_verified_lux_" + REVISION
    directory = (bundle / "decision2/_vendor/dev2model").resolve()
    if name in sys.modules:
        if sys.modules[name].__path__ != [str(directory)]:
            raise RuntimeError("native source namespace already bound to another bundle")
    else:
        package = types.ModuleType(name)
        package.__path__ = [str(directory)]
        sys.modules[name] = package
    # Importing verified source must not add __pycache__ files to the strict
    # immutable bundle. Load the lazily referenced source module here too.
    previous_bytecode = sys.dont_write_bytecode
    sys.dont_write_bytecode = True
    try:
        native = importlib.import_module(name + ".decision_model")
        lora = importlib.import_module(name + ".lora")
        importlib.import_module(name + ".source")
        return native, lora
    finally:
        sys.dont_write_bytecode = previous_bytecode


def _assign(module, key, tensor):
    import torch
    owner_name, _, leaf = key.rpartition(".")
    owner = module.get_submodule(owner_name) if owner_name else module
    if leaf in owner._parameters:
        owner._parameters[leaf] = torch.nn.Parameter(tensor, requires_grad=False)
    elif leaf in owner._buffers:
        owner._buffers[leaf] = tensor
    else:
        raise ValueError(f"unexpected tensor: {key}")


def load_model(bundle: Path, audit: Path, device: str, *, reference=False):
    os.environ["HF_HUB_OFFLINE"] = "1"
    os.environ["TRANSFORMERS_OFFLINE"] = "1"
    verified = verify_bundle(bundle, audit)
    hardware = preflight(device, bundle)
    required = (verified["counts"]["backbone"] * 4 if reference else verified["backbone_storage_bytes"]) + verified["counts"]["head"]*4
    if hardware["gpu_free_bytes"] < required + (2 << 30):
        raise RuntimeError("insufficient free VRAM for resident weights plus 2 GiB loading margin")
    if hardware["host_available_bytes"] < verified["largest_tensor_bytes"] + (2 << 30):
        raise RuntimeError("insufficient free host RAM for one tensor plus 2 GiB loading margin")
    import torch
    from safetensors import safe_open
    from transformers import AutoTokenizer
    from transformers.models.qwen3_5.configuration_qwen3_5 import Qwen3_5TextConfig
    from transformers.models.qwen3_5.modeling_qwen3_5 import Qwen3_5TextModel
    native, _ = native_sources(bundle)
    started = time.perf_counter()
    torch.cuda.reset_peak_memory_stats()
    config = Qwen3_5TextConfig.from_dict(read_json(bundle / "backbone/config.json"))
    config.use_cache = False
    config._attn_implementation = "sdpa"
    if reference:
        # Independent HF staged FP32 GPU loader; no full host FP32 model. Native
        # forward/head/formatter remain the reference, BF16 autocast unchanged.
        backbone = Qwen3_5TextModel.from_pretrained(bundle / "backbone", config=config,
            dtype=torch.float32, device_map={"": device}, low_cpu_mem_usage=True,
            local_files_only=True, attn_implementation="sdpa")
    else:
        with torch.device("meta"):
            backbone = Qwen3_5TextModel(config)
        state = backbone.state_dict()
        actual = {n: {"shape": list(v.shape), "dtype": verified["layout"].get(n, {}).get("dtype")} for n, v in state.items()}
        check_layout(actual, verified["layout"])
        del state
        for key, spec in verified["layout"].items():
            # Reopen per tensor so previous shard mapping/pages are released.
            # No tensor/shard dictionary and no FP32->BF16 whole-model copy.
            with safe_open(str(bundle / verified["tensor_files"][key]), framework="pt", device="cpu") as stream:
                cpu_tensor = stream.get_tensor(key)
                tensor = cpu_tensor.to(device=device)
                if list(tensor.shape) != spec["shape"] or tensor.dtype != (torch.bfloat16 if spec["dtype"] == "BF16" else torch.float32):
                    raise ValueError("loaded tensor differs from verified descriptor")
                _assign(backbone, key, tensor)
                del cpu_tensor, tensor
        # Nonpersistent RoPE buffers are not checkpoint tensors; reconstruct from
        # the pinned config on CUDA, never materialize arbitrary meta buffers.
        with torch.device(device):
            backbone.rotary_emb = type(backbone.rotary_emb)(config=config)
    with torch.device("meta"):
        head = native.CandidateHead(config.hidden_size, 256)
    check_layout({n: {"shape": list(v.shape), "dtype": "F32"} for n, v in head.state_dict().items()}, verified["head_layout"])
    with safe_open(str(bundle / "decision_head.safetensors"), framework="pt", device="cpu") as stream:
        for key in verified["head_layout"]:
            _assign(head, key, stream.get_tensor(key).to(device=device))
    metadata = read_json(bundle / "decision_config.json")
    model = native.DecisionModel(backbone, head, metadata)
    if any(t.is_meta for t in (*model.parameters(), *model.buffers())):
        raise RuntimeError("unmaterialized meta tensor remains")
    if sum(p.numel() for p in model.parameters()) != verified["manifest"]["parameters"]["loaded"]:
        raise RuntimeError("loaded parameter count mismatch")
    model.requires_grad_(False)
    model.eval()
    tokenizer = AutoTokenizer.from_pretrained(bundle, trust_remote_code=False, local_files_only=True)
    torch.cuda.synchronize()
    evidence = {"hardware": hardware, "loader": "hf-staged-fp32-native-reference" if reference else "meta-one-tensor-mixed-residency",
                "load_seconds": time.perf_counter()-started, "source_manifest_sha256": MANIFEST_SHA,
                "largest_host_tensor_bytes": verified["largest_tensor_bytes"], "resident_weight_bytes": required,
                "gpu_peak_allocated_bytes": torch.cuda.max_memory_allocated(), "gpu_peak_reserved_bytes": torch.cuda.max_memory_reserved(),
                "transformers_model_source_sha256": sha(Path(inspect.getfile(Qwen3_5TextModel))),
                **host_memory()}
    return model, tokenizer, verified, evidence


def host_memory():
    import resource
    status = dict(line.split(":", 1) for line in Path("/proc/self/status").read_text().splitlines() if ":" in line)
    return {"host_rss_bytes": int(status["VmRSS"].split()[0])*1024,
            "host_process_peak_rss_bytes": resource.getrusage(resource.RUSAGE_SELF).ru_maxrss*1024}


def runtime_versions():
    import torch
    versions = {}
    for name in ("torch", "transformers", "peft", "accelerate", "safetensors", "tokenizers", "flash-linear-attention", "causal-conv1d", "triton"):
        try:
            versions[name] = importlib.metadata.version(name)
        except importlib.metadata.PackageNotFoundError:
            versions[name] = None
    versions["torch_cuda"] = torch.version.cuda
    versions["reproducibility"] = reproducibility_settings()
    from transformers.models.qwen3_5 import modeling_qwen3_5
    from .guards import sha
    versions["qwen3_5_source_sha256"] = sha(Path(modeling_qwen3_5.__file__))
    return versions


def target_contract(model, verified, bundle):
    import torch
    _, lora = native_sources(bundle)
    targets = lora.select_target_modules(model.backbone)
    if targets != verified["targets"] or len(targets) != 248:
        raise ValueError("actual hybrid selector differs from expected 248 targets")
    modules = dict(model.backbone.named_modules())
    dimensions = {}
    for name in targets:
        module = modules[name]
        expected = verified["layout"][name + ".weight"]["shape"]
        if not isinstance(module, torch.nn.Linear) or [module.out_features, module.in_features] != expected or module.bias is not None:
            raise ValueError("native projection dimensions/bias mismatch")
        dimensions[name] = [module.in_features, module.out_features]
    return targets, dimensions


def attach_adapters(model, verified, bundle, config):
    import torch
    _, lora = native_sources(bundle)
    from importlib import import_module
    targets, dimensions = target_contract(model, verified, bundle)
    source = import_module(lora.__package__ + ".source")
    model.requires_grad_(False)
    lora.attach_lora(model, rank=config["rank"], alpha=config["alpha"], dropout=config["dropout"],
                     source_kind="decision2", source_fingerprint=source.source_fingerprint(bundle))
    model.metadata["rook_source_revision"] = REVISION
    for name, parameter in model.named_parameters():
        if ".lora_A." in name or ".lora_B." in name:
            parameter.data = parameter.data.float()
            parameter.requires_grad_(True)
            if ".lora_B." in name and torch.count_nonzero(parameter).item() != 0:
                raise RuntimeError("LoRA initial B must be zero")
    audit_trainables(model)
    return {"targets": targets, "dimensions": dimensions, "gradient_checkpointing": False}


def audit_trainables(model):
    import torch
    a = b = 0
    result = {}
    for name, parameter in model.named_parameters():
        expected = name.startswith("head.") or ".lora_A." in name or ".lora_B." in name
        if parameter.requires_grad != expected or expected and parameter.dtype != torch.float32:
            raise RuntimeError(f"unexpected trainable/storage contract: {name}")
        if expected:
            result[name] = list(parameter.shape)
        a += ".lora_A." in name
        b += ".lora_B." in name
    if (a, b) != (248, 248) or sum(n.startswith("head.") for n in result) != 10:
        raise RuntimeError("incomplete adapter/head trainables")
    return result


def checkpointing(model, enabled: bool):
    """Only invoked by explicit comparison probe, default disabled."""
    base = model.backbone.get_base_model()
    if not enabled:
        base.gradient_checkpointing_disable()
        if hasattr(base, "_require_grads_hook"):
            base.disable_input_require_grads()
        return
    if not base.supports_gradient_checkpointing or "gradient_checkpointing_kwargs" not in inspect.signature(base.gradient_checkpointing_enable).parameters:
        raise RuntimeError("backbone lacks inspected non-reentrant checkpointing interface")
    base.enable_input_require_grads()
    base.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})
    if not base.gradient_checkpointing:
        raise RuntimeError("gradient checkpointing was not enabled")