"""Local HuggingFace LLM implementation."""

from __future__ import annotations

import asyncio

from roohai.base import LLMModel


class LocalLLM(LLMModel):
    META = {
        "display_name": "Local LLM (TinyLlama)",
        "description": "Run a HuggingFace language model locally. Requires GPU for reasonable speed.",
        "kind": "local",
        "pip_dependencies": ["transformers>=4.46.0,<5.0.0", "torch", "accelerate"],
        "hf_models": ["TinyLlama/TinyLlama-1.1B-Chat-v1.0"],
        "estimated_download_mb": 2200,
        "config_fields": [
            {
                "name": "model_id",
                "label": "Model ID",
                "type": "text",
                "required": False,
                "default": "TinyLlama/TinyLlama-1.1B-Chat-v1.0",
            },
        ],
    }

    def __init__(self, model_id: str = "TinyLlama/TinyLlama-1.1B-Chat-v1.0", **kwargs) -> None:
        self.model_id = model_id
        self._pipeline = None

    def load(self) -> None:
        from transformers import pipeline

        self._pipeline = pipeline(
            "text-generation",
            model=self.model_id,
            device_map="auto",
            dtype="auto",
        )

    def unload(self) -> None:
        if self._pipeline is not None:
            from roohai.base import _release_torch_memory
            del self._pipeline
            self._pipeline = None
            _release_torch_memory()

    @property
    def is_loaded(self) -> bool:
        return self._pipeline is not None

    def _generate(self, message: str, history: list[dict] | None, system_prompt: str | None) -> str:
        if not self.is_loaded:
            raise RuntimeError("Model not loaded. Call load() first.")

        messages = []
        if system_prompt:
            messages.append({"role": "system", "content": system_prompt})
        if history:
            messages.extend(history)
        messages.append({"role": "user", "content": message})

        result = self._pipeline(messages, max_new_tokens=256, do_sample=True, temperature=0.7)
        return result[0]["generated_text"][-1]["content"]

    async def chat(
        self,
        message: str,
        history: list[dict] | None = None,
        system_prompt: str | None = None,
        session_id: str | None = None,
    ) -> str:
        return await asyncio.to_thread(self._generate, message, history, system_prompt)
