import pytest
import torch
from unittest.mock import Mock
from swarms.models.huggingface import HuggingFaceLLM  


@pytest.fixture
def mock_torch():
    return Mock()


@pytest.fixture
def mock_autotokenizer():
    return Mock()


@pytest.fixture
def mock_automodelforcausallm():
    return Mock()


@pytest.fixture
def mock_bitsandbytesconfig():
    return Mock()


@pytest.fixture
def hugging_face_llm(mock_torch, mock_autotokenizer, mock_automodelforcausallm, mock_bitsandbytesconfig):
    HuggingFaceLLM.torch = mock_torch
    HuggingFaceLLM.AutoTokenizer = mock_autotokenizer
    HuggingFaceLLM.AutoModelForCausalLM = mock_automodelforcausallm
    HuggingFaceLLM.BitsAndBytesConfig = mock_bitsandbytesconfig

    return HuggingFaceLLM(model_id='test')


def test_init(hugging_face_llm, mock_autotokenizer, mock_automodelforcausallm):
    assert hugging_face_llm.model_id == 'test'
    mock_autotokenizer.from_pretrained.assert_called_once_with('test')
    mock_automodelforcausallm.from_pretrained.assert_called_once_with('test', quantization_config=None)


def test_init_with_quantize(hugging_face_llm, mock_autotokenizer, mock_automodelforcausallm, mock_bitsandbytesconfig):
    quantization_config = {
        'load_in_4bit': True,
        'bnb_4bit_use_double_quant': True,
        'bnb_4bit_quant_type': "nf4",
        'bnb_4bit_compute_dtype': torch.bfloat16
    }
    mock_bitsandbytesconfig.return_value = quantization_config

    HuggingFaceLLM(model_id='test', quantize=True)

    mock_bitsandbytesconfig.assert_called_once_with(**quantization_config)
    mock_autotokenizer.from_pretrained.assert_called_once_with('test')
    mock_automodelforcausallm.from_pretrained.assert_called_once_with('test', quantization_config=quantization_config)


def test_generate_text(hugging_face_llm):
    prompt_text = 'test prompt'
    expected_output = 'test output'
    hugging_face_llm.tokenizer.encode.return_value = torch.tensor([0])  # Mock tensor
    hugging_face_llm.model.generate.return_value = torch.tensor([0])  # Mock tensor
    hugging_face_llm.tokenizer.decode.return_value = expected_output

    output = hugging_face_llm.generate_text(prompt_text)

    assert output == expected_output