"""Provider implementation for LiteLLM."""

import asyncio
import dataclasses
import logging
import threading
from collections.abc import Iterator, Sequence
from typing import Any

import langcore as lx
import litellm
from langcore.core.base_model import BaseLanguageModel
from langcore.core.types import ScoredOutput
from litellm.exceptions import (
    APIConnectionError as LiteLLMConnectionError,
    APIError as LiteLLMAPIError,
    AuthenticationError as LiteLLMAuthenticationError,
    BadRequestError as LiteLLMBadRequestError,
    NotFoundError as LiteLLMNotFoundError,
    PermissionDeniedError as LiteLLMPermissionDeniedError,
    RateLimitError as LiteLLMRateLimitError,
    Timeout as LiteLLMTimeout,
)

logger = logging.getLogger(__name__)

# Keys consumed internally by the provider and must not be
# forwarded to ``litellm.completion()`` / ``litellm.acompletion()``.
_INTERNAL_KEYS: frozenset[str] = frozenset({"max_workers", "pass_num"})


@dataclasses.dataclass
class UsageStats:
    """Token usage statistics returned by the LLM.

    Instances are available via :pyattr:`LiteLLMLanguageModel.last_usage`
    (most recent call) and :pyattr:`LiteLLMLanguageModel.total_usage`
    (cumulative across all calls since the last :pymeth:`reset_usage`).
    """

    prompt_tokens: int = 0
    completion_tokens: int = 0
    total_tokens: int = 0

    def __iadd__(self, other: "UsageStats") -> "UsageStats":
        """Accumulate usage from *other* into this instance."""
        self.prompt_tokens += other.prompt_tokens
        self.completion_tokens += other.completion_tokens
        self.total_tokens += other.total_tokens
        return self


@lx.providers.registry.register(r"^litellm", priority=10)
class LiteLLMLanguageModel(BaseLanguageModel):
    """LangCore provider backed by LiteLLM's unified API.

    Routes to any model supported by LiteLLM (OpenAI, Anthropic,
    Google, Mistral, Ollama, vLLM, Azure, Bedrock, …).  Model IDs
    must carry a ``litellm/`` or ``litellm-`` prefix so that the
    LangCore provider registry can dispatch to this class.

    Registration pattern: ``r"^litellm"`` (priority 10).

    Examples::

        litellm/gpt-4o
        litellm/anthropic/claude-3-opus
        litellm/ollama/llama3
        litellm-azure/gpt-4o
    """

    def __init__(self, model_id: str, **kwargs: Any) -> None:
        """Initialize the LiteLLM provider.

        Args:
            model_id: The model identifier (e.g., 'gpt-4',
                'claude-3-opus', 'llama-2-7b-chat').
            **kwargs: Any parameters supported by
                litellm.completion(), including:
                - api_key: API key for authentication. If not
                    provided, LiteLLM uses provider-specific
                    environment variables (OPENAI_API_KEY,
                    ANTHROPIC_API_KEY, GOOGLE_API_KEY, etc.)
                - api_base: Custom API base URL
                - temperature: Sampling temperature (0.0-1.0)
                - max_tokens: Maximum tokens to generate
                - top_p: Top-p sampling parameter
                - frequency_penalty: Frequency penalty
                - presence_penalty: Presence penalty
                - timeout: Request timeout in seconds
                - max_workers (int): Maximum concurrent async
                    requests (default: 10). Consumed internally,
                    not forwarded to LiteLLM.
        """
        super().__init__()

        # Remove litellm prefix for actual model calls
        if model_id.startswith("litellm/"):
            self.model_id = model_id[8:]  # Remove 'litellm/' prefix
        elif model_id.startswith("litellm-"):
            self.model_id = model_id[8:]  # Remove 'litellm-' prefix
        else:
            self.model_id = model_id

        self.original_model_id = model_id

        # Pop internal keys before storing provider kwargs so they
        # are never forwarded to ``litellm.completion()``.
        self._max_workers: int = kwargs.pop("max_workers", 10)
        self.provider_kwargs = kwargs

        # Lazily initialised in ``_get_semaphore`` to avoid binding
        # to an event loop that may not exist yet at construction.
        self._semaphore: asyncio.Semaphore | None = None

        # Token-usage tracking — thread-safe via a lock.
        self._usage_lock = threading.Lock()
        self._last_usage = UsageStats()
        self._total_usage = UsageStats()

        logger.info("Initialized LiteLLM provider for model: %s", self.model_id)

    # ------------------------------------------------------------------
    # Token-usage public API
    # ------------------------------------------------------------------

    @property
    def last_usage(self) -> UsageStats:
        """Token usage from the most recent ``infer`` / ``async_infer`` call."""
        with self._usage_lock:
            return dataclasses.replace(self._last_usage)

    @property
    def total_usage(self) -> UsageStats:
        """Cumulative token usage since construction or last ``reset_usage``."""
        with self._usage_lock:
            return dataclasses.replace(self._total_usage)

    def reset_usage(self) -> None:
        """Reset both ``last_usage`` and ``total_usage`` to zero."""
        with self._usage_lock:
            self._last_usage = UsageStats()
            self._total_usage = UsageStats()

    def _record_usage(self, response: Any) -> UsageStats | None:
        """Extract token usage from *response* and update accumulators."""
        usage = getattr(response, "usage", None)
        if usage is None:
            return None
        stats = UsageStats(
            prompt_tokens=getattr(usage, "prompt_tokens", 0) or 0,
            completion_tokens=getattr(usage, "completion_tokens", 0) or 0,
            total_tokens=getattr(usage, "total_tokens", 0) or 0,
        )
        logger.debug(
            "Token usage for %s: prompt=%d, completion=%d, total=%d",
            self.model_id,
            stats.prompt_tokens,
            stats.completion_tokens,
            stats.total_tokens,
        )
        with self._usage_lock:
            self._last_usage = stats
            self._total_usage += stats
        return stats

    @property
    def _litellm_kwargs(self) -> dict[str, Any]:
        """Provider kwargs with internal-only keys stripped.

        Ensures keys like ``max_workers`` that are consumed by
        the provider itself are never forwarded to
        ``litellm.completion()`` / ``litellm.acompletion()``.

        Note: ``__init__`` already pops keys listed in
        ``_INTERNAL_KEYS`` from kwargs, so this filter is a
        belt-and-suspenders defence against future direct mutation
        of ``provider_kwargs``.  When adding new internal keys,
        update **both** the ``pop`` in ``__init__`` and the
        ``_INTERNAL_KEYS`` set.
        """
        return {
            k: v for k, v in self.provider_kwargs.items() if k not in _INTERNAL_KEYS
        }

    def _get_semaphore(self) -> asyncio.Semaphore:
        """Return the concurrency-limiting semaphore.

        Lazily initialised so it is bound to the running event
        loop, not whichever loop (if any) existed at
        construction time.
        """
        if self._semaphore is None:
            self._semaphore = asyncio.Semaphore(self._max_workers)
        return self._semaphore

    def _parse_response(self, response: Any) -> list[ScoredOutput]:
        """Convert a LiteLLM response to a list of ScoredOutput.

        Centralises the response-extraction logic so that both
        ``infer()`` and ``async_infer()`` share the same code path.
        Also records token usage via ``_record_usage``.

        Args:
            response: The response object returned by
                ``litellm.completion()`` or
                ``litellm.acompletion()``.

        Returns:
            A single-element list of ScoredOutput with score 1.0
            on success, or score 0.0 for empty/missing content.
        """
        stats = self._record_usage(response)
        usage_dict = None
        if stats is not None:
            usage_dict = {
                "prompt_tokens": stats.prompt_tokens,
                "completion_tokens": stats.completion_tokens,
                "total_tokens": stats.total_tokens,
            }

        if response.choices:
            content = response.choices[0].message.content
            if content:
                return [ScoredOutput(score=1.0, output=content, usage=usage_dict)]
            logger.warning(
                "Empty response from LiteLLM for model %s",
                self.model_id,
            )
            return [ScoredOutput(score=0.0, output="", usage=usage_dict)]
        logger.error(
            "No choices in response from LiteLLM for model %s",
            self.model_id,
        )
        return [ScoredOutput(score=0.0, output="", usage=usage_dict)]

    def infer(
        self, batch_prompts: Sequence[str], **kwargs: Any
    ) -> Iterator[Sequence[ScoredOutput]]:
        """Run inference on a batch of prompts.

        Args:
            batch_prompts: List of prompts to process.
            **kwargs: Additional inference parameters that override
                instance defaults.  ``pass_num`` (int) is consumed
                internally: when >= 1 the call includes
                ``cache={"no-cache": True}`` so that repeat
                extraction passes are never served from the LiteLLM
                response cache.

        Yields:
            Lists of ScoredOutput objects, one per prompt.
        """
        pass_num: int = kwargs.pop("pass_num", 0)

        # Build per-call kwargs: instance defaults + cache bypass.
        call_kwargs = dict(self._litellm_kwargs)
        if pass_num >= 1:
            call_kwargs["cache"] = {"no-cache": True}

        logger.info(
            "Running sync inference for %d prompt(s) on %s",
            len(batch_prompts),
            self.model_id,
        )

        for prompt in batch_prompts:
            try:
                logger.debug(
                    "Calling LiteLLM completion for model %s",
                    self.model_id,
                )

                messages = [{"role": "user", "content": prompt}]

                response = litellm.completion(
                    model=self.model_id,
                    messages=messages,
                    **call_kwargs,
                )

                yield self._parse_response(response)

            except (
                LiteLLMAPIError,
                LiteLLMAuthenticationError,
                LiteLLMBadRequestError,
                LiteLLMConnectionError,
                LiteLLMNotFoundError,
                LiteLLMPermissionDeniedError,
                LiteLLMRateLimitError,
                LiteLLMTimeout,
            ) as e:
                logger.warning(
                    "LiteLLM API error for model %s: %s: %s",
                    self.model_id,
                    type(e).__name__,
                    e,
                )
                yield [ScoredOutput(score=0.0, output="LLM inference failed")]
            except Exception:
                logger.exception(
                    "Unexpected error during LiteLLM inference for model %s",
                    self.model_id,
                )
                yield [ScoredOutput(score=0.0, output="LLM inference failed")]

    async def async_infer(
        self, batch_prompts: Sequence[str], **kwargs: Any
    ) -> list[Sequence[ScoredOutput]]:
        """Native async inference using ``litellm.acompletion``.

        Uses a shared ``asyncio.Semaphore`` for concurrency control
        instead of ``ThreadPoolExecutor``, eliminating thread creation
        overhead while providing explicit back-pressure across
        concurrent batches.

        Args:
            batch_prompts: List of prompts to process.
            **kwargs: Additional inference parameters.
                ``pass_num`` (int) is consumed internally: when >= 1
                the call includes ``cache={"no-cache": True}`` so
                that repeat extraction passes bypass the LiteLLM
                response cache.

        Returns:
            List of lists of ScoredOutput objects, one per prompt.
        """
        pass_num: int = kwargs.pop("pass_num", 0)
        semaphore = self._get_semaphore()

        # Build per-call kwargs: instance defaults + cache bypass.
        call_kwargs = dict(self._litellm_kwargs)
        if pass_num >= 1:
            call_kwargs["cache"] = {"no-cache": True}

        logger.info(
            "Running async inference for %d prompt(s) on %s",
            len(batch_prompts),
            self.model_id,
        )

        async def _process_single(prompt: str) -> list[ScoredOutput]:
            async with semaphore:
                try:
                    logger.debug(
                        "Calling LiteLLM acompletion for model %s",
                        self.model_id,
                    )
                    messages = [{"role": "user", "content": prompt}]

                    response = await litellm.acompletion(
                        model=self.model_id,
                        messages=messages,
                        **call_kwargs,
                    )

                    return self._parse_response(response)

                except (
                    LiteLLMAPIError,
                    LiteLLMAuthenticationError,
                    LiteLLMBadRequestError,
                    LiteLLMConnectionError,
                    LiteLLMNotFoundError,
                    LiteLLMPermissionDeniedError,
                    LiteLLMRateLimitError,
                    LiteLLMTimeout,
                ) as e:
                    logger.warning(
                        "LiteLLM API error for model %s: %s: %s",
                        self.model_id,
                        type(e).__name__,
                        e,
                    )
                    return [
                        ScoredOutput(
                            score=0.0,
                            output="LLM inference failed",
                        )
                    ]
                except Exception:
                    logger.exception(
                        "Unexpected error during LiteLLM acompletion for model %s",
                        self.model_id,
                    )
                    return [
                        ScoredOutput(
                            score=0.0,
                            output="LLM inference failed",
                        )
                    ]

        tasks = [_process_single(prompt) for prompt in batch_prompts]
        return list(await asyncio.gather(*tasks))
