"""Adapter for patching the OpenAI Chat Completion API."""

import copy
import functools
import logging
from typing import Any, Dict, Optional

from klira.sdk.adapters.llm_base_adapter import BaseLLMAdapter

# Try to import OpenAI client for patching
try:
    import openai
    import openai.resources.chat.completions as completions_module

    OPENAI_CLIENT_AVAILABLE = True
except ImportError:
    openai = None  # type: ignore[assignment]
    completions_module = None  # type: ignore[assignment]
    OPENAI_CLIENT_AVAILABLE = False

logger = logging.getLogger("klira.adapters.openai_completion")


class OpenAICompletionAdapter(BaseLLMAdapter):
    """Patches OpenAI Chat Completion API calls for guideline injection.

    Patches at the CLASS level (Completions.create and AsyncCompletions.create)
    so that ALL client instances (OpenAI, AsyncOpenAI) are instrumented,
    regardless of how they are instantiated or imported.
    """

    is_available = OPENAI_CLIENT_AVAILABLE

    def __init__(self) -> None:
        super().__init__()
        from typing import Callable

        self.original_create: Optional[Callable[..., Any]] = None
        self.original_acreate: Optional[Callable[..., Any]] = None
        self._patched: bool = False

    def patch(self) -> None:
        """Patch the OpenAI Completions and AsyncCompletions classes."""
        if self._patched:
            logger.debug("OpenAI client already patched")
            return

        if not OPENAI_CLIENT_AVAILABLE or completions_module is None:
            logger.debug("OpenAI client not available. Skipping patch.")
            return

        logger.debug(
            "OpenAICompletionAdapter: Attempting to patch OpenAI completions classes..."
        )

        self._patch_sync_create()
        self._patch_async_create()

        self._patched = True
        logger.info("OpenAI adapter patching complete.")

    def _patch_sync_create(self) -> None:
        """Patch the synchronous Completions.create class method."""
        if not OPENAI_CLIENT_AVAILABLE or completions_module is None:
            return

        try:
            target_cls = getattr(completions_module, "Completions", None)
            if target_cls is None:
                logger.warning("Could not find Completions class to patch")
                return

            method_name = "create"
            if hasattr(target_cls, method_name) and not hasattr(
                getattr(target_cls, method_name), "_klira_augmented"
            ):
                original_create = getattr(target_cls, method_name)
                self.original_create = original_create

                adapter_self = self

                @functools.wraps(original_create)
                def patched_create(*args: Any, **kwargs: Any) -> Any:
                    return adapter_self._instrumented_create(
                        original_create, *args, **kwargs
                    )

                setattr(patched_create, "_klira_augmented", True)
                setattr(target_cls, method_name, patched_create)
                logger.info(
                    "Successfully patched Completions.create for augmentation."
                )
            elif hasattr(getattr(target_cls, method_name, None), "_klira_augmented"):
                logger.debug("Completions.create already patched.")
            else:
                logger.warning("Could not find Completions.create to patch.")
        except Exception as e:
            logger.error(f"Error patching sync OpenAI client: {e}", exc_info=True)

    def _patch_async_create(self) -> None:
        """Patch the asynchronous AsyncCompletions.create class method."""
        if not OPENAI_CLIENT_AVAILABLE or completions_module is None:
            return

        try:
            async_target_cls = getattr(completions_module, "AsyncCompletions", None)
            if async_target_cls is None:
                logger.debug("AsyncCompletions class not found, skipping async patch.")
                return

            method_name = "create"
            if not hasattr(async_target_cls, method_name):
                logger.debug(
                    "Could not find create method on AsyncCompletions, skipping patch."
                )
                return

            target_method = getattr(async_target_cls, method_name)
            if not hasattr(target_method, "_klira_augmented"):
                original_acreate = target_method
                self.original_acreate = original_acreate

                adapter_self = self

                @functools.wraps(original_acreate)
                async def patched_acreate(*args: Any, **kwargs: Any) -> Any:
                    return await adapter_self._instrumented_acreate(
                        original_acreate, *args, **kwargs
                    )

                setattr(patched_acreate, "_klira_augmented", True)
                setattr(async_target_cls, method_name, patched_acreate)
                logger.info(
                    "Successfully patched AsyncCompletions.create for augmentation."
                )
            elif hasattr(target_method, "_klira_augmented"):
                logger.debug("AsyncCompletions.create already patched.")

        except Exception as e:
            logger.error(f"Error patching async OpenAI client: {e}", exc_info=True)

    def _instrumented_create(
        self, original_fn: Any, *args: Any, **kwargs: Any
    ) -> Any:
        """Synchronous instrumented wrapper for OpenAI completions."""
        from klira.sdk.tracing.tracing import get_tracer as get_traceloop_tracer
        from opentelemetry.trace import StatusCode
        from opentelemetry import trace as otel_trace

        model = kwargs.get("model", "unknown")
        messages = kwargs.get("messages", [])
        original_messages = copy.deepcopy(messages) if messages else None

        if self._is_in_unified_trace():
            tracer = otel_trace.get_tracer("klira")
        else:
            tracer = get_traceloop_tracer()

        with tracer.start_as_current_span("klira.llm.openai") as llm_span:
            try:
                self._create_llm_request_span(
                    llm_span,
                    provider="openai",
                    model=model,
                    messages=messages,
                    temperature=kwargs.get("temperature"),
                    max_tokens=kwargs.get("max_tokens"),
                )

                from klira.sdk.tracing.tracing import set_framework_metadata

                set_framework_metadata(llm_span, "openai")

                self._add_augmentation_attributes(
                    llm_span, original_messages=original_messages
                )

                self._inject_guidelines_into_messages(kwargs)

            except Exception as e:
                logger.error(f"Policy injection error: {str(e)}")

            try:
                result = original_fn(*args, **kwargs)

                self._add_llm_response_attributes(llm_span, result, "openai")
                llm_span.set_status(StatusCode.OK)

            except Exception as e:
                llm_span.set_status(StatusCode.ERROR, str(e))
                llm_span.record_exception(e)
                raise

            return result

    async def _instrumented_acreate(
        self, original_fn: Any, *args: Any, **kwargs: Any
    ) -> Any:
        """Async instrumented wrapper for OpenAI completions."""
        from klira.sdk.tracing.tracing import get_tracer as get_traceloop_tracer
        from opentelemetry.trace import StatusCode
        from opentelemetry import trace as otel_trace

        model = kwargs.get("model", "unknown")
        messages = kwargs.get("messages", [])
        original_messages = copy.deepcopy(messages) if messages else None

        if self._is_in_unified_trace():
            tracer = otel_trace.get_tracer("klira")
        else:
            tracer = get_traceloop_tracer()

        with tracer.start_as_current_span("klira.llm.openai") as llm_span:
            try:
                self._create_llm_request_span(
                    llm_span,
                    provider="openai",
                    model=model,
                    messages=messages,
                    temperature=kwargs.get("temperature"),
                    max_tokens=kwargs.get("max_tokens"),
                )

                from klira.sdk.tracing.tracing import set_framework_metadata

                set_framework_metadata(llm_span, "openai")

                self._add_augmentation_attributes(
                    llm_span, original_messages=original_messages
                )

                self._inject_guidelines_into_messages(kwargs)

            except Exception as e:
                logger.error(f"Policy injection error in async completion: {str(e)}")

            try:
                result = await original_fn(*args, **kwargs)

                self._add_llm_response_attributes(llm_span, result, "openai")
                llm_span.set_status(StatusCode.OK)

            except Exception as e:
                llm_span.set_status(StatusCode.ERROR, str(e))
                llm_span.record_exception(e)
                raise

            return result

    def _inject_guidelines_into_messages(self, kwargs: Dict[str, Any]) -> None:
        """Inject policy guidelines into the messages."""
        messages = kwargs.get("messages")
        if not messages:
            return

        system_msg = next(
            (m for m in messages if m.get("role") == "system"), None
        )
        if system_msg and isinstance(system_msg.get("content"), str):
            from klira.sdk.guardrails.engine import GuardrailsEngine

            guidelines = GuardrailsEngine.get_current_guidelines()
            if guidelines:
                augmentation_text = "\n\nIMPORTANT POLICY DIRECTIVES:\n"
                augmentation_text += "\n".join([f"• {g}" for g in guidelines])
                system_msg["content"] += augmentation_text
                GuardrailsEngine.clear_current_guidelines()
                logger.info(
                    f"Injected {len(guidelines)} policy guidelines into system prompt"
                )

    def _extract_prompt_text(self, messages: Any, provider: str) -> str:
        """Extract prompt text from OpenAI messages."""
        if isinstance(messages, list):
            parts = []
            for msg in messages:
                if isinstance(msg, dict) and "content" in msg:
                    parts.append(f"{msg.get('role', 'user')}: {msg['content']}")
            return "\n".join(parts)
        return str(messages)

    def _extract_response_text(self, response: Any, provider: str) -> str:
        """Extract response text from OpenAI response."""
        try:
            if hasattr(response, "choices") and response.choices:
                choice = response.choices[0]
                if hasattr(choice, "message") and hasattr(choice.message, "content"):
                    return str(choice.message.content) if choice.message.content else ""
            return ""
        except Exception as e:
            logger.debug(f"Error extracting OpenAI response text: {e}")
            return ""

    def _extract_token_usage(
        self, response: Any, provider: str
    ) -> Optional[Dict[str, int]]:
        """Extract token usage from OpenAI response."""
        try:
            usage = getattr(response, "usage", None)
            if usage is None:
                return None
            input_tokens = getattr(usage, "prompt_tokens", None)
            output_tokens = getattr(usage, "completion_tokens", None)
            if input_tokens is not None or output_tokens is not None:
                result: Dict[str, int] = {}
                if input_tokens is not None:
                    result["input_tokens"] = int(input_tokens)
                if output_tokens is not None:
                    result["output_tokens"] = int(output_tokens)
                return result
            return None
        except Exception as e:
            logger.debug(f"Error extracting OpenAI token usage: {e}")
            return None

    def _extract_finish_reason(self, response: Any, provider: str) -> Optional[str]:
        """Extract finish reason from OpenAI response."""
        try:
            if hasattr(response, "choices") and response.choices:
                return response.choices[0].finish_reason
            return None
        except Exception as e:
            logger.debug(f"Error extracting OpenAI finish reason: {e}")
            return None

    def _extract_response_model(self, response: Any, provider: str) -> Optional[str]:
        """Extract model name from OpenAI response."""
        try:
            if hasattr(response, "model") and response.model is not None:
                return str(response.model)
            return None
        except Exception as e:
            logger.debug(f"Error extracting OpenAI response model: {e}")
            return None

    def unpatch_llm_client(self) -> bool:
        """Restore original OpenAI client methods."""
        if not OPENAI_CLIENT_AVAILABLE or not self._patched:
            return False

        try:
            if completions_module is None:
                return False  # type: ignore[unreachable]

            # Restore Completions.create
            if self.original_create:
                target_cls = getattr(completions_module, "Completions", None)
                if target_cls:
                    setattr(target_cls, "create", self.original_create)
                    self.original_create = None

            # Restore AsyncCompletions.create
            if self.original_acreate:
                async_target_cls = getattr(
                    completions_module, "AsyncCompletions", None
                )
                if async_target_cls:
                    setattr(async_target_cls, "create", self.original_acreate)
                    self.original_acreate = None

            self._patched = False
            logger.info("Successfully unpatched OpenAI client methods")
            return True

        except Exception as e:
            logger.error(f"Failed to unpatch OpenAI client: {e}")
            return False

    def get_client_info(self) -> Dict[str, Any]:
        """Get information about the OpenAI client."""
        if not OPENAI_CLIENT_AVAILABLE:
            return {"available": False, "reason": "openai not installed"}

        try:
            return {
                "available": True,
                "library": "openai",
                "version": getattr(openai, "__version__", "unknown"),
                "patched": self._patched,
            }
        except Exception as e:
            return {"available": False, "reason": str(e)}
