"""Tests for model exporters."""

import pytest
import tempfile
from pathlib import Path

from model_exporter.exporters import get_exporter, GPT2Exporter


def test_gpt2_exporter_init():
    """Test GPT2Exporter initialization."""
    exporter = GPT2Exporter('gpt2')
    assert exporter.model_name == 'gpt2'
    assert exporter.model is None
    assert exporter.tokenizer is None


def test_get_exporter():
    """Test getting exporter by model name."""
    exporter = get_exporter('gpt2')
    assert isinstance(exporter, GPT2Exporter)
    
    exporter = get_exporter('gpt2-medium')
    assert isinstance(exporter, GPT2Exporter)


def test_get_exporter_unsupported():
    """Test error handling for unsupported models."""
    with pytest.raises(ValueError):
        get_exporter('unsupported-model')


@pytest.mark.slow
def test_gpt2_load_model():
    """Test loading GPT-2 model (slow test)."""
    exporter = GPT2Exporter('distilgpt2')  # Use smallest model for testing
    exporter.load_model()
    
    assert exporter.model is not None
    assert exporter.tokenizer is not None
    assert exporter.config is not None


@pytest.mark.slow
def test_gpt2_export_binary():
    """Test exporting GPT-2 to binary format (slow test)."""
    with tempfile.TemporaryDirectory() as tmpdir:
        exporter = GPT2Exporter('distilgpt2')
        exporter.load_model()
        
        output_dir = Path(tmpdir) / 'weights'
        exporter.export(output_dir, 'binary')
        
        # Check files exist
        assert (output_dir / 'config.json').exists()
        assert (output_dir / 'model_weights.bin').exists()
        assert (output_dir / 'vocab' / 'vocab.json').exists()
        assert (output_dir / 'vocab' / 'merges.txt').exists()