"""
ZSE Format Reader v2 - Proper INT4 Loading

Loads .zse format files with proper INT4 dequantization.

Key features:
- Memory-mapped file access
- Direct GPU loading
- On-the-fly INT4 dequantization
- Fast cold start
"""

import json
import mmap
import struct
import tempfile
import base64
from pathlib import Path
from typing import Dict, List, Optional, Tuple, Any, Union
import numpy as np

import torch
import torch.nn as nn
from tqdm import tqdm

from .spec import (
    ZSEHeader,
    TensorInfo,
    TensorDType,
    QuantizationType,
    decode_header,
    zse_dtype_to_torch,
)


def dequantize_int4_zse(
    packed: torch.Tensor,
    scales: torch.Tensor,
    group_size: int = 128,
    dtype: torch.dtype = torch.float16,
) -> torch.Tensor:
    """
    Dequantize INT4 packed tensor back to FP16.
    
    Args:
        packed: [out_features, in_features//2] as UINT8
        scales: [out_features, num_groups] as FP16
        group_size: quantization group size
        dtype: output dtype
    """
    out_features = packed.shape[0]
    in_features = packed.shape[1] * 2
    
    # Unpack INT4 values
    low = (packed & 0x0F).to(torch.int8) - 8  # Back to [-7, 7]
    high = ((packed >> 4) & 0x0F).to(torch.int8) - 8
    
    # Interleave back to original order
    unpacked = torch.zeros(out_features, in_features, dtype=torch.int8, device=packed.device)
    unpacked[:, 0::2] = low
    unpacked[:, 1::2] = high
    
    # Reshape for group-wise dequantization
    num_groups = in_features // group_size
    unpacked_grouped = unpacked.view(out_features, num_groups, group_size)
    scales_expanded = scales.unsqueeze(-1)  # [out_features, num_groups, 1]
    
    # Dequantize
    dequantized = (unpacked_grouped.float() * scales_expanded).view(out_features, in_features)
    
    return dequantized.to(dtype)


def dequantize_int8_zse(
    quantized: torch.Tensor,
    scales: torch.Tensor,
    dtype: torch.dtype = torch.float16,
) -> torch.Tensor:
    """Dequantize INT8 tensor back to FP16."""
    # Per-channel scale
    if scales.ndim == 1:
        scales = scales.unsqueeze(1)
    return (quantized.float() * scales).to(dtype)


class ZSEReaderV2:
    """
    Reader for .zse format with proper INT4 support.
    """
    
    def __init__(
        self,
        path: Union[str, Path],
        device: str = "cuda",
    ):
        self.path = Path(path)
        self.device = device
        
        if not self.path.exists():
            raise FileNotFoundError(f"File not found: {self.path}")
        
        # Open and mmap file
        self._file = open(self.path, 'rb')
        self._mmap = mmap.mmap(self._file.fileno(), 0, access=mmap.ACCESS_READ)
        self._data = self._mmap
        
        # Parse header
        self.header, self._header_end = decode_header(self._data)
        
        # Build tensor lookup
        self._tensor_map = {t.name: t for t in self.header.tensors}
    
    def close(self):
        if self._mmap:
            self._mmap.close()
        self._file.close()
    
    def __enter__(self):
        return self
    
    def __exit__(self, *args):
        self.close()
    
    @property
    def quantization(self) -> str:
        return self.header.quantization
    
    @property
    def num_layers(self) -> int:
        return self.header.num_hidden_layers
    
    def get_info(self) -> Dict[str, Any]:
        """Get model info."""
        return {
            "path": str(self.path),
            "size_gb": self.path.stat().st_size / 1e9,
            "architecture": self.header.architecture,
            "num_layers": self.header.num_hidden_layers,
            "hidden_size": self.header.hidden_size,
            "quantization": self.header.quantization,
            "num_tensors": len(self.header.tensors),
        }
    
    def load_tokenizer(self):
        """Load embedded tokenizer."""
        from transformers import AutoTokenizer
        
        offset = self.header.tokenizer_offset
        size_bytes = self._data[offset:offset + 4]
        tokenizer_size = struct.unpack('<I', size_bytes)[0]
        
        tokenizer_json = self._data[offset + 4:offset + 4 + tokenizer_size].decode('utf-8')
        tokenizer_data = json.loads(tokenizer_json)
        
        # Extract to temp dir
        temp_dir = tempfile.mkdtemp(prefix="zse_tok_")
        for filename, content_b64 in tokenizer_data.items():
            content = base64.b64decode(content_b64)
            with open(Path(temp_dir) / filename, 'wb') as f:
                f.write(content)
        
        return AutoTokenizer.from_pretrained(temp_dir)
    
    def _read_raw_tensor(self, info: TensorInfo) -> torch.Tensor:
        """Read raw tensor bytes."""
        raw = self._data[info.offset:info.offset + info.size]
        
        # Map dtype
        np_dtype_map = {
            TensorDType.FLOAT32: np.float32,
            TensorDType.FLOAT16: np.float16,
            TensorDType.INT8: np.int8,
            TensorDType.INT4: np.uint8,  # INT4 is packed as uint8
            TensorDType.UINT8: np.uint8,
        }
        np_dtype = np_dtype_map.get(info.dtype, np.float16)
        
        array = np.frombuffer(raw, dtype=np_dtype).reshape(info.shape)
        return torch.from_numpy(array.copy())
    
    def load_tensor_dequantized(
        self,
        name: str,
        device: Optional[str] = None,
    ) -> torch.Tensor:
        """Load a tensor, dequantizing if needed."""
        device = device or self.device
        info = self._tensor_map.get(name)
        
        if info is None:
            raise KeyError(f"Tensor not found: {name}")
        
        tensor = self._read_raw_tensor(info)
        
        # Handle quantized tensors
        if info.quant_type != QuantizationType.NONE:
            scales_name = f"{name}.scales"
            if scales_name in self._tensor_map:
                scales_info = self._tensor_map[scales_name]
                scales = self._read_raw_tensor(scales_info).to(device)
                tensor = tensor.to(device)
                
                if info.dtype == TensorDType.INT4:
                    tensor = dequantize_int4_zse(tensor, scales, info.group_size)
                elif info.dtype == TensorDType.INT8:
                    tensor = dequantize_int8_zse(tensor, scales)
            else:
                tensor = tensor.to(device).half()
        else:
            tensor = tensor.to(device)
            if tensor.dtype != torch.float16 and tensor.is_floating_point():
                tensor = tensor.half()
        
        return tensor
    
    def load_state_dict_dequantized(
        self,
        device: Optional[str] = None,
        progress: bool = True,
    ) -> Dict[str, torch.Tensor]:
        """Load full state dict with dequantization."""
        device = device or self.device
        state_dict = {}
        
        # Get weight tensor names (skip scales)
        weight_names = [t.name for t in self.header.tensors if not t.name.endswith(".scales")]
        
        if progress:
            weight_names = tqdm(weight_names, desc="Loading tensors")
        
        for name in weight_names:
            tensor = self.load_tensor_dequantized(name, device)
            state_dict[name] = tensor
        
        return state_dict


class QuantizedLinearZSE(nn.Module):
    """
    Quantized linear layer that stores INT4 weights.
    Dequantizes on-the-fly during forward pass.
    """
    
    def __init__(
        self,
        in_features: int,
        out_features: int,
        group_size: int = 128,
        bias: bool = False,
    ):
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.group_size = group_size
        
        # Packed INT4 weights: [out_features, in_features//2]
        self.register_buffer('weight_packed', torch.zeros(out_features, in_features // 2, dtype=torch.uint8))
        
        # Scales: [out_features, num_groups]
        num_groups = (in_features + group_size - 1) // group_size
        self.register_buffer('weight_scales', torch.zeros(out_features, num_groups, dtype=torch.float16))
        
        if bias:
            self.register_buffer('bias', torch.zeros(out_features, dtype=torch.float16))
        else:
            self.register_buffer('bias', None)
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Dequantize weights on-the-fly
        weight = dequantize_int4_zse(
            self.weight_packed,
            self.weight_scales,
            self.group_size,
            dtype=x.dtype,
        )
        
        # Trim if we padded during quantization
        if weight.shape[1] > self.in_features:
            weight = weight[:, :self.in_features]
        
        return nn.functional.linear(x, weight, self.bias)


def load_zse_model(
    zse_path: Union[str, Path],
    device: str = "cuda",
) -> Tuple[nn.Module, Any, Dict]:
    """
    Load a .zse file into a usable model.
    
    Returns:
        (model, tokenizer, info)
    """
    from transformers import AutoConfig, AutoModelForCausalLM
    
    print(f"Loading {zse_path}...")
    
    with ZSEReaderV2(zse_path, device) as reader:
        info = reader.get_info()
        print(f"  Architecture: {info['architecture']}")
        print(f"  Quantization: {info['quantization']}")
        print(f"  File size: {info['size_gb']:.2f} GB")
        
        # Load tokenizer
        print("  Loading tokenizer...")
        tokenizer = reader.load_tokenizer()
        
        # Load state dict (dequantized to FP16 for now)
        # In future: keep as INT4 with QuantizedLinearZSE
        print("  Loading weights...")
        state_dict = reader.load_state_dict_dequantized(device)
        
        # Create model from config metadata
        print("  Creating model...")
        
        # Load HF config from source
        config = AutoConfig.from_pretrained(
            reader.header.source_model,
            trust_remote_code=True,
        )
        
        # Create model on meta device
        with torch.device("meta"):
            model = AutoModelForCausalLM.from_config(config, trust_remote_code=True)
        
        # Load weights
        model.load_state_dict(state_dict, assign=True, strict=False)
        
        return model, tokenizer, info
