"""OpenAI SDK wrapper with automatic Mandatum logging."""

import time
from typing import TYPE_CHECKING, Any, Dict, Optional

from .base import calculate_cost, extract_mandatum_params

if TYPE_CHECKING:
    from ..client import Mandatum

try:
    from openai import OpenAI as OriginalOpenAI
    from openai import AsyncOpenAI as OriginalAsyncOpenAI

    OPENAI_AVAILABLE = True
except ImportError:
    OPENAI_AVAILABLE = False
    OriginalOpenAI = object
    OriginalAsyncOpenAI = object


class OpenAIWrapper:
    """Wrapper for OpenAI SDK."""

    def __init__(self, mandatum_client: "Mandatum"):
        if not OPENAI_AVAILABLE:
            raise ImportError(
                "OpenAI SDK is not installed. Install it with: pip install openai"
            )
        self.mandatum = mandatum_client

    def OpenAI(self, **kwargs):
        """Return wrapped OpenAI client."""
        return WrappedOpenAI(self.mandatum, **kwargs)

    def AsyncOpenAI(self, **kwargs):
        """Return wrapped AsyncOpenAI client."""
        return WrappedAsyncOpenAI(self.mandatum, **kwargs)


class WrappedOpenAI(OriginalOpenAI):
    """Wrapped OpenAI client that logs to Mandatum."""

    def __init__(self, mandatum_client: "Mandatum", **kwargs):
        super().__init__(**kwargs)
        self.mandatum = mandatum_client

        # Import logging service
        from ..logging_service import LoggingService

        self._logger = LoggingService(
            api_key=mandatum_client.api_key,
            base_url=mandatum_client.base_url,
            organization_id=mandatum_client.organization_id,
            debug=mandatum_client.debug,
        )

        # Replace chat.completions with wrapped version
        self.chat.completions = WrappedCompletions(
            self.chat.completions, mandatum_client, self._logger
        )


class WrappedCompletions:
    """Wrapped chat completions that intercepts create() calls."""

    def __init__(self, original, mandatum_client: "Mandatum", logger):
        self._original = original
        self.mandatum = mandatum_client
        self._logger = logger

    def create(self, **kwargs):
        """Intercept create() call to log to Mandatum."""
        # Extract Mandatum-specific params
        provider_kwargs, mandatum_params = extract_mandatum_params(kwargs)

        # Capture start time
        start_time = time.time()
        error = None

        try:
            # Make actual OpenAI call
            response = self._original.create(**provider_kwargs)

            # Calculate latency
            latency_ms = int((time.time() - start_time) * 1000)

            # Extract token counts
            input_tokens = response.usage.prompt_tokens if response.usage else 0
            output_tokens = response.usage.completion_tokens if response.usage else 0

            # Calculate cost
            cost = calculate_cost(
                provider="openai",
                model=provider_kwargs.get("model", ""),
                input_tokens=input_tokens,
                output_tokens=output_tokens,
            )

            # Extract response content
            response_content = ""
            if response.choices and len(response.choices) > 0:
                response_content = response.choices[0].message.content or ""

            # Log to Mandatum asynchronously
            if self.mandatum.async_logging:
                self._logger.log_request(
                    provider="openai",
                    model=provider_kwargs.get("model", ""),
                    request_data={
                        "messages": provider_kwargs.get("messages", []),
                        "temperature": provider_kwargs.get("temperature"),
                        "max_tokens": provider_kwargs.get("max_tokens"),
                        "top_p": provider_kwargs.get("top_p"),
                        "frequency_penalty": provider_kwargs.get("frequency_penalty"),
                        "presence_penalty": provider_kwargs.get("presence_penalty"),
                    },
                    response_data={"content": response_content},
                    latency_ms=latency_ms,
                    input_tokens=input_tokens,
                    output_tokens=output_tokens,
                    cost=cost,
                    environment=mandatum_params.get("environment"),
                    tags=mandatum_params.get("tags"),
                    metadata=mandatum_params.get("metadata"),
                    prompt_id=mandatum_params.get("prompt_id"),
                )

            if self.mandatum.debug:
                print(
                    f"[Mandatum] OpenAI request logged: "
                    f"model={provider_kwargs.get('model')}, "
                    f"latency={latency_ms}ms, "
                    f"tokens={input_tokens}/{output_tokens}, "
                    f"cost=${cost:.6f}"
                )

            return response

        except Exception as e:
            latency_ms = int((time.time() - start_time) * 1000)
            error = str(e)

            # Log error to Mandatum
            if self.mandatum.async_logging:
                self._logger.log_request(
                    provider="openai",
                    model=provider_kwargs.get("model", ""),
                    request_data={
                        "messages": provider_kwargs.get("messages", []),
                        "temperature": provider_kwargs.get("temperature"),
                        "max_tokens": provider_kwargs.get("max_tokens"),
                    },
                    response_data={},
                    latency_ms=latency_ms,
                    input_tokens=0,
                    output_tokens=0,
                    cost=0.0,
                    environment=mandatum_params.get("environment"),
                    tags=mandatum_params.get("tags"),
                    metadata=mandatum_params.get("metadata"),
                    prompt_id=mandatum_params.get("prompt_id"),
                    error=error,
                )

            # Re-raise the original exception
            raise

    def __getattr__(self, name):
        """Forward other attributes to original completions."""
        return getattr(self._original, name)


class WrappedAsyncOpenAI(OriginalAsyncOpenAI):
    """Wrapped AsyncOpenAI client (async support)."""

    def __init__(self, mandatum_client: "Mandatum", **kwargs):
        super().__init__(**kwargs)
        self.mandatum = mandatum_client

        # Import logging service
        from ..logging_service import LoggingService

        self._logger = LoggingService(
            api_key=mandatum_client.api_key,
            base_url=mandatum_client.base_url,
            organization_id=mandatum_client.organization_id,
            debug=mandatum_client.debug,
        )

        # Replace chat.completions with wrapped version
        self.chat.completions = WrappedAsyncCompletions(
            self.chat.completions, mandatum_client, self._logger
        )


class WrappedAsyncCompletions:
    """Wrapped async chat completions."""

    def __init__(self, original, mandatum_client: "Mandatum", logger):
        self._original = original
        self.mandatum = mandatum_client
        self._logger = logger

    async def create(self, **kwargs):
        """Intercept async create() call."""
        # Extract Mandatum-specific params
        provider_kwargs, mandatum_params = extract_mandatum_params(kwargs)

        # Capture start time
        start_time = time.time()
        error = None

        try:
            # Make actual OpenAI call
            response = await self._original.create(**provider_kwargs)

            # Calculate latency
            latency_ms = int((time.time() - start_time) * 1000)

            # Extract token counts
            input_tokens = response.usage.prompt_tokens if response.usage else 0
            output_tokens = response.usage.completion_tokens if response.usage else 0

            # Calculate cost
            cost = calculate_cost(
                provider="openai",
                model=provider_kwargs.get("model", ""),
                input_tokens=input_tokens,
                output_tokens=output_tokens,
            )

            # Extract response content
            response_content = ""
            if response.choices and len(response.choices) > 0:
                response_content = response.choices[0].message.content or ""

            # Log to Mandatum
            if self.mandatum.async_logging:
                self._logger.log_request(
                    provider="openai",
                    model=provider_kwargs.get("model", ""),
                    request_data={
                        "messages": provider_kwargs.get("messages", []),
                        "temperature": provider_kwargs.get("temperature"),
                        "max_tokens": provider_kwargs.get("max_tokens"),
                    },
                    response_data={"content": response_content},
                    latency_ms=latency_ms,
                    input_tokens=input_tokens,
                    output_tokens=output_tokens,
                    cost=cost,
                    environment=mandatum_params.get("environment"),
                    tags=mandatum_params.get("tags"),
                    metadata=mandatum_params.get("metadata"),
                    prompt_id=mandatum_params.get("prompt_id"),
                )

            return response

        except Exception as e:
            latency_ms = int((time.time() - start_time) * 1000)
            error = str(e)

            # Log error
            if self.mandatum.async_logging:
                self._logger.log_request(
                    provider="openai",
                    model=provider_kwargs.get("model", ""),
                    request_data={
                        "messages": provider_kwargs.get("messages", []),
                    },
                    response_data={},
                    latency_ms=latency_ms,
                    input_tokens=0,
                    output_tokens=0,
                    cost=0.0,
                    error=error,
                )

            raise

    def __getattr__(self, name):
        """Forward other attributes to original completions."""
        return getattr(self._original, name)
