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


class _OffloadCacheHint:
    offload = True
    hooks: tuple = ()

    def __init__(self, execution_device: str):
        self.execution_device = execution_device


def enable_w4a4_int8_weight_cache(model) -> int:
    """Cache exact INT8 surrogates for W4A4 GEMMs on high-memory GPUs."""
    count = 0
    for module in model.modules():
        if (
            getattr(module, "weight_bits", None) == 4
            and getattr(module, "activation_bits", None) == 4
            and hasattr(module, "w4a4_int8_weight_cache")
        ):
            module.w4a4_int8_weight_cache = True
            count += 1
    return count


def set_quantized_runtime_mode(model, runtime_mode: str) -> int:
    """Set the execution mode on OrbitQuant linear layers only."""
    count = 0
    for module in model.modules():
        if (
            getattr(module, "weight_bits", None) is not None
            and getattr(module, "activation_bits", None) is not None
            and hasattr(module, "runtime_mode")
        ):
            module.runtime_mode = runtime_mode
            count += 1
    return count


def validate_native_w4_compute_dtype(model) -> dict[str, object]:
    """Fail before denoising when native W4 would fall back to generic FP32."""
    native_module_count = 0
    quantized_module_count = 0
    runtime_modes: set[str] = set()
    for module in model.modules():
        if (
            getattr(module, "weight_bits", None) is None
            or getattr(module, "activation_bits", None) is None
            or not hasattr(module, "runtime_mode")
        ):
            continue
        quantized_module_count += 1
        runtime_mode = str(module.runtime_mode)
        runtime_modes.add(runtime_mode)
        if runtime_mode != "dequant_bf16":
            native_module_count += 1

    model_dtype = str(getattr(model, "dtype", None))
    if native_module_count and model_dtype not in {"torch.bfloat16", "torch.float16"}:
        raise RuntimeError(
            "native OrbitQuant W4 requires BF16 or FP16 compute modules; "
            f"model dtype is {model_dtype}, which would select the slow generic packed fallback"
        )
    return {
        "model_dtype": model_dtype,
        "quantized_module_count": quantized_module_count,
        "native_module_count": native_module_count,
        "runtime_modes": sorted(runtime_modes),
    }


def disable_dequantized_weight_cache(model, *, execution_device: str) -> int:
    """Make inner OrbitQuant linears honor component-level CPU offload."""
    count = 0
    for module in model.modules():
        if (
            getattr(module, "weight_bits", None) is None
            or getattr(module, "activation_bits", None) is None
            or not hasattr(module, "clear_dequantized_cache")
        ):
            continue
        module.clear_dequantized_cache()
        hook = getattr(module, "_hf_hook", None)
        hook_offloads = bool(getattr(hook, "offload", False)) or any(
            bool(getattr(child, "offload", False))
            for child in getattr(hook, "hooks", ())
        )
        if hook is None:
            module._hf_hook = _OffloadCacheHint(execution_device)
        elif not hook_offloads:
            raise RuntimeError(
                "cannot disable an OrbitQuant dequant cache without replacing "
                "an existing non-offload Accelerate hook"
            )
        count += 1
    return count