"""Model exporters registry."""

from .base import BaseExporter
from .gpt2 import GPT2Exporter


EXPORTER_REGISTRY = {
    'gpt2': GPT2Exporter,
    'gpt2-medium': GPT2Exporter,
    'gpt2-large': GPT2Exporter,
    'gpt2-xl': GPT2Exporter,
    'distilgpt2': GPT2Exporter,
    # 'llama': LlamaExporter,  # To be implemented
    # 'bert': BertExporter,    # To be implemented
}


def get_exporter(model_name: str, **kwargs) -> BaseExporter:
    """
    Get appropriate exporter for a model.
    
    Args:
        model_name: HuggingFace model name
        **kwargs: Additional arguments for exporter
    
    Returns:
        Model exporter instance
    
    Raises:
        ValueError: If model not supported
    """
    # Try exact match first
    if model_name in EXPORTER_REGISTRY:
        return EXPORTER_REGISTRY[model_name](model_name, **kwargs)
    
    # Try prefix match
    for key, exporter_class in EXPORTER_REGISTRY.items():
        if model_name.startswith(key):
            return exporter_class(model_name, **kwargs)
    
    available = ', '.join(set(EXPORTER_REGISTRY.values().__class__.__name__))
    raise ValueError(
        f"Model '{model_name}' not supported. "
        f"Available exporters: {available}"
    )


def list_supported_models() -> dict:
    """List all supported model families."""
    families = {}
    for model_name, exporter_class in EXPORTER_REGISTRY.items():
        family = exporter_class.__name__.replace('Exporter', '')
        if family not in families:
            families[family] = []
        if model_name not in families[family]:
            families[family].append(model_name)
    return families


__all__ = [
    'BaseExporter',
    'GPT2Exporter',
    'get_exporter',
    'list_supported_models',
]