"""Base exporter interface for model weight extraction."""

from abc import ABC, abstractmethod
from pathlib import Path
from typing import Dict, Any, Optional
import torch
from transformers import PreTrainedModel, PreTrainedTokenizer


class BaseExporter(ABC):
    """Base class for model exporters."""
    
    def __init__(self, model_name: str, cache_dir: Optional[str] = None):
        """
        Initialize the exporter.
        
        Args:
            model_name: HuggingFace model name or local path
            cache_dir: Optional cache directory for downloaded models
        """
        self.model_name = model_name
        self.cache_dir = cache_dir
        self.model: Optional[PreTrainedModel] = None
        self.tokenizer: Optional[PreTrainedTokenizer] = None
        self.config: Optional[Dict[str, Any]] = None
    
    @abstractmethod
    def load_model(self) -> None:
        """Load the model and tokenizer from HuggingFace."""
        pass
    
    @abstractmethod
    def get_state_dict(self) -> Dict[str, torch.Tensor]:
        """
        Get the model's state dictionary.
        
        Returns:
            Dictionary mapping parameter names to tensors
        """
        pass
    
    def get_config(self) -> Dict[str, Any]:
        """
        Get model configuration.
        
        Returns:
            Configuration dictionary
        """
        if self.model is None:
            raise ValueError("Model not loaded. Call load_model() first.")
        return self.model.config.to_dict()
    
    def get_vocab(self) -> Dict[str, int]:
        """
        Get tokenizer vocabulary.
        
        Returns:
            Vocabulary dictionary
        """
        if self.tokenizer is None:
            raise ValueError("Tokenizer not loaded. Call load_model() first.")
        return self.tokenizer.get_vocab()
    
    def export(self, output_dir: Path, format_name: str, **kwargs) -> None:
        """
        Export model weights in specified format.
        
        Args:
            output_dir: Directory to save exported weights
            format_name: Export format ('binary', 'gguf', 'safetensors', 'onnx')
            **kwargs: Additional format-specific arguments
        """
        from model_exporter.formats import get_format_exporter
        
        if self.model is None:
            self.load_model()
        
        output_dir = Path(output_dir)
        output_dir.mkdir(parents=True, exist_ok=True)
        
        format_exporter = get_format_exporter(format_name)
        format_exporter.export(
            state_dict=self.get_state_dict(),
            config=self.get_config(),
            tokenizer=self.tokenizer,
            output_dir=output_dir,
            **kwargs
        )
    
    def print_model_info(self) -> None:
        """Print model architecture and parameter information."""
        if self.model is None:
            raise ValueError("Model not loaded. Call load_model() first.")
        
        config = self.get_config()
        total_params = sum(p.numel() for p in self.model.parameters())
        
        print("\n=== Model Information ===")
        print(f"Model: {self.model_name}")
        print(f"Total parameters: {total_params:,}")
        print(f"\nConfiguration:")
        for key, value in config.items():
            if not key.startswith('_'):
                print(f"  {key}: {value}")