"""
Tests for HuggingFace adapter (using mock models to avoid downloads).
"""

import pytest
import torch
import torch.nn as nn
import numpy as np

from model_osmosis.adapters.huggingface import HuggingFaceExpert


class MockHFOutput:
    """Mock HuggingFace model output."""
    def __init__(self, last_hidden_state):
        self.last_hidden_state = last_hidden_state


class MockHFModel(nn.Module):
    """Mock HuggingFace model that accepts inputs_embeds."""
    def __init__(self, hidden_size):
        super().__init__()
        self.layer = nn.Linear(hidden_size, hidden_size)
        self.hidden_size = hidden_size

    def forward(self, inputs_embeds=None, **kwargs):
        out = self.layer(inputs_embeds)
        return MockHFOutput(last_hidden_state=out)


class TestHuggingFaceExpert:
    def test_forward_2d_input(self):
        """2D input (batch, hidden) should work."""
        mock = MockHFModel(hidden_size=64)
        expert = HuggingFaceExpert(mock, hidden_size=64, name="test")

        x = torch.rand(5, 64)
        with torch.no_grad():
            out = expert(x)
        assert out.shape == (5, 64)

    def test_forward_3d_input(self):
        """3D input (batch, seq, hidden) should work."""
        mock = MockHFModel(hidden_size=64)
        expert = HuggingFaceExpert(mock, hidden_size=64)

        x = torch.rand(5, 10, 64)
        with torch.no_grad():
            out = expert(x)
        assert out.shape == (5, 64)

    def test_weights_frozen(self):
        """Expert weights should be frozen (no grad)."""
        mock = MockHFModel(hidden_size=64)
        expert = HuggingFaceExpert(mock, hidden_size=64)

        for p in expert.parameters():
            assert not p.requires_grad

    def test_expert_name(self):
        mock = MockHFModel(hidden_size=64)
        expert = HuggingFaceExpert(mock, hidden_size=64, name="code_expert")
        assert expert.name == "code_expert"
        assert expert.hidden_size == 64


class TestHuggingFaceExpertStacking:
    def test_stack_hf_experts(self):
        """HuggingFaceExpert models can be stacked."""
        from model_osmosis.stacking import ModelStacker

        experts = []
        for i, name in enumerate(["code", "math", "chat"]):
            mock = MockHFModel(hidden_size=64)
            expert = HuggingFaceExpert(mock, hidden_size=64, name=name)
            experts.append(expert)

        sample = np.random.rand(30, 64).astype(np.float32)

        stacker = ModelStacker(verbose=False)
        result = stacker.stack(
            experts,
            sample_data=sample,
            expert_names=["code", "math", "chat"],
            top_k=2,
        )

        x = torch.rand(5, 64)
        with torch.no_grad():
            out = result(x)
        assert out.shape == (5, 64)

    def test_stack_mixed_hf_and_pytorch(self):
        """Can mix HuggingFace experts with regular PyTorch models."""
        from model_osmosis.stacking import ModelStacker

        # Regular PyTorch model
        pytorch_model = nn.Sequential(
            nn.Linear(64, 32), nn.ReLU(), nn.Linear(32, 64)
        )

        # HuggingFace expert
        mock = MockHFModel(hidden_size=64)
        hf_expert = HuggingFaceExpert(mock, hidden_size=64, name="hf")

        sample = np.random.rand(30, 64).astype(np.float32)

        stacker = ModelStacker(verbose=False)
        result = stacker.stack(
            [pytorch_model, hf_expert],
            sample_data=sample,
            expert_names=["pytorch", "huggingface"],
        )

        x = torch.rand(5, 64)
        with torch.no_grad():
            out = result(x)
        assert out.shape == (5, 64)
