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

import time
from typing import TYPE_CHECKING, Any, Dict

from .base import calculate_cost, extract_mandatum_params

if TYPE_CHECKING:
    from ..client import Mandatum

try:
    from anthropic import Anthropic as OriginalAnthropic
    from anthropic import AsyncAnthropic as OriginalAsyncAnthropic

    ANTHROPIC_AVAILABLE = True
except ImportError:
    ANTHROPIC_AVAILABLE = False
    OriginalAnthropic = object
    OriginalAsyncAnthropic = object


class AnthropicWrapper:
    """Wrapper for Anthropic SDK."""

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

    def Anthropic(self, **kwargs):
        """Return wrapped Anthropic client."""
        return WrappedAnthropic(self.mandatum, **kwargs)

    def AsyncAnthropic(self, **kwargs):
        """Return wrapped AsyncAnthropic client."""
        return WrappedAsyncAnthropic(self.mandatum, **kwargs)


class WrappedAnthropic(OriginalAnthropic):
    """Wrapped Anthropic 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 messages with wrapped version
        self.messages = WrappedMessages(self.messages, mandatum_client, self._logger)


class WrappedMessages:
    """Wrapped messages API 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 Anthropic call
            response = self._original.create(**provider_kwargs)

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

            # Extract token counts
            input_tokens = response.usage.input_tokens if hasattr(response, "usage") else 0
            output_tokens = response.usage.output_tokens if hasattr(response, "usage") else 0

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

            # Extract response content
            response_content = ""
            if response.content and len(response.content) > 0:
                response_content = response.content[0].text if hasattr(response.content[0], "text") else ""

            # Log to Mandatum asynchronously
            if self.mandatum.async_logging:
                self._logger.log_request(
                    provider="anthropic",
                    model=provider_kwargs.get("model", ""),
                    request_data={
                        "messages": provider_kwargs.get("messages", []),
                        "temperature": provider_kwargs.get("temperature"),
                        "max_tokens": provider_kwargs.get("max_tokens"),
                        "system": provider_kwargs.get("system"),
                    },
                    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] Anthropic 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="anthropic",
                    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 messages."""
        return getattr(self._original, name)


class WrappedAsyncAnthropic(OriginalAsyncAnthropic):
    """Wrapped AsyncAnthropic 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 messages with wrapped version
        self.messages = WrappedAsyncMessages(self.messages, mandatum_client, self._logger)


class WrappedAsyncMessages:
    """Wrapped async messages API."""

    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 Anthropic 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.input_tokens if hasattr(response, "usage") else 0
            output_tokens = response.usage.output_tokens if hasattr(response, "usage") else 0

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

            # Extract response content
            response_content = ""
            if response.content and len(response.content) > 0:
                response_content = response.content[0].text if hasattr(response.content[0], "text") else ""

            # Log to Mandatum
            if self.mandatum.async_logging:
                self._logger.log_request(
                    provider="anthropic",
                    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="anthropic",
                    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 messages."""
        return getattr(self._original, name)
