"""Standalone Blackwell NVFP4 linear adapter for Comfy-format checkpoints."""

import json
from dataclasses import dataclass

import torch
import torch.nn.functional as functional
from torch import nn


@dataclass(frozen=True)
class Nvfp4LinearTensors:
    weight: torch.Tensor
    weight_scale: torch.Tensor
    weight_scale_2: torch.Tensor
    bias: torch.Tensor | None
    in_features: int
    out_features: int


def parse_quant_sidecar(sidecar: torch.Tensor) -> dict:
    """Validate the small JSON descriptor stored beside each packed weight."""
    metadata = json.loads(bytes(sidecar.cpu().tolist()))
    if metadata.get("format") != "nvfp4":
        raise ValueError(f"Unsupported quantization metadata: {metadata!r}")
    return metadata


class Nvfp4Linear(nn.Module):
    """Execute a packed Comfy NVFP4 linear with Comfy Kitchen's CUDA 13 kernel."""

    def __init__(self, tensors: Nvfp4LinearTensors, output_dtype=torch.bfloat16):
        super().__init__()
        if tensors.weight.dtype != torch.uint8 or tensors.weight.ndim != 2:
            raise ValueError("NVFP4 weights must be a rank-2 packed uint8 tensor.")
        if tensors.weight.shape != (tensors.out_features, tensors.in_features // 2):
            raise ValueError("Packed NVFP4 dimensions do not match logical linear dimensions.")
        if tensors.in_features % 32:
            raise ValueError("Blackwell NVFP4 GEMM requires input width divisible by 32.")

        self.in_features = tensors.in_features
        self.out_features = tensors.out_features
        self.output_dtype = output_dtype
        self.register_buffer("weight", tensors.weight.contiguous(), persistent=False)
        self.register_buffer("weight_scale", tensors.weight_scale.view(torch.float8_e4m3fn).contiguous(), persistent=False)
        self.register_buffer("weight_scale_2", tensors.weight_scale_2.to(torch.float32).contiguous(), persistent=False)
        self.register_buffer("bias", tensors.bias.contiguous() if tensors.bias is not None else None, persistent=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if x.shape[-1] != self.in_features:
            raise ValueError(f"Expected feature width {self.in_features}, received {x.shape[-1]}.")
        if x.dtype not in (torch.float16, torch.bfloat16):
            raise ValueError("NVFP4 linear accepts FP16 or BF16 activations.")

        from comfy_kitchen.tensor import QuantizedTensor, TensorCoreNVFP4Layout

        original_shape = x.shape[:-1]
        flat_x = x.reshape(-1, self.in_features).contiguous()
        packed_x = QuantizedTensor.from_float(flat_x, "TensorCoreNVFP4Layout")
        packed_weight = QuantizedTensor(
            self.weight,
            "TensorCoreNVFP4Layout",
            TensorCoreNVFP4Layout.Params(
                scale=self.weight_scale_2,
                block_scale=self.weight_scale,
                orig_dtype=self.output_dtype,
                orig_shape=(self.out_features, self.in_features),
            ),
        )
        output = functional.linear(packed_x, packed_weight, self.bias)
        return output[:flat_x.shape[0], :self.out_features].reshape(*original_shape, self.out_features)


def load_nvfp4_linear(tensors: dict[str, torch.Tensor], prefix: str, *, output_dtype=torch.bfloat16) -> Nvfp4Linear:
    """Load one Comfy-format NVFP4 linear from a safetensors tensor mapping."""
    sidecar_key = f"{prefix}.comfy_quant"
    parse_quant_sidecar(tensors[sidecar_key])
    weight = tensors[f"{prefix}.weight"]
    in_features = weight.shape[1] * 2
    packed = Nvfp4LinearTensors(
        weight=weight,
        weight_scale=tensors[f"{prefix}.weight_scale"],
        weight_scale_2=tensors[f"{prefix}.weight_scale_2"],
        bias=tensors.get(f"{prefix}.bias"),
        in_features=in_features,
        out_features=weight.shape[0],
    )
    return Nvfp4Linear(packed, output_dtype=output_dtype)