# 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
import logging
from abc import ABC, abstractmethod
from collections.abc import AsyncIterator
from typing import Any

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

from klira.sdk.guardrails.augmentation import build_augmented_messages, verify_augmentation
from klira.sdk.guardrails.guideline_context import get_current_guidelines, reset_current_guidelines

logger = logging.getLogger(__name__)

_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)

    def _get_and_clear_guidelines(self) -> list[str] | None:
        """Read current guidelines from context and clear them.

        Clears immediately to prevent double-injection if multiple LLM calls
        happen within the same decorated function.

        Returns:
            The guidelines list, or None if no guidelines are active.
        """
        guidelines = get_current_guidelines()
        if guidelines is not None:
            # Clear immediately to prevent double-injection on subsequent
            # LLM calls within the same decorated function.
            # The decorator's try/finally will do the token-based reset.
            reset_current_guidelines()
        return guidelines

    def _augment_standard_messages(self, kwargs: dict[str, Any], guidelines: list[str]) -> None:
        """Augment standard messages-format kwargs with guidelines.

        Modifies kwargs["messages"] in place with guideline injection.

        Args:
            kwargs: The LLM call kwargs containing a "messages" key.
            guidelines: Guideline strings to inject.
        """
        messages = kwargs.get("messages", [])
        augmented = build_augmented_messages(messages, guidelines)
        kwargs["messages"] = augmented
        if verify_augmentation(messages, augmented, guidelines):
            logger.debug("Guidelines augmented into messages")

