"""Anthropic integration with auto-instrumentation via monkey patching."""

import functools
import logging
from typing import Any

from fenra._context import get_context
from fenra._core import enqueue_transaction

logger = logging.getLogger(__name__)

_patched_sync = False
_patched_async = False
_patched_stream_sync = False
_patched_stream_async = False

# Store original methods for unpatching
_original_create: Any = None
_original_create_async: Any = None
_original_stream: Any = None
_original_stream_async: Any = None


def _detect_web_search(response: Any, tools: list[Any] | None = None) -> dict[str, Any] | None:
    """
    Detect web search tool usage from response content blocks.

    Returns a requests usage entry if web search was used, None otherwise.
    """
    if not hasattr(response, "content") or response.content is None:
        return None

    # Check for web_search_tool_result blocks in response content
    search_count = sum(
        1
        for block in response.content
        if getattr(block, "type", None) == "web_search_tool_result"
    )

    # Fallback: check for tool_use blocks with name == "web_search"
    if search_count == 0:
        search_count = sum(
            1
            for block in response.content
            if getattr(block, "type", None) == "tool_use"
            and getattr(block, "name", None) == "web_search"
        )

    # If still no count but web search tool was in the request, default to 1
    if search_count == 0 and tools:
        for tool in tools:
            tool_type = tool.get("type", "") if isinstance(tool, dict) else getattr(tool, "type", "")
            if "web_search" in str(tool_type):
                search_count = 1
                break

    if search_count > 0:
        return {
            "type": "requests",
            "metrics": {"count": search_count, "request_type": "web_search"},
        }

    return None


def _track_message_response(
    response: Any,
    model: str,
    context: dict[str, Any],
    tools: list[Any] | None = None,
) -> None:
    """Extract usage from Anthropic Message response and queue transaction."""
    if not hasattr(response, "usage") or response.usage is None:
        return

    usage = response.usage

    input_tokens = getattr(usage, "input_tokens", 0) or 0
    output_tokens = getattr(usage, "output_tokens", 0) or 0

    metrics: dict[str, Any] = {
        "input_tokens": input_tokens,
        "output_tokens": output_tokens,
        "total_tokens": input_tokens + output_tokens,
    }

    # Include cached tokens if present (prompt caching)
    cache_read = getattr(usage, "cache_read_input_tokens", 0) or 0
    cache_creation = getattr(usage, "cache_creation_input_tokens", 0) or 0
    if cache_read or cache_creation:
        metrics["cached_tokens"] = cache_read + cache_creation

    usage_entries: list[dict[str, Any]] = [
        {
            "type": "tokens",
            "metrics": metrics,
        }
    ]

    # Detect web search tool usage
    web_search_usage = _detect_web_search(response, tools)
    if web_search_usage:
        usage_entries.append(web_search_usage)

    transaction: dict[str, Any] = {
        "provider": "anthropic",
        "model": model,
        "usage": usage_entries,
        "context": context,
    }

    # Include raw usage for debugging if available
    try:
        transaction["provider_usage_raw"] = {
            "input_tokens": getattr(usage, "input_tokens", None),
            "output_tokens": getattr(usage, "output_tokens", None),
            "cache_read_input_tokens": getattr(usage, "cache_read_input_tokens", None),
            "cache_creation_input_tokens": getattr(usage, "cache_creation_input_tokens", None),
        }
    except Exception:
        pass

    enqueue_transaction(transaction)


def _track_stream(
    stream: Any,
    model: str,
    context: dict[str, Any],
    tools: list[Any] | None = None,
) -> Any:
    """
    Wrap a streaming response to track usage from the final chunk.

    For Anthropic's Stream[RawMessageStreamEvent], we need to accumulate
    usage from both message_start and message_delta events. The message_start
    contains initial input_tokens, while message_delta at the end contains
    the final output_tokens.
    """
    accumulated_message = None
    final_usage = None

    def wrapped_stream() -> Any:
        nonlocal accumulated_message, final_usage
        try:
            for event in stream:
                if hasattr(event, "type"):
                    if event.type == "message_start" and hasattr(event, "message"):
                        accumulated_message = event.message
                    elif event.type == "message_delta" and hasattr(event, "usage"):
                        # Capture usage from delta - this has final output_tokens
                        if event.usage:
                            final_usage = event.usage
                yield event
        finally:
            # After stream is exhausted, track usage
            if accumulated_message:
                # Merge final_usage into accumulated_message if available
                if final_usage and hasattr(accumulated_message, "usage") and accumulated_message.usage:
                    # Update output_tokens from delta
                    output_tokens = getattr(final_usage, "output_tokens", None)
                    if output_tokens is not None:
                        accumulated_message.usage.output_tokens = output_tokens
                try:
                    _track_message_response(accumulated_message, model, context, tools)
                except Exception as e:
                    logger.error(
                        f"Error tracking Anthropic streaming usage: {e}", exc_info=True
                    )

    return wrapped_stream()


async def _track_stream_async(
    stream: Any,
    model: str,
    context: dict[str, Any],
    tools: list[Any] | None = None,
) -> Any:
    """
    Wrap an async streaming response to track usage from the final chunk.

    For Anthropic's async stream, we need to accumulate usage from both
    message_start and message_delta events. The message_start contains
    initial input_tokens, while message_delta at the end contains
    the final output_tokens.
    """
    accumulated_message = None
    final_usage = None

    try:
        async for event in stream:
            if hasattr(event, "type"):
                if event.type == "message_start" and hasattr(event, "message"):
                    accumulated_message = event.message
                elif event.type == "message_delta" and hasattr(event, "usage"):
                    # Capture usage from delta - this has final output_tokens
                    if event.usage:
                        final_usage = event.usage
            yield event
    finally:
        # After stream is exhausted, track usage
        if accumulated_message:
            # Merge final_usage into accumulated_message if available
            if final_usage and hasattr(accumulated_message, "usage") and accumulated_message.usage:
                # Update output_tokens from delta
                output_tokens = getattr(final_usage, "output_tokens", None)
                if output_tokens is not None:
                    accumulated_message.usage.output_tokens = output_tokens
            try:
                _track_message_response(accumulated_message, model, context, tools)
            except Exception as e:
                logger.error(
                    f"Error tracking Anthropic async streaming usage: {e}", exc_info=True
                )


class _TrackedMessageStreamManager:
    """
    Wrapper for MessageStreamManager that tracks usage after stream completion.
    """

    def __init__(
        self,
        manager: Any,
        model: str,
        context: dict[str, Any],
        tools: list[Any] | None = None,
    ):
        self._manager = manager
        self._model = model
        self._context = context
        self._tools = tools

    def __enter__(self) -> Any:
        self._stream = self._manager.__enter__()
        return _TrackedMessageStream(
            self._stream, self._model, self._context, self._tools
        )

    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> Any:
        return self._manager.__exit__(exc_type, exc_val, exc_tb)


class _TrackedMessageStream:
    """
    Wrapper for MessageStream that tracks usage when get_final_message is called.
    """

    def __init__(
        self,
        stream: Any,
        model: str,
        context: dict[str, Any],
        tools: list[Any] | None = None,
    ):
        self._stream = stream
        self._model = model
        self._context = context
        self._tools = tools
        self._tracked = False

    def __iter__(self) -> Any:
        return iter(self._stream)

    def __getattr__(self, name: str) -> Any:
        attr = getattr(self._stream, name)

        if name == "get_final_message" and not self._tracked:

            @functools.wraps(attr)
            def tracked_get_final_message(*args: Any, **kwargs: Any) -> Any:
                result = attr(*args, **kwargs)
                if not self._tracked:
                    self._tracked = True
                    try:
                        _track_message_response(
                            result, self._model, self._context, self._tools
                        )
                    except Exception as e:
                        logger.error(
                            f"Error tracking Anthropic stream usage: {e}", exc_info=True
                        )
                return result

            return tracked_get_final_message

        return attr


class _TrackedAsyncMessageStreamManager:
    """
    Wrapper for AsyncMessageStreamManager that tracks usage after stream completion.
    """

    def __init__(
        self,
        manager: Any,
        model: str,
        context: dict[str, Any],
        tools: list[Any] | None = None,
    ):
        self._manager = manager
        self._model = model
        self._context = context
        self._tools = tools

    async def __aenter__(self) -> Any:
        self._stream = await self._manager.__aenter__()
        return _TrackedAsyncMessageStream(
            self._stream, self._model, self._context, self._tools
        )

    async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> Any:
        return await self._manager.__aexit__(exc_type, exc_val, exc_tb)


class _TrackedAsyncMessageStream:
    """
    Wrapper for AsyncMessageStream that tracks usage when get_final_message is called.
    """

    def __init__(
        self,
        stream: Any,
        model: str,
        context: dict[str, Any],
        tools: list[Any] | None = None,
    ):
        self._stream = stream
        self._model = model
        self._context = context
        self._tools = tools
        self._tracked = False

    def __aiter__(self) -> Any:
        return self._stream.__aiter__()

    def __getattr__(self, name: str) -> Any:
        attr = getattr(self._stream, name)

        if name == "get_final_message" and not self._tracked:

            @functools.wraps(attr)
            async def tracked_get_final_message(*args: Any, **kwargs: Any) -> Any:
                result = await attr(*args, **kwargs)
                if not self._tracked:
                    self._tracked = True
                    try:
                        _track_message_response(
                            result, self._model, self._context, self._tools
                        )
                    except Exception as e:
                        logger.error(
                            f"Error tracking Anthropic async stream usage: {e}",
                            exc_info=True,
                        )
                return result

            return tracked_get_final_message

        return attr


# ============================================================================
# Messages.create Patching (sync)
# ============================================================================


def patch_anthropic() -> None:
    """
    Patch Anthropic Messages client to auto-track usage.

    This patches the synchronous `Messages.create` method.
    """
    global _patched_sync, _original_create

    if _patched_sync:
        logger.debug("Anthropic Messages sync already patched, skipping")
        return

    try:
        from anthropic.resources.messages.messages import Messages
    except ImportError:
        logger.warning(
            "Anthropic SDK not installed. Install with: pip install 'fenra[anthropic]'"
        )
        return

    original_create = Messages.create
    _original_create = original_create

    @functools.wraps(original_create)
    def patched_create(self: Messages, *args: Any, **kwargs: Any) -> Any:
        is_streaming = kwargs.get("stream", False)

        response = original_create(self, *args, **kwargs)

        # Don't track streaming responses here - they need special handling
        if is_streaming:
            context = get_context()
            model = kwargs.get("model") or getattr(response, "model", "unknown")
            tools = kwargs.get("tools")
            try:
                return _track_stream(response, model, context, tools)
            except Exception as e:
                logger.error(f"Error wrapping Anthropic stream: {e}", exc_info=True)
                return response

        context = get_context()
        model = kwargs.get("model") or getattr(response, "model", "unknown")
        tools = kwargs.get("tools")

        try:
            _track_message_response(response, model, context, tools)
        except Exception as e:
            logger.error(f"Error tracking Anthropic usage: {e}", exc_info=True)

        return response

    Messages.create = patched_create  # type: ignore[assignment]
    _patched_sync = True
    logger.info("Anthropic Messages SDK patched for auto-instrumentation")


# ============================================================================
# AsyncMessages.create Patching (async)
# ============================================================================


def patch_anthropic_async() -> None:
    """
    Patch Anthropic Messages async client to auto-track usage.

    This patches the asynchronous `AsyncMessages.create` method.
    """
    global _patched_async, _original_create_async

    if _patched_async:
        logger.debug("Anthropic Messages async already patched, skipping")
        return

    try:
        from anthropic.resources.messages.messages import AsyncMessages
    except ImportError:
        logger.warning(
            "Anthropic SDK not installed. Install with: pip install 'fenra[anthropic]'"
        )
        return

    original_create = AsyncMessages.create
    _original_create_async = original_create

    @functools.wraps(original_create)
    async def patched_create(self: AsyncMessages, *args: Any, **kwargs: Any) -> Any:
        is_streaming = kwargs.get("stream", False)

        response = await original_create(self, *args, **kwargs)

        # Don't track streaming responses here - they need special handling
        if is_streaming:
            context = get_context()
            model = kwargs.get("model") or getattr(response, "model", "unknown")
            tools = kwargs.get("tools")
            try:
                return _track_stream_async(response, model, context, tools)
            except Exception as e:
                logger.error(
                    f"Error wrapping Anthropic async stream: {e}", exc_info=True
                )
                return response

        context = get_context()
        model = kwargs.get("model") or getattr(response, "model", "unknown")
        tools = kwargs.get("tools")

        try:
            _track_message_response(response, model, context, tools)
        except Exception as e:
            logger.error(f"Error tracking Anthropic async usage: {e}", exc_info=True)

        return response

    AsyncMessages.create = patched_create  # type: ignore[assignment]
    _patched_async = True
    logger.info("Anthropic Messages async SDK patched for auto-instrumentation")


# ============================================================================
# Messages.stream Patching (sync)
# ============================================================================


def patch_anthropic_stream() -> None:
    """
    Patch Anthropic Messages.stream to auto-track usage.

    This patches the synchronous `Messages.stream` method which returns
    a MessageStreamManager context manager.
    """
    global _patched_stream_sync, _original_stream

    if _patched_stream_sync:
        logger.debug("Anthropic Messages.stream sync already patched, skipping")
        return

    try:
        from anthropic.resources.messages.messages import Messages
    except ImportError:
        logger.warning(
            "Anthropic SDK not installed. Install with: pip install 'fenra[anthropic]'"
        )
        return

    if not hasattr(Messages, "stream"):
        logger.debug("Anthropic Messages has no stream method, skipping")
        return

    original_stream = Messages.stream
    _original_stream = original_stream

    @functools.wraps(original_stream)
    def patched_stream(self: Messages, *args: Any, **kwargs: Any) -> Any:
        manager = original_stream(self, *args, **kwargs)

        context = get_context()
        model = kwargs.get("model", "unknown")
        tools = kwargs.get("tools")

        try:
            return _TrackedMessageStreamManager(manager, model, context, tools)
        except Exception as e:
            logger.error(f"Error wrapping Anthropic stream manager: {e}", exc_info=True)
            return manager

    Messages.stream = patched_stream  # type: ignore[assignment]
    _patched_stream_sync = True
    logger.info("Anthropic Messages.stream SDK patched for auto-instrumentation")


# ============================================================================
# AsyncMessages.stream Patching (async)
# ============================================================================


def patch_anthropic_stream_async() -> None:
    """
    Patch Anthropic AsyncMessages.stream to auto-track usage.

    This patches the asynchronous `AsyncMessages.stream` method which returns
    an AsyncMessageStreamManager context manager.
    """
    global _patched_stream_async, _original_stream_async

    if _patched_stream_async:
        logger.debug("Anthropic AsyncMessages.stream async already patched, skipping")
        return

    try:
        from anthropic.resources.messages.messages import AsyncMessages
    except ImportError:
        logger.warning(
            "Anthropic SDK not installed. Install with: pip install 'fenra[anthropic]'"
        )
        return

    if not hasattr(AsyncMessages, "stream"):
        logger.debug("Anthropic AsyncMessages has no stream method, skipping")
        return

    original_stream = AsyncMessages.stream
    _original_stream_async = original_stream

    @functools.wraps(original_stream)
    def patched_stream(self: AsyncMessages, *args: Any, **kwargs: Any) -> Any:
        manager = original_stream(self, *args, **kwargs)

        context = get_context()
        model = kwargs.get("model", "unknown")
        tools = kwargs.get("tools")

        try:
            return _TrackedAsyncMessageStreamManager(manager, model, context, tools)
        except Exception as e:
            logger.error(
                f"Error wrapping Anthropic async stream manager: {e}", exc_info=True
            )
            return manager

    AsyncMessages.stream = patched_stream  # type: ignore[assignment]
    _patched_stream_async = True
    logger.info("Anthropic AsyncMessages.stream SDK patched for auto-instrumentation")


# ============================================================================
# Unpatch Functions
# ============================================================================


def unpatch_anthropic() -> None:
    """Restore original Anthropic Messages create method."""
    global _patched_sync, _original_create

    if not _patched_sync or _original_create is None:
        return

    try:
        from anthropic.resources.messages.messages import Messages

        Messages.create = _original_create  # type: ignore[assignment]
        _patched_sync = False
        _original_create = None
        logger.info("Anthropic Messages SDK unpatched")
    except ImportError:
        pass


def unpatch_anthropic_async() -> None:
    """Restore original Anthropic Messages async create method."""
    global _patched_async, _original_create_async

    if not _patched_async or _original_create_async is None:
        return

    try:
        from anthropic.resources.messages.messages import AsyncMessages

        AsyncMessages.create = _original_create_async  # type: ignore[assignment]
        _patched_async = False
        _original_create_async = None
        logger.info("Anthropic Messages async SDK unpatched")
    except ImportError:
        pass


def unpatch_anthropic_stream() -> None:
    """Restore original Anthropic Messages.stream method."""
    global _patched_stream_sync, _original_stream

    if not _patched_stream_sync or _original_stream is None:
        return

    try:
        from anthropic.resources.messages.messages import Messages

        Messages.stream = _original_stream  # type: ignore[assignment]
        _patched_stream_sync = False
        _original_stream = None
        logger.info("Anthropic Messages.stream SDK unpatched")
    except ImportError:
        pass


def unpatch_anthropic_stream_async() -> None:
    """Restore original Anthropic AsyncMessages.stream method."""
    global _patched_stream_async, _original_stream_async

    if not _patched_stream_async or _original_stream_async is None:
        return

    try:
        from anthropic.resources.messages.messages import AsyncMessages

        AsyncMessages.stream = _original_stream_async  # type: ignore[assignment]
        _patched_stream_async = False
        _original_stream_async = None
        logger.info("Anthropic AsyncMessages.stream SDK unpatched")
    except ImportError:
        pass


def unpatch_anthropic_all() -> None:
    """Restore all original Anthropic SDK methods."""
    unpatch_anthropic()
    unpatch_anthropic_async()
    unpatch_anthropic_stream()
    unpatch_anthropic_stream_async()
