"""Plain binary format exporter for C++ inference."""

import struct
import json
from pathlib import Path
from typing import Dict, Any, Optional
import torch
import numpy as np
from tqdm import tqdm


class BinaryFormatExporter:
    """Export model weights to plain binary format."""
    
    DTYPE_MAP = {
        np.float32: 0,
        np.float16: 1,
        np.int32: 2,
        np.int64: 3
    }
    
    @staticmethod
    def export(
        state_dict: Dict[str, torch.Tensor],
        config: Dict[str, Any],
        tokenizer: Any,
        output_dir: Path,
        save_vocab: bool = True,
        **kwargs
    ) -> None:
        """
        Export weights to binary format.
        
        Args:
            state_dict: Model state dictionary
            config: Model configuration
            tokenizer: Model tokenizer
            output_dir: Output directory
            save_vocab: Whether to save vocabulary files
        """
        output_dir = Path(output_dir)
        output_dir.mkdir(parents=True, exist_ok=True)
        
        # Save config
        config_file = output_dir / "config.json"
        with open(config_file, 'w') as f:
            json.dump(config, f, indent=2)
        print(f"✓ Saved config to {config_file}")
        
        # Save vocabulary if requested
        if save_vocab and tokenizer is not None:
            BinaryFormatExporter._save_vocabulary(tokenizer, output_dir)
        
        # Save weights
        weights_file = output_dir / "model_weights.bin"
        BinaryFormatExporter._save_weights(state_dict, weights_file)
        
        file_size_mb = weights_file.stat().st_size / 1024 / 1024
        print(f"✓ Exported weights to {weights_file}")
        print(f"  File size: {file_size_mb:.2f} MB")
    
    @staticmethod
    def _save_vocabulary(tokenizer: Any, output_dir: Path) -> None:
        """Save tokenizer vocabulary and merges."""
        vocab_dir = output_dir / "vocab"
        vocab_dir.mkdir(exist_ok=True)
        
        # Save vocab
        vocab = tokenizer.get_vocab()
        vocab_file = vocab_dir / "vocab.json"
        with open(vocab_file, 'w') as f:
            json.dump(vocab, f)
        
        # Save merges (for BPE tokenizers)
        if hasattr(tokenizer, 'bpe_ranks'):
            merges_file = vocab_dir / "merges.txt"
            with open(merges_file, 'w') as f:
                for merge in tokenizer.bpe_ranks.keys():
                    f.write(f"{merge[0]} {merge[1]}\n")
            print(f"✓ Saved vocabulary to {vocab_file}")
            print(f"✓ Saved merges to {merges_file}")
        else:
            print(f"✓ Saved vocabulary to {vocab_file}")
    
    @staticmethod
    def _save_weights(state_dict: Dict[str, torch.Tensor], output_file: Path) -> None:
        """
        Save weights in binary format.
        
        Binary format structure:
        - 4 bytes: number of tensors (uint32)
        For each tensor:
          - 4 bytes: name length (uint32)
          - N bytes: name (utf-8)
          - 4 bytes: number of dimensions (uint32)
          - 4*ndim bytes: shape (uint32 for each dimension)
          - 4 bytes: dtype code (uint32)
          - 8 bytes: data size in bytes (uint64)
          - N bytes: tensor data
        """
        print(f"\nExporting {len(state_dict)} tensors to binary format...")
        
        total_params = sum(t.numel() for t in state_dict.values())
        print(f"Total parameters: {total_params:,}")
        
        with open(output_file, 'wb') as f:
            # Write number of tensors
            num_tensors = len(state_dict)
            f.write(struct.pack('I', num_tensors))
            
            # Write each tensor
            for name, tensor in tqdm(state_dict.items(), desc="Writing tensors"):
                weight = tensor.cpu().numpy()
                
                # Write name
                name_bytes = name.encode('utf-8')
                f.write(struct.pack('I', len(name_bytes)))
                f.write(name_bytes)
                
                # Write shape
                f.write(struct.pack('I', len(weight.shape)))
                for dim in weight.shape:
                    f.write(struct.pack('I', dim))
                
                # Write dtype
                dtype_code = BinaryFormatExporter.DTYPE_MAP.get(weight.dtype.type, 0)
                f.write(struct.pack('I', dtype_code))
                
                # Write data
                weight_bytes = weight.tobytes()
                f.write(struct.pack('Q', len(weight_bytes)))
                f.write(weight_bytes)
    
    @staticmethod
    def load_info(weights_file: Path, max_tensors: int = 5) -> None:
        """
        Load and print information about a binary weights file.
        
        Args:
            weights_file: Path to binary weights file
            max_tensors: Maximum number of tensors to display info for
        """
        dtype_names = {0: 'float32', 1: 'float16', 2: 'int32', 3: 'int64'}
        
        with open(weights_file, 'rb') as f:
            num_tensors = struct.unpack('I', f.read(4))[0]
            print(f"Number of tensors: {num_tensors}\n")
            
            for i in range(min(max_tensors, num_tensors)):
                # Read name
                name_len = struct.unpack('I', f.read(4))[0]
                name = f.read(name_len).decode('utf-8')
                
                # Read shape
                ndims = struct.unpack('I', f.read(4))[0]
                shape = struct.unpack(f'{ndims}I', f.read(4 * ndims))
                
                # Read dtype
                dtype = struct.unpack('I', f.read(4))[0]
                
                # Read data size and skip data
                data_size = struct.unpack('Q', f.read(8))[0]
                f.seek(data_size, 1)
                
                print(f"{i+1}. {name}")
                print(f"   Shape: {shape}")
                print(f"   Dtype: {dtype_names.get(dtype, 'unknown')}")
                print(f"   Size: {data_size} bytes\n")
            
            if num_tensors > max_tensors:
                print(f"... and {num_tensors - max_tensors} more tensors")