from __future__ import annotations

import logging
from typing import Any

from leshy.buffs.base import Buff
from leshy.core.registry import buff_registry
from leshy.core.tokens import token_tracker

logger = logging.getLogger(__name__)


def _get_litellm() -> Any:
    try:
        import litellm
    except ImportError:
        raise ImportError(
            "The 'litellm' package is required for LLM buffs. "
            "Install it with: pip install 'leshy[llm]'"
        )
    return litellm


@buff_registry.register("rephrase")
class RephraseBuff(Buff):
    """Rephrase payload using an LLM while preserving intent."""

    name = "rephrase"
    description = "Rephrase using LLM (requires leshy[llm])"

    def __init__(self, **kwargs: Any) -> None:
        super().__init__(**kwargs)
        self._model = kwargs.get("model", "gpt-4o-mini")
        self._api_key: str | None = kwargs.get("api_key")
        self._style = kwargs.get("style", "casual")

    async def transform(self, content: str) -> str:
        litellm = _get_litellm()
        try:
            response = await litellm.acompletion(
                model=self._model,
                api_key=self._api_key,
                messages=[
                    {
                        "role": "system",
                        "content": (
                            f"Rephrase the following text in a {self._style} style. "
                            "Keep the core meaning and intent exactly the same. "
                            "Only output the rephrased text, nothing else."
                        ),
                    },
                    {"role": "user", "content": content},
                ],
                temperature=0.7,
                max_tokens=1000,
            )
            if response.usage:
                token_tracker.add(
                    response.usage.prompt_tokens,
                    response.usage.completion_tokens,
                )
            return response.choices[0].message.content or content
        except ImportError:
            raise
        except Exception as e:
            logger.warning("Rephrase buff failed: %s", e)
            return content


@buff_registry.register("translate")
class TranslateBuff(Buff):
    """Translate payload to another language using an LLM."""

    name = "translate"
    description = "Translate using LLM (requires leshy[llm])"

    def __init__(self, **kwargs: Any) -> None:
        super().__init__(**kwargs)
        self._model = kwargs.get("model", "gpt-4o-mini")
        self._api_key: str | None = kwargs.get("api_key")
        self._language = kwargs.get("language", "Spanish")

    async def transform(self, content: str) -> str:
        litellm = _get_litellm()
        try:
            response = await litellm.acompletion(
                model=self._model,
                api_key=self._api_key,
                messages=[
                    {
                        "role": "system",
                        "content": (
                            f"Translate the following text to {self._language}. "
                            "Keep the meaning and intent exactly the same. "
                            "Only output the translated text, nothing else."
                        ),
                    },
                    {"role": "user", "content": content},
                ],
                temperature=0.3,
                max_tokens=1000,
            )
            if response.usage:
                token_tracker.add(
                    response.usage.prompt_tokens,
                    response.usage.completion_tokens,
                )
            return response.choices[0].message.content or content
        except ImportError:
            raise
        except Exception as e:
            logger.warning("Translate buff failed: %s", e)
            return content
