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

"""Ollama adapter.

Patches ollama.Client.chat to create klira.llm.ollama 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 = "ollama"


class OllamaAdapter(BaseLLMAdapter):
    """Adapter for Ollama API."""

    def patch(self, client_class: type) -> None:
        """Patch the Ollama client's chat method."""
        if self._is_already_patched(client_class):
            logger.debug("Ollama Client already patched, skipping")
            return

        adapter = self

        original_chat = client_class.chat  # type: ignore[attr-defined]

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

            with adapter.create_llm_span(PROVIDER) as span:
                adapter.set_request_attributes(span, model, messages)

                response = original_chat(self_client, *args, **kwargs)

                # Ollama response is a dict
                if isinstance(response, dict):
                    resp_model = response.get("model", model)
                    input_tokens = response.get("prompt_eval_count")
                    output_tokens = response.get("eval_count")
                    done_reason = response.get("done_reason")

                    adapter.set_response_attributes(
                        span,
                        model=resp_model,
                        input_tokens=input_tokens,
                        output_tokens=output_tokens,
                        finish_reasons=[done_reason] if done_reason else [],
                    )

                    # Capture output
                    message = response.get("message", {})
                    content = message.get("content", "")
                    if content:
                        if len(content) > adapter.OUTPUT_TRUNCATION_LIMIT:
                            content = content[: adapter.OUTPUT_TRUNCATION_LIMIT]
                        span.set_attribute("klira.output", content)

                return response

        self._mark_function(patched_chat)
        client_class.chat = patched_chat  # 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, "chat", None)
        if method is None or not getattr(method, "_klira_patched", False):
            raise RuntimeError(
                f"Ollama adapter: patch not applied to {client_class}"
            )
        return True
