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

"""Abstract base class for LLM client adapters.

Every LLM adapter (OpenAI, Anthropic, Gemini, Ollama, LiteLLM) inherits
from this class and implements `patch()` and `verify_patch()`.

The base class provides shared utilities for span creation, attribute
setting, and outbound guardrails execution so that all adapters produce
identical span structures. (Learning #1: define trace schema as contract.)
"""

from __future__ import annotations

import contextlib
from abc import ABC, abstractmethod
from collections.abc import AsyncIterator
from typing import Any

from opentelemetry import trace
from opentelemetry.trace import Span, Tracer

_KLIRA_PATCH_SENTINEL = "_klira_patched"


class BaseLLMAdapter(ABC):
    """Abstract base for LLM client adapters.

    Subclasses must implement:
        - patch(): Monkey-patch the client's API call method.
        - verify_patch(): Confirm the patch took effect.
    """

    PROMPT_TRUNCATION_LIMIT: int = 10000
    OUTPUT_TRUNCATION_LIMIT: int = 5000

    def __init__(self) -> None:
        self._tracer: Tracer = trace.get_tracer("klira")

    def _is_already_patched(self, target_cls: type) -> bool:
        """Check if the target class has already been patched by Klira."""
        return getattr(target_cls, _KLIRA_PATCH_SENTINEL, False)

    def _mark_as_patched(self, target_cls: type) -> None:
        """Mark the target class as patched by Klira."""
        setattr(target_cls, _KLIRA_PATCH_SENTINEL, True)

    def _mark_function(self, func: Any) -> None:
        """Mark a wrapper function as a Klira patch."""
        func._klira_patched = True  # type: ignore[attr-defined]

    @abstractmethod
    def patch(self, client_class: type) -> None:
        """Monkey-patch the client's API call method.

        After patching, must call verify_patch() to confirm it took effect.
        Raises RuntimeError if verification fails. (Learning #14)
        """

    @abstractmethod
    def verify_patch(self, client_class: type) -> bool:
        """Confirm the patch took effect.

        Returns True if the patch is active. Raises RuntimeError if not.
        """

    def create_llm_span(self, provider: str, operation: str = "") -> Any:
        """Create a klira.llm.{provider} span using start_as_current_span.

        Args:
            provider: LLM provider name (lowercase, e.g., 'openai', 'anthropic').
            operation: Optional operation suffix (e.g., '.completion', '.responses').

        Returns:
            A context manager that yields a Span.
        """
        span_name = f"klira.llm.{provider}{operation}"
        return self._tracer.start_as_current_span(
            span_name,
            attributes={
                "klira.entity_type": "llm",
                "gen_ai.system": provider,
            },
        )

    @contextlib.asynccontextmanager
    async def create_async_llm_span(
        self, provider: str, operation: str = ""
    ) -> AsyncIterator[Span]:
        """Create a klira.llm.{provider} span for async LLM calls.

        Identical to create_llm_span but usable with ``async with``.

        Args:
            provider: LLM provider name (lowercase).
            operation: Optional operation suffix.

        Yields:
            The active Span.
        """
        span_name = f"klira.llm.{provider}{operation}"
        with self._tracer.start_as_current_span(
            span_name,
            attributes={
                "klira.entity_type": "llm",
                "gen_ai.system": provider,
            },
        ) as span:
            yield span

    def set_request_attributes(
        self,
        span: Span,
        model: str,
        messages: list[dict[str, Any]] | None = None,
    ) -> None:
        """Set standard request attributes on an LLM span.

        Args:
            span: The active LLM span.
            model: Model name (no provider prefix).
            messages: Optional message list for prompt capture.
        """
        span.set_attribute("gen_ai.request.model", model)
        if messages is not None:
            prompt_str = str(messages)
            if len(prompt_str) > self.PROMPT_TRUNCATION_LIMIT:
                prompt_str = prompt_str[: self.PROMPT_TRUNCATION_LIMIT]
            span.set_attribute("gen_ai.prompt", prompt_str)

    def set_response_attributes(
        self,
        span: Span,
        *,
        model: str | None = None,
        input_tokens: int | None = None,
        output_tokens: int | None = None,
        finish_reasons: list[str] | None = None,
    ) -> None:
        """Set standard response attributes on an LLM span.

        Args:
            span: The active LLM span.
            model: Actual model used in response.
            input_tokens: Input token count.
            output_tokens: Output token count.
            finish_reasons: Finish reasons from LLM.
        """
        if model is not None:
            span.set_attribute("gen_ai.response.model", model)
        if input_tokens is not None:
            span.set_attribute("gen_ai.usage.input_tokens", input_tokens)
        if output_tokens is not None:
            span.set_attribute("gen_ai.usage.output_tokens", output_tokens)
        if finish_reasons is not None:
            span.set_attribute("gen_ai.response.finish_reasons", finish_reasons)

