# 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 Responses API adapter.

Patches the OpenAI Responses API (responses.create) to create
klira.llm.openai.responses 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 = ".responses"


class OpenAIResponsesAdapter(BaseLLMAdapter):
    """Adapter for OpenAI Responses API."""

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

    def _resolve_responses_cls(self) -> type | None:
        """Resolve the real Responses resource class."""
        try:
            from openai.resources.responses.responses import Responses

            return Responses
        except (ImportError, AttributeError):
            return None

    def _extract_response_attrs(self, response: Any, model: str) -> dict[str, Any]:
        """Extract common response attributes from an OpenAI response."""
        resp_model = getattr(response, "model", model)
        usage = getattr(response, "usage", None)
        input_tokens = getattr(usage, "input_tokens", None) if usage else None
        output_tokens = getattr(usage, "output_tokens", None) if usage else None
        return {
            "model": resp_model,
            "input_tokens": input_tokens,
            "output_tokens": output_tokens,
        }

    def _extract_output_text(self, response: Any) -> str:
        """Extract output text from response output items."""
        output_items = getattr(response, "output", [])
        parts: list[str] = []
        for item in output_items:
            content_list = getattr(item, "content", [])
            for content in content_list:
                text = getattr(content, "text", None)
                if text:
                    parts.append(text)
        return "".join(parts)

    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 Responses resource's create method.

        Patches both sync ``Responses.create`` and async
        ``AsyncResponses.create``.
        """
        adapter = self

        responses_cls = self._resolve_responses_cls() or client_class

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

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

        def patched_create(self_responses: Any, *args: Any, **kwargs: Any) -> Any:
            model = kwargs.get("model", "unknown")

            with adapter.create_llm_span(PROVIDER, OPERATION) as span:
                adapter.set_request_attributes(span, model)
                response = original_create(self_responses, *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)
        responses_cls.create = patched_create  # type: ignore[attr-defined]

        # Patch AsyncResponses.create
        try:
            from openai.resources.responses.responses import AsyncResponses

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

                async def patched_async_create(
                    self_responses: Any, *args: Any, **kwargs: Any
                ) -> Any:
                    model = kwargs.get("model", "unknown")

                    async with adapter.create_async_llm_span(PROVIDER, OPERATION) as span:
                        adapter.set_request_attributes(span, model)
                        response = await original_async_create(
                            self_responses, *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)
                AsyncResponses.create = patched_async_create  # type: ignore[attr-defined]
                self._mark_as_patched(AsyncResponses)
        except (ImportError, AttributeError):
            pass

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

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