"""Modal GPU EZ 텍스트 생성 모델을 위한 LangChain LLM 래퍼."""

from __future__ import annotations

import asyncio
from typing import Any

from langchain_core.callbacks import CallbackManagerForLLMRun
from langchain_core.language_models.llms import BaseLLM
from langchain_core.outputs import Generation, LLMResult
from pydantic import ConfigDict, Field

from langchain_modal_gpu_ez._utils import parse_text_generation_output


class ModalGpuEzLLM(BaseLLM):
    """modal_gpu_ez를 사용하는 LangChain LLM 구현.

    Example:
        .. code-block:: python

            from langchain_modal_gpu_ez import ModalGpuEzLLM

            llm = ModalGpuEzLLM(model_id="distilgpt2", gpu="T4")
            result = llm.invoke("Once upon a time")
    """

    model_id: str
    gpu: str = "T4"
    max_new_tokens: int = 256
    model_kwargs: dict[str, Any] = Field(default_factory=dict)

    model_config = ConfigDict(
        extra="forbid",
        protected_namespaces=(),
    )

    @property
    def _llm_type(self) -> str:
        return "modal-gpu-ez"

    def _generate(
        self,
        prompts: list[str],
        stop: list[str] | None = None,
        run_manager: CallbackManagerForLLMRun | None = None,
        **kwargs: Any,
    ) -> LLMResult:
        """프롬프트 리스트에 대해 텍스트를 생성한다."""
        import modal_gpu_ez as mg

        merged_kwargs = {**self.model_kwargs, **kwargs}
        generations: list[list[Generation]] = []

        for prompt in prompts:
            raw = mg.use(
                self.model_id,
                self.gpu,
                input=prompt,
                max_new_tokens=self.max_new_tokens,
                **merged_kwargs,
            )
            text = parse_text_generation_output(raw, prompt, stop)
            generations.append([Generation(text=text)])

        return LLMResult(generations=generations)

    async def _agenerate(
        self,
        prompts: list[str],
        stop: list[str] | None = None,
        run_manager: Any | None = None,
        **kwargs: Any,
    ) -> LLMResult:
        """비동기로 프롬프트 리스트에 대해 텍스트를 병렬 생성한다."""
        import modal_gpu_ez as mg

        merged_kwargs = {**self.model_kwargs, **kwargs}

        async def _generate_one(prompt: str) -> list[Generation]:
            raw = await mg.async_use(
                self.model_id,
                self.gpu,
                input=prompt,
                max_new_tokens=self.max_new_tokens,
                **merged_kwargs,
            )
            text = parse_text_generation_output(raw, prompt, stop)
            return [Generation(text=text)]

        generations = await asyncio.gather(
            *[_generate_one(p) for p in prompts]
        )
        return LLMResult(generations=list(generations))
