"""ModalGpuEzLLM 유닛 테스트."""

from __future__ import annotations

from unittest.mock import patch

import pytest

from langchain_modal_gpu_ez import ModalGpuEzLLM
from langchain_modal_gpu_ez._utils import (
    apply_stop_words,
    parse_text_generation_output,
    strip_prompt_prefix,
)


class TestUtils:
    """출력 파싱 유틸리티 테스트."""

    def test_strip_prompt_prefix(self) -> None:
        assert strip_prompt_prefix("Hello", "Hello world") == " world"

    def test_strip_prompt_prefix_no_match(self) -> None:
        assert strip_prompt_prefix("Hi", "Hello world") == "Hello world"

    def test_apply_stop_words(self) -> None:
        assert apply_stop_words("Hello\n\nHuman: hi", ["\n\nHuman:"]) == "Hello"

    def test_apply_stop_words_none(self) -> None:
        assert apply_stop_words("Hello world", None) == "Hello world"

    def test_apply_stop_words_multiple(self) -> None:
        text = "Hello STOP1 world STOP2 end"
        assert apply_stop_words(text, ["STOP1", "STOP2"]) == "Hello "

    def test_parse_text_generation_output(self) -> None:
        raw = [{"generated_text": "Hello world is great"}]
        result = parse_text_generation_output(raw, "Hello", None)
        assert result == " world is great"

    def test_parse_translation_output(self) -> None:
        raw = ["안녕하세요"]
        result = parse_text_generation_output(raw, "Hello", None)
        assert result == "안녕하세요"

    def test_parse_with_stop(self) -> None:
        raw = [{"generated_text": "Hello world\nHuman: stop"}]
        result = parse_text_generation_output(raw, "Hello", ["\nHuman:"])
        assert result == " world"


class TestModalGpuEzLLM:
    """ModalGpuEzLLM 클래스 테스트."""

    def test_llm_type(self) -> None:
        llm = ModalGpuEzLLM(model_id="distilgpt2")
        assert llm._llm_type == "modal-gpu-ez"

    def test_default_params(self) -> None:
        llm = ModalGpuEzLLM(model_id="distilgpt2")
        assert llm.gpu == "T4"
        assert llm.max_new_tokens == 256
        assert llm.model_kwargs == {}

    def test_custom_params(self) -> None:
        llm = ModalGpuEzLLM(
            model_id="meta-llama/Llama-2-7b-hf",
            gpu="H100",
            max_new_tokens=512,
            model_kwargs={"temperature": 0.7},
        )
        assert llm.model_id == "meta-llama/Llama-2-7b-hf"
        assert llm.gpu == "H100"
        assert llm.max_new_tokens == 512
        assert llm.model_kwargs == {"temperature": 0.7}

    def test_extra_fields_forbidden(self) -> None:
        with pytest.raises(Exception):
            ModalGpuEzLLM(model_id="distilgpt2", unknown_field="value")  # type: ignore[call-arg]

    @patch("modal_gpu_ez.use")
    def test_generate(self, mock_use: object) -> None:
        from unittest.mock import MagicMock

        mock_use_fn = MagicMock(
            return_value=[{"generated_text": "Hello world is great"}]
        )
        with patch("modal_gpu_ez.use", mock_use_fn):
            llm = ModalGpuEzLLM(model_id="distilgpt2")
            result = llm._generate(["Hello"])

        assert len(result.generations) == 1
        assert result.generations[0][0].text == " world is great"
        mock_use_fn.assert_called_once_with(
            "distilgpt2",
            "T4",
            input="Hello",
            max_new_tokens=256,
        )

    @patch("modal_gpu_ez.use")
    def test_generate_with_stop(self, mock_use: object) -> None:
        from unittest.mock import MagicMock

        mock_use_fn = MagicMock(
            return_value=[{"generated_text": "Hello world\nHuman: stop"}]
        )
        with patch("modal_gpu_ez.use", mock_use_fn):
            llm = ModalGpuEzLLM(model_id="distilgpt2")
            result = llm._generate(["Hello"], stop=["\nHuman:"])

        assert result.generations[0][0].text == " world"

    @patch("modal_gpu_ez.use")
    def test_generate_multiple_prompts(self, mock_use: object) -> None:
        from unittest.mock import MagicMock

        mock_use_fn = MagicMock(
            side_effect=[
                [{"generated_text": "AAA BBB"}],
                [{"generated_text": "CCC DDD"}],
            ]
        )
        with patch("modal_gpu_ez.use", mock_use_fn):
            llm = ModalGpuEzLLM(model_id="distilgpt2")
            result = llm._generate(["AAA", "CCC"])

        assert len(result.generations) == 2
        assert result.generations[0][0].text == " BBB"
        assert result.generations[1][0].text == " DDD"

    @patch("modal_gpu_ez.use")
    def test_invoke(self, mock_use: object) -> None:
        from unittest.mock import MagicMock

        mock_use_fn = MagicMock(
            return_value=[{"generated_text": "Hello world"}]
        )
        with patch("modal_gpu_ez.use", mock_use_fn):
            llm = ModalGpuEzLLM(model_id="distilgpt2")
            result = llm.invoke("Hello")

        assert isinstance(result, str)
        assert result == " world"
