"""GPT-2 model exporter."""

from typing import Dict, Optional
import torch
from transformers import GPT2LMHeadModel, GPT2Tokenizer

from .base import BaseExporter


class GPT2Exporter(BaseExporter):
    """Exporter for GPT-2 family models."""
    
    SUPPORTED_MODELS = [
        'gpt2',
        'gpt2-medium',
        'gpt2-large',
        'gpt2-xl',
        'distilgpt2'
    ]
    
    def __init__(self, model_name: str = 'gpt2', cache_dir: Optional[str] = None):
        """
        Initialize GPT-2 exporter.
        
        Args:
            model_name: GPT-2 model variant
            cache_dir: Optional cache directory
        """
        if model_name not in self.SUPPORTED_MODELS:
            print(f"Warning: {model_name} not in tested models: {self.SUPPORTED_MODELS}")
        
        super().__init__(model_name, cache_dir)
    
    def load_model(self) -> None:
        """Load GPT-2 model and tokenizer from HuggingFace."""
        print(f"Loading {self.model_name} from HuggingFace...")
        
        self.model = GPT2LMHeadModel.from_pretrained(
            self.model_name,
            cache_dir=self.cache_dir
        )
        self.tokenizer = GPT2Tokenizer.from_pretrained(
            self.model_name,
            cache_dir=self.cache_dir
        )
        
        self.config = self.model.config.to_dict()
        
        total_params = sum(p.numel() for p in self.model.parameters())
        print(f"✓ Loaded {self.model_name} ({total_params:,} parameters)")
    
    def get_state_dict(self) -> Dict[str, torch.Tensor]:
        """
        Get GPT-2 state dictionary.
        
        Returns:
            State dictionary with all model weights
        """
        if self.model is None:
            raise ValueError("Model not loaded. Call load_model() first.")
        
        return self.model.state_dict()
    
    def get_merges(self) -> list:
        """
        Get BPE merges from tokenizer.
        
        Returns:
            List of merge tuples
        """
        if self.tokenizer is None:
            raise ValueError("Tokenizer not loaded. Call load_model() first.")
        
        return list(self.tokenizer.bpe_ranks.keys())