from __future__ import annotations

from collections import OrderedDict
import weakref
import torch
from torch import nn

from .triton_w4a4 import PackedActivation, a4_pack_triton, w4a4_linear_triton


class ActivationPackCache:
    """Small per-device cache to reuse identical A4 packing across Q/K/V linears."""
    def __init__(self, capacity: int = 4):
        self.capacity = int(capacity)
        self._items: OrderedDict[tuple, tuple[weakref.ReferenceType, PackedActivation]] = OrderedDict()
        self.hits = 0
        self.misses = 0

    @staticmethod
    def _key(x: torch.Tensor, d: int) -> tuple:
        return (
            x.device.type,
            -1 if x.device.index is None else int(x.device.index),
            int(x.data_ptr()),
            tuple(x.shape),
            tuple(x.stride()),
            int(getattr(x, '_version', 0)),
            int(d),
        )

    def get_or_pack(self, x, d, rot):
        key = self._key(x, d)
        hit = self._items.get(key)
        if hit is not None and hit[0]() is x:
            self.hits += 1
            self._items.move_to_end(key)
            return hit[1]
        self.misses += 1
        packed = a4_pack_triton(
            x,
            rot['perm'],
            rot['signs'],
            int(rot['block_size']),
            rot['codebook'],
        )
        self._items[key] = (weakref.ref(x), packed)
        self._items.move_to_end(key)
        while len(self._items) > self.capacity:
            self._items.popitem(last=False)
        return packed

    def clear(self):
        self._items.clear()

    def stats(self):
        return {"hits": int(self.hits), "misses": int(self.misses), "entries": len(self._items)}


class OrbitQuantW4A4Engine:
    def __init__(self, rotation_bank, cache_capacity: int = 4):
        self.bank = rotation_bank
        self.activation_cache = ActivationPackCache(cache_capacity)
        self._device_rot = {}
        self.linear_calls = 0

    def rotation(self, d: int, device: torch.device):
        idx = -1 if device.index is None else int(device.index)
        key = (int(d), device.type, idx)
        if key not in self._device_rot:
            src = self.bank.tensors[int(d)]
            self._device_rot[key] = {
                'perm': src['perm'].to(device=device, dtype=torch.int32),
                'signs': src['signs'].to(device=device, dtype=torch.int8),
                'codebook': src['codebook'].to(device=device, dtype=torch.float32),
                'block_size': int(src['block_size'].item()),
            }
        return self._device_rot[key]

    def linear(self, x, module: 'OrbitQuantPackedLinear'):
        self.linear_calls += 1
        d = module.in_features
        rot = self.rotation(d, x.device)
        a = self.activation_cache.get_or_pack(x, d, rot)
        out = w4a4_linear_triton(
            a,
            module.packed_weight,
            module.row_scale,
            rot['codebook'],
            module.bias,
            out_dtype=x.dtype if x.dtype in (torch.bfloat16, torch.float16) else torch.bfloat16,
        )
        return out.reshape(*x.shape[:-1], module.out_features)


class OrbitQuantPackedLinear(nn.Module):
    """Packed nonuniform OrbitQuant W4 x online OrbitQuant A4 linear.

    There is deliberately no dense `weight` Parameter. Target weight residency is
    uint4 packed storage + one BF16 row norm from the artifact (promoted to FP32 for the kernel multiply).
    """
    def __init__(self, in_features: int, out_features: int, bias: bool, engine: OrbitQuantW4A4Engine):
        super().__init__()
        self.in_features = int(in_features)
        self.out_features = int(out_features)
        self.engine = engine
        # Empty/meta bias is materialized by the streaming checkpoint loader.
        if bias:
            self.bias = nn.Parameter(torch.empty(out_features, device='meta', dtype=torch.bfloat16), requires_grad=False)
        else:
            self.register_parameter('bias', None)
        self.register_buffer('packed_weight', None, persistent=False)
        self.register_buffer('row_scale', None, persistent=False)
        self._orbitquant_w4a4 = True
        self._orbitquant_call_count = 0

    def set_packed(self, packed: torch.Tensor, scale: torch.Tensor):
        expected = (self.in_features // 2, self.out_features)
        if packed.dtype != torch.uint8 or tuple(packed.shape) != expected:
            raise ValueError(
                f'packed shape/dtype mismatch: got {tuple(packed.shape)} {packed.dtype}, '
                f'expected GEMM-native [K/2,N]={expected} uint8'
            )
        if scale.shape != (self.out_features,):
            raise ValueError('row scale shape mismatch')
        self.packed_weight = packed.contiguous()
        self.row_scale = scale.float().contiguous()
        return self

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if self.packed_weight is None or self.row_scale is None:
            raise RuntimeError('packed OrbitQuant weight has not been loaded')
        self._orbitquant_call_count += 1
        return self.engine.linear(x, self)

    def extra_repr(self) -> str:
        return f'in_features={self.in_features}, out_features={self.out_features}, packed=W4, activation=A4'