"""
Tests for Hugging Face Transformers Integration.
"""

import sys
import os
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

import torch
import tempfile
import pytest
from transformers import AutoConfig, AutoModelForCausalLM
from src.hf_model import QTensorFormerConfig, QTensorFormerForCausalLM


def test_hf_config_and_model_forward():
    config = QTensorFormerConfig(
        vocab_size=100,
        d_model=32,
        n_layers=2,
        n_heads=2,
        tt_rank=4,
        max_seq_len=32,
        use_quantum=True,
    )
    model = QTensorFormerForCausalLM(config)

    input_ids = torch.randint(0, 100, (2, 8))
    out = model(input_ids)

    assert out.logits.shape == (2, 8, 100)
    print("✓ test_hf_config_and_model_forward passed")


def test_hf_generation():
    config = QTensorFormerConfig(
        vocab_size=100,
        d_model=32,
        n_layers=2,
        n_heads=2,
        tt_rank=4,
        max_seq_len=32,
    )
    model = QTensorFormerForCausalLM(config)

    input_ids = torch.tensor([[1, 2, 3]])
    generated = model.generate(input_ids, max_new_tokens=5, do_sample=False)

    assert generated.shape == (1, 8), f"Expected shape (1, 8), got {generated.shape}"
    print("✓ test_hf_generation passed")


def test_hf_save_and_load():
    config = QTensorFormerConfig(
        vocab_size=50,
        d_model=32,
        n_layers=2,
        n_heads=2,
        tt_rank=4,
    )
    model = QTensorFormerForCausalLM(config)

    with tempfile.TemporaryDirectory() as tmpdir:
        model.save_pretrained(tmpdir)

        # Load back with AutoModelForCausalLM
        loaded = AutoModelForCausalLM.from_pretrained(tmpdir)
        assert isinstance(loaded, QTensorFormerForCausalLM)

        model.eval()
        loaded.eval()
        inp = torch.randint(0, 50, (1, 4))
        with torch.no_grad():
            out1 = model(inp).logits
            out2 = loaded(inp).logits
        assert torch.allclose(out1, out2, atol=1e-4)

    print("✓ test_hf_save_and_load passed")


if __name__ == "__main__":
    test_hf_config_and_model_forward()
    test_hf_generation()
    test_hf_save_and_load()
    print("All Hugging Face integration tests passed!")