"""Safetensors format exporter."""

import json
from pathlib import Path
from typing import Dict, Any
import torch
from safetensors.torch import save_file, load_file
from safetensors import safe_open


class SafetensorsFormatExporter:
    """Export model weights to safetensors format."""
    
    @staticmethod
    def export(
        state_dict: Dict[str, torch.Tensor],
        config: Dict[str, Any],
        tokenizer: Any,
        output_dir: Path,
        **kwargs
    ) -> None:
        """
        Export weights to safetensors format.
        
        Args:
            state_dict: Model state dictionary
            config: Model configuration
            tokenizer: Model tokenizer
            output_dir: Output directory
        """
        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 weights in safetensors format
        weights_file = output_dir / "model.safetensors"
        
        # Convert state dict to CPU tensors and handle shared tensors
        cpu_state_dict = {}
        seen_data_ptrs = {}
        
        for name, tensor in state_dict.items():
            cpu_tensor = tensor.cpu()
            data_ptr = cpu_tensor.data_ptr()
            
            # Check if we've seen this memory location before
            if data_ptr in seen_data_ptrs:
                original_name = seen_data_ptrs[data_ptr]
                print(f"  Note: {name} shares memory with {original_name}, cloning...")
                cpu_state_dict[name] = cpu_tensor.clone()
            else:
                seen_data_ptrs[data_ptr] = name
                cpu_state_dict[name] = cpu_tensor
        
        # Save using safetensors
        save_file(cpu_state_dict, weights_file)
        
        file_size_mb = weights_file.stat().st_size / 1024 / 1024
        total_params = sum(t.numel() for t in state_dict.values())
        
        print(f"✓ Exported {len(cpu_state_dict)} tensors to {weights_file}")
        print(f"  Total parameters: {total_params:,}")
        print(f"  File size: {file_size_mb:.2f} MB")
    
    @staticmethod
    def load_info(weights_file: Path, max_tensors: int = 5) -> None:
        """
        Load and print information about a safetensors file.
        
        Args:
            weights_file: Path to safetensors file
            max_tensors: Maximum number of tensors to display info for
        """
        with safe_open(weights_file, framework="pt", device="cpu") as f:
            keys = f.keys()
            num_tensors = len(keys)
            
            print(f"Number of tensors: {num_tensors}\n")
            
            total_params = 0
            for i, key in enumerate(keys):
                if i < max_tensors:
                    tensor = f.get_tensor(key)
                    shape = tuple(tensor.shape)
                    dtype = str(tensor.dtype).replace('torch.', '')
                    numel = tensor.numel()
                    total_params += numel
                    
                    print(f"{i+1}. {key}")
                    print(f"   Shape: {shape}")
                    print(f"   Dtype: {dtype}")
                    print(f"   Elements: {numel:,}\n")
                else:
                    # Just count params for remaining tensors
                    tensor = f.get_tensor(key)
                    total_params += tensor.numel()
            
            if num_tensors > max_tensors:
                print(f"... and {num_tensors - max_tensors} more tensors")
            
            print(f"\nTotal parameters: {total_params:,}")