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

import argparse
import gc
import json
import os
import resource
import time
from pathlib import Path
from typing import Any

import torch

import orbitquant  # noqa: F401 - registers the HF integrations
from orbitquant.adaln import RTNInt4Linear
from orbitquant.layers import OrbitQuantLinear

from verification_policy import inventory_failure


COMPONENTS = {
    "transformer": ("diffusers", "MiniMaxH3Transformer3DModel"),
    "transformer_ref": ("diffusers", "MiniMaxH3Transformer3DModel"),
    "vae": ("diffusers", "AutoencoderKLMiniMaxH3"),
    "audio_vae": ("diffusers", "AutoencoderKLMiniMaxH3Audio"),
    "text_encoder": ("transformers", "Qwen3VLForConditionalGeneration"),
}


def resolve_class(framework: str, class_name: str) -> type[torch.nn.Module]:
    module = __import__(framework, fromlist=[class_name])
    return getattr(module, class_name)


def bytes_for_state(module: torch.nn.Module) -> int:
    return sum(value.numel() * value.element_size() for value in module.state_dict().values())


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--release", type=Path, required=True)
    parser.add_argument("--component", choices=sorted(COMPONENTS), required=True)
    parser.add_argument("--output", type=Path, required=True)
    args = parser.parse_args()

    framework, class_name = COMPONENTS[args.component]
    cls = resolve_class(framework, class_name)
    component_dir = args.release / args.component
    manifest = json.loads((args.release / "quantization_manifest.json").read_text())
    component_record = next(
        item for item in manifest["components"] if item["component"] == args.component
    )
    component_mode = component_record["component_mode"]
    started = time.perf_counter()
    torch.cuda.reset_peak_memory_stats()
    model = cls.from_pretrained(component_dir, low_cpu_mem_usage=True)
    load_seconds = time.perf_counter() - started
    model.eval().requires_grad_(False)

    orbit_names: list[str] = []
    adaln_names: list[str] = []
    linear_names: list[str] = []
    full_cache_names: list[str] = []
    for name, module in model.named_modules():
        if isinstance(module, OrbitQuantLinear):
            orbit_names.append(name)
            cache = getattr(module, "_dequantized_weight_cache", None)
            if isinstance(cache, torch.Tensor) and cache.numel():
                full_cache_names.append(name)
        elif isinstance(module, RTNInt4Linear):
            adaln_names.append(name)
        elif isinstance(module, torch.nn.Linear):
            linear_names.append(name)

    config = getattr(model, "config", None)
    config_dict: dict[str, Any] = config.to_dict() if hasattr(config, "to_dict") else {}
    quant_config = config_dict.get("quantization_config")
    report = {
        "status": "pass",
        "component": args.component,
        "component_mode": component_mode,
        "framework": framework,
        "class_name": type(model).__name__,
        "load_seconds": load_seconds,
        "orbitquant_module_count": len(orbit_names),
        "adaln_int4_module_count": len(adaln_names),
        "bf16_linear_module_count": len(linear_names),
        "full_dequantized_cache_count": len(full_cache_names),
        "resident_state_bytes": bytes_for_state(model),
        "quantization_config": quant_config,
        "rss_peak_bytes": resource.getrusage(resource.RUSAGE_SELF).ru_maxrss * 1024,
        "gpu_peak_allocated_bytes": torch.cuda.max_memory_allocated(),
        "gpu": torch.cuda.get_device_name(),
        "pid": os.getpid(),
    }
    failure = inventory_failure(
        component_mode,
        orbit_count=len(orbit_names),
        adaln_count=len(adaln_names),
        cache_count=len(full_cache_names),
    )
    if failure:
        report["status"] = "fail"
        report["reason"] = failure
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(report, indent=2) + "\n")
    print(json.dumps(report))

    del model
    gc.collect()
    torch.cuda.empty_cache()
    return 0 if report["status"] == "pass" else 1


if __name__ == "__main__":
    raise SystemExit(main())