import functools
import logging
from typing import Any

from opentelemetry.trace import Status, StatusCode

from ...contexts import GenerationContext
from ...sdk import VerseSDK
from .utils import (
    format_messages_for_tracing,
    get_completion_choices,
    get_finish_reason,
    get_request_parameters,
    get_stream_text,
    get_usage_from_response,
)


def configure_anthropic(sdk: VerseSDK):
    """Configure Anthropic to use the Verse SDK."""
    try:
        import anthropic
    except ImportError:
        logging.warning("Anthropic SDK not installed, skipping integration")
        return

    original_create = anthropic.resources.messages.Messages.create
    original_create_async = anthropic.resources.messages.AsyncMessages.create
    original_stream = anthropic.resources.messages.Messages.stream

    @functools.wraps(original_create)
    def create_with_tracing(self, *args, **kwargs):
        """Wrap synchronous messages.create with tracing."""
        model = kwargs.get("model", "unknown")
        messages = kwargs.get("messages", [])
        stream = kwargs.get("stream", False)

        generation_ctx = sdk.create_generation(
            name=f"Anthropic: {'Chat' if messages else 'Completion'}",
            model=model,
            usage=get_request_parameters(kwargs),
        )

        generation_ctx.provider("anthropic").operation(
            "chat" if messages else "text_completion"
        ).request_params(get_request_parameters(kwargs))

        if messages:
            formatted_messages = format_messages_for_tracing(messages)
            generation_ctx.messages(formatted_messages)

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

            if stream:
                return _wrap_stream_response(response, generation_ctx, sdk)

            _process_response(response, generation_ctx)
            generation_ctx.span.end()
            return response

        except Exception as e:
            generation_ctx.error(e)
            generation_ctx.span.end()
            raise

    @functools.wraps(original_create_async)
    async def create_async_with_tracing(self, *args, **kwargs):
        """Wrap asynchronous messages.create with tracing."""
        model = kwargs.get("model", "unknown")
        messages = kwargs.get("messages", [])
        stream = kwargs.get("stream", False)

        generation_ctx = sdk.create_generation(
            name=f"Anthropic: {'Chat' if messages else 'Completion'}",
            model=model,
            usage=get_request_parameters(kwargs),
        )

        generation_ctx.provider("anthropic").operation(
            "chat" if messages else "text_completion"
        ).request_params(get_request_parameters(kwargs))

        if messages:
            formatted_messages = format_messages_for_tracing(messages)
            generation_ctx.messages(formatted_messages)

        try:
            response = await original_create_async(self, *args, **kwargs)

            if stream:
                return _wrap_async_stream_response(response, generation_ctx, sdk)

            _process_response(response, generation_ctx)
            generation_ctx.span.end()
            return response

        except Exception as e:
            generation_ctx.error(e)
            generation_ctx.span.end()
            raise

    @functools.wraps(original_stream)
    def stream_with_tracing(self, *args, **kwargs):
        """Wrap synchronous messages.stream with tracing."""
        model = kwargs.get("model", "unknown")
        messages = kwargs.get("messages", [])

        generation_ctx = sdk.create_generation(
            name="Anthropic: Chat",
            model=model,
            usage=get_request_parameters(kwargs),
        )

        generation_ctx.provider("anthropic").operation("chat").request_params(
            get_request_parameters(kwargs)
        )

        if messages:
            formatted_messages = format_messages_for_tracing(messages)
            generation_ctx.messages(formatted_messages)

        try:
            stream = original_stream(self, *args, **kwargs)
            return _wrap_stream_context_manager(stream, generation_ctx, sdk)
        except Exception as e:
            generation_ctx.error(e)
            generation_ctx.span.end()
            raise

    class TracedStreamContextManager:
        """Wrapper for MessageStream context manager that traces events."""

        def __init__(self, stream_manager, generation_ctx: GenerationContext, sdk: VerseSDK):
            self._stream_manager = stream_manager
            self._generation_ctx = generation_ctx
            self._sdk = sdk
            self._finalized = False
            self._inner_stream = None

        def __enter__(self):
            self._inner_stream = self._stream_manager.__enter__()
            return self

        def __exit__(self, exc_type, exc_val, exc_tb):
            try:
                return self._stream_manager.__exit__(exc_type, exc_val, exc_tb)
            finally:
                self._finalize()

        def __iter__(self):
            """Iterate over stream events, emit deltas, and finalize span."""
            if self._inner_stream is None:
                raise RuntimeError(
                    "TracedStreamContextManager must be entered with 'with' before iteration"
                )
            try:
                for event in self._inner_stream:
                    if (
                        event.type == "content_block_delta"
                        and hasattr(event.delta, "text")
                        and event.delta.text is not None
                    ):
                        self._sdk.event(
                            "gen_ai.stream.delta",
                            metadata={"content": event.delta.text[:512]},
                            level="info",
                            span=self._generation_ctx.span,
                        )
                    yield event
            except Exception as e:
                self._generation_ctx.error(e)
                raise
            finally:
                self._finalize()

        def get_final_message(self):
            """Get the final message from the stream."""
            if self._inner_stream is None:
                raise RuntimeError(
                    "TracedStreamContextManager must be entered with 'with' before calling get_final_message()"
                )
            return self._inner_stream.get_final_message()

        def _finalize(self):
            """Finalize the span using get_final_message if available."""
            if self._finalized:
                return
            self._finalized = True

            try:
                if self._inner_stream:
                    final_message = self._inner_stream.get_final_message()
                    if final_message:
                        _process_response(final_message, self._generation_ctx)
            except Exception as e:
                logging.warning("Error finalizing Anthropic stream", exc_info=e)
            finally:
                self._generation_ctx.span.end()

    def _wrap_stream_context_manager(stream, generation_ctx: GenerationContext, sdk: VerseSDK):
        """Wrap streaming context manager with tracing proxy."""
        return TracedStreamContextManager(stream, generation_ctx, sdk)

    anthropic.resources.messages.Messages.create = create_with_tracing
    anthropic.resources.messages.AsyncMessages.create = create_async_with_tracing
    anthropic.resources.messages.Messages.stream = stream_with_tracing


def _process_response(response: Any, generation_ctx: GenerationContext):
    """Process non-streaming response and add attributes to span."""
    try:
        if hasattr(response, "model"):
            generation_ctx.model_used(response.model)

        if hasattr(response, "id"):
            generation_ctx.response_id(response.id)

        finish_reason = get_finish_reason(response)
        if finish_reason:
            generation_ctx.reason(finish_reason)

        choices = get_completion_choices(response)
        if choices:
            generation_ctx.completions(choices)

        usage = get_usage_from_response(response)
        if usage:
            generation_ctx.usage(usage)

        generation_ctx.status(Status(StatusCode.OK))

    except Exception as e:
        logging.warning("Error processing Anthropic response", exc_info=e)


class TracedMessageStream:
    """Proxy wrapper for MessageStream that ensures span finalization."""

    def __init__(self, stream, generation_ctx: GenerationContext, sdk: VerseSDK):
        self._stream = stream
        self._generation_ctx = generation_ctx
        self._sdk = sdk
        self._finalized = False

    def __getattr__(self, name):
        """Delegate attribute access to the underlying stream."""
        return getattr(self._stream, name)

    def __iter__(self):
        """Iterate over stream chunks, emit deltas, and finalize span."""
        try:
            for chunk in self._stream:
                text = get_stream_text(chunk)
                if text:
                    self._sdk.event(
                        "gen_ai.stream.delta",
                        metadata={"content": text[:512]},
                        level="info",
                        span=self._generation_ctx.span,
                    )
                yield chunk
        except Exception as e:
            self._generation_ctx.error(e)
            raise
        finally:
            self._finalize()

    def __enter__(self):
        """Enter context manager."""
        self._stream.__enter__()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """Exit context manager and finalize span."""
        try:
            return self._stream.__exit__(exc_type, exc_val, exc_tb)
        finally:
            self._finalize()

    def _finalize(self):
        """Finalize the span using get_final_message if available."""
        if self._finalized:
            return
        self._finalized = True

        try:
            if hasattr(self._stream, "get_final_message"):
                final_message = self._stream.get_final_message()
                _process_response(final_message, self._generation_ctx)
            elif hasattr(self._stream, "response"):
                _process_response(self._stream.response, self._generation_ctx)
        except Exception as e:
            logging.warning("Error finalizing Anthropic stream", exc_info=e)
        finally:
            self._generation_ctx.span.end()


class TracedAsyncMessageStream:
    """Proxy wrapper for AsyncMessageStream that ensures span finalization."""

    def __init__(self, stream, generation_ctx: GenerationContext, sdk: VerseSDK):
        self._stream = stream
        self._generation_ctx = generation_ctx
        self._sdk = sdk
        self._finalized = False

    def __getattr__(self, name):
        """Delegate attribute access to the underlying stream."""
        return getattr(self._stream, name)

    async def __aiter__(self):
        """Iterate over stream chunks, emit deltas, and finalize span."""
        try:
            async for chunk in self._stream:
                text = get_stream_text(chunk)
                if text:
                    self._sdk.event(
                        "gen_ai.stream.delta",
                        metadata={"content": text[:512]},
                        level="info",
                        span=self._generation_ctx.span,
                    )
                yield chunk
        except Exception as e:
            self._generation_ctx.error(e)
            raise
        finally:
            await self._finalize()

    async def __aenter__(self):
        """Enter async context manager."""
        await self._stream.__aenter__()
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        """Exit async context manager and finalize span."""
        try:
            return await self._stream.__aexit__(exc_type, exc_val, exc_tb)
        finally:
            await self._finalize()

    async def _finalize(self):
        """Finalize the span using get_final_message if available."""
        if self._finalized:
            return
        self._finalized = True

        try:
            if hasattr(self._stream, "get_final_message"):
                final_message = await self._stream.get_final_message()
                _process_response(final_message, self._generation_ctx)
            elif hasattr(self._stream, "response"):
                _process_response(self._stream.response, self._generation_ctx)
        except Exception as e:
            logging.warning("Error finalizing Anthropic async stream", exc_info=e)
        finally:
            self._generation_ctx.span.end()


def _wrap_stream_response(stream, generation_ctx: GenerationContext, sdk: VerseSDK):
    """Wrap streaming response with tracing proxy."""
    return TracedMessageStream(stream, generation_ctx, sdk)


def _wrap_async_stream_response(
    stream, generation_ctx: GenerationContext, sdk: VerseSDK
):
    """Wrap async streaming response with tracing proxy."""
    return TracedAsyncMessageStream(stream, generation_ctx, sdk)
