# 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.

"""OpenAI Chat Completions adapter.

Patches openai.ChatCompletion.create / openai.chat.completions.create
to create klira.llm.openai.completion spans with standard gen_ai attributes.
"""

from __future__ import annotations

import logging
from typing import Any

from klira.sdk.adapters.base_llm import BaseLLMAdapter

logger = logging.getLogger(__name__)

PROVIDER = "openai"
OPERATION = ".completion"


class OpenAICompletionAdapter(BaseLLMAdapter):
    """Adapter for OpenAI Chat Completions API."""

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

    def _resolve_completions_cls(self) -> type | None:
        """Resolve the real Completions resource class."""
        try:
            from openai.resources.chat.completions.completions import Completions

            return Completions
        except (ImportError, AttributeError):
            return None

    def _extract_response_attrs(self, response: Any, model: str) -> dict[str, Any]:
        """Extract common response attributes from an OpenAI completion."""
        resp_model = getattr(response, "model", model)
        usage = getattr(response, "usage", None)
        input_tokens = getattr(usage, "prompt_tokens", None) if usage else None
        output_tokens = (
            getattr(usage, "completion_tokens", None) if usage else None
        )
        choices = getattr(response, "choices", [])
        finish_reasons = [getattr(c, "finish_reason", "unknown") for c in choices]
        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 = getattr(response, "choices", [])
        if choices:
            content = getattr(
                getattr(choices[0], "message", None), "content", None
            )
            if content:
                return str(content)
        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 the OpenAI Completions resource's create method.

        Patches both sync ``Completions.create`` and async
        ``AsyncCompletions.create`` so that LLM spans are captured
        regardless of whether the caller uses sync or async.
        """
        adapter = self

        completions_cls = self._resolve_completions_cls() or client_class

        if self._is_already_patched(completions_cls):
            logger.debug("OpenAI Completions already patched, skipping")
            return

        original_create = completions_cls.create  # type: ignore[attr-defined]

        def patched_create(self_completions: Any, *args: Any, **kwargs: Any) -> Any:
            model = kwargs.get("model", "unknown")
            messages = kwargs.get("messages", [])

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

        self._mark_function(patched_create)
        completions_cls.create = patched_create  # type: ignore[attr-defined]

        # Patch AsyncCompletions.create (used by async callers)
        try:
            from openai.resources.chat.completions.completions import AsyncCompletions

            if not self._is_already_patched(AsyncCompletions):
                original_async_create = AsyncCompletions.create  # type: ignore[attr-defined]

                async def patched_async_create(
                    self_completions: Any, *args: Any, **kwargs: Any
                ) -> Any:
                    model = kwargs.get("model", "unknown")
                    messages = kwargs.get("messages", [])

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

                self._mark_function(patched_async_create)
                AsyncCompletions.create = patched_async_create  # type: ignore[attr-defined]
                self._mark_as_patched(AsyncCompletions)
        except (ImportError, AttributeError):
            pass

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

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