# Copyright 2024 Klira AI
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""LiteLLM adapter.

Patches litellm.completion to create klira.llm.litellm spans
with standard gen_ai attributes.

Handles falsy Pydantic model correctly for token extraction (PROD-437 fix).
"""

from __future__ import annotations

import logging
from typing import Any

from klira.sdk.adapters.base_llm import BaseLLMAdapter

logger = logging.getLogger(__name__)

PROVIDER = "litellm"


def _safe_get_attr(obj: Any, attr: str, default: Any = None) -> Any:
    """Safely get attribute from potentially falsy Pydantic model.

    LiteLLM uses Pydantic models that can be falsy (empty) but still
    have attributes. Using `if usage:` would skip valid zero-token responses.
    """
    if obj is None:
        return default
    return getattr(obj, attr, default)


class LiteLLMAdapter(BaseLLMAdapter):
    """Adapter for LiteLLM."""

    # ------------------------------------------------------------------
    # Helpers shared by sync and async patches
    # ------------------------------------------------------------------

    def _extract_response_attrs(self, response: Any, model: str) -> dict[str, Any]:
        """Extract common response attributes from a LiteLLM response."""
        resp_model = _safe_get_attr(response, "model", model)
        usage = _safe_get_attr(response, "usage")
        input_tokens = _safe_get_attr(usage, "prompt_tokens")
        output_tokens = _safe_get_attr(usage, "completion_tokens")

        choices = _safe_get_attr(response, "choices", [])
        finish_reasons = []
        if choices:
            for choice in choices:
                reason = _safe_get_attr(choice, "finish_reason")
                if reason:
                    finish_reasons.append(str(reason))

        return {
            "model": resp_model,
            "input_tokens": input_tokens,
            "output_tokens": output_tokens,
            "finish_reasons": finish_reasons,
        }

    def _extract_output_text(self, response: Any) -> str:
        """Extract output text from the first choice."""
        choices = _safe_get_attr(response, "choices", [])
        if choices:
            message = _safe_get_attr(choices[0], "message")
            content = _safe_get_attr(message, "content", "")
            return content or ""
        return ""

    def _set_span_output(self, span: Any, output_text: str) -> None:
        """Set klira.output attribute on the span."""
        if output_text:
            if len(output_text) > self.OUTPUT_TRUNCATION_LIMIT:
                output_text = output_text[: self.OUTPUT_TRUNCATION_LIMIT]
            span.set_attribute("klira.output", output_text)

    # ------------------------------------------------------------------
    # patch()
    # ------------------------------------------------------------------

    def patch(self, client_class: type) -> None:
        """Patch litellm.completion and litellm.acompletion functions."""
        if self._is_already_patched(client_class):
            logger.debug("LiteLLM module already patched, skipping")
            return

        adapter = self

        # LiteLLM is a module, not a class — client_class here is the module
        original_completion = client_class.completion  # type: ignore[attr-defined]

        def patched_completion(*args: Any, **kwargs: Any) -> Any:
            model = kwargs.get("model", args[0] if args else "unknown")
            messages = kwargs.get("messages", [])

            with adapter.create_llm_span(PROVIDER) as span:
                adapter.set_request_attributes(span, str(model), messages)
                response = original_completion(*args, **kwargs)
                attrs = adapter._extract_response_attrs(response, str(model))
                adapter.set_response_attributes(span, **attrs)
                adapter._set_span_output(span, adapter._extract_output_text(response))
                return response

        self._mark_function(patched_completion)
        client_class.completion = patched_completion  # type: ignore[attr-defined]

        # Patch acompletion (async version)
        original_acompletion = getattr(client_class, "acompletion", None)
        if original_acompletion is not None:

            async def patched_acompletion(*args: Any, **kwargs: Any) -> Any:
                model = kwargs.get("model", args[0] if args else "unknown")
                messages = kwargs.get("messages", [])

                async with adapter.create_async_llm_span(PROVIDER) as span:
                    adapter.set_request_attributes(span, str(model), messages)
                    response = await original_acompletion(*args, **kwargs)
                    attrs = adapter._extract_response_attrs(response, str(model))
                    adapter.set_response_attributes(span, **attrs)
                    adapter._set_span_output(
                        span, adapter._extract_output_text(response)
                    )
                    return response

            self._mark_function(patched_acompletion)
            client_class.acompletion = patched_acompletion  # type: ignore[attr-defined]

        self._mark_as_patched(client_class)
        self.verify_patch(client_class)

    def verify_patch(self, client_class: type) -> bool:
        """Verify the patch took effect by inspecting the actual method."""
        method = getattr(client_class, "completion", None)
        if method is None or not getattr(method, "_klira_patched", False):
            raise RuntimeError(
                f"LiteLLM adapter: patch not applied to {client_class}"
            )
        return True
