import argparse
from unittest.mock import MagicMock, patch

import pytest

from debgpt.llm import AbstractLLM, HuggingFaceLLM, create_llm


def test_abstract_llm():
    llm = AbstractLLM()

    with pytest.raises(NotImplementedError):
        llm.generate([{"role": "user", "content": "hi"}])

    with pytest.raises(NotImplementedError):
        llm.generate_stream([{"role": "user", "content": "hi"}])

    with pytest.raises(NotImplementedError):
        llm([{"role": "user", "content": "hi"}])

    with pytest.raises(NotImplementedError):
        llm.chat()


@patch("debgpt.llm.AutoTokenizer")
@patch("debgpt.llm.transformers.pipeline")
@patch("debgpt.llm.AutoModelForCausalLM")
def test_create_llm(mock_auto_model, mock_pipeline, mock_tokenizer):
    args = argparse.Namespace(
        model_id="test-model",
        device="cpu",
        precision="fp32",
        max_new_tokens=100,
        temperature=0.8,
        top_p=0.9,
    )

    llm = create_llm(args)
    assert isinstance(llm, HuggingFaceLLM)
    assert llm.model_id == "test-model"
    assert llm.kwargs["max_new_tokens"] == 100
    assert llm.kwargs["temperature"] == 0.8
    assert llm.kwargs["top_p"] == 0.9


@patch("debgpt.llm.AutoTokenizer")
@patch("debgpt.llm.transformers.pipeline")
@patch("debgpt.llm.AutoModelForCausalLM")
def test_create_llm_no_temp(mock_auto_model, mock_pipeline, mock_tokenizer):
    args = argparse.Namespace(
        model_id="test-model-2",
        device="cpu",
        precision="fp16",
        max_new_tokens=50,
        temperature=None,
        top_p=None,
    )

    llm = create_llm(args)
    assert isinstance(llm, HuggingFaceLLM)
    assert llm.kwargs["max_new_tokens"] == 50
    # check default fallback from HuggingFaceLLM init
    assert llm.kwargs["temperature"] == 0.7
    assert llm.kwargs["top_p"] == 0.95
