import functools
import logging
from collections.abc import AsyncIterable

from ...sdk import VerseSDK


def patch_pydantic_ai(sdk: VerseSDK):
    """Patch Pydantic AI to use the Verse SDK."""

    from pydantic_ai import Agent, RunContext
    from pydantic_ai.messages import (
        AgentStreamEvent,
        FunctionToolResultEvent,
        PartDeltaEvent,
    )

    from .utils import (
        get_finish_reason,
        get_generation_messages,
        get_llm_info_from_model,
        get_operation_type,
        get_usage,
    )

    async def handle_events(ctx: RunContext, events: AsyncIterable[AgentStreamEvent]):
        """
        Handle Pydantic events to create LLM generation spans.
        Pydantic dispatches more events than we need, so the delta and finish reason checks are necessary to avoid duplicates.
        """
        try:
            contains_generation = get_finish_reason(ctx) == "stop"

            async for event in events:
                if isinstance(event, (PartDeltaEvent, FunctionToolResultEvent)):
                    contains_generation = True

            if not contains_generation:
                return

            op = get_operation_type(ctx)
            vendor, model_name = get_llm_info_from_model(ctx.model)
            prompts, completions = get_generation_messages(ctx)

            generation_ctx = sdk.create_generation(
                f"Pydantic AI: {op.capitalize()}",
                model=model_name,
                usage=get_usage(ctx.usage),
                vendor=vendor,
            )

            generation_ctx.provider(vendor).operation(op)

            if prompts:
                generation_ctx.messages(prompts)

            if completions:
                generation_ctx.completions(completions)

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

            generation_ctx.span.end()
        except Exception as e:
            logging.warning("Pydantic streaming event dropped", exc_info=e)

    def patch():
        """Setup monkey-patches for Pydantic AI if not already patched."""
        if not getattr(Agent, "_verse_patched", False):
            patch_run()
            patch_run_stream()
            Agent._verse_patched = True

    def patch_run():
        """Setup a monkey patch for Agent.run() that captures the LLM call in a generation span."""
        original_run_method = Agent.run

        @functools.wraps(original_run_method)
        async def run__patched(self, *args, **kwargs):
            try:
                return await original_run_method(
                    self, *args, **kwargs, event_stream_handler=handle_events
                )
            except Exception as e:
                logging.warning("Error patching Agent.run()", exc_info=e)
                return original_run_method(self, *args, **kwargs)

        Agent.run = run__patched

    def patch_run_stream():
        """Setup a monkey patch for Agent.run_stream() that captures streaming LLM calls."""
        original_run_stream_method = Agent.run_stream

        @functools.wraps(original_run_stream_method)
        def run_stream__patched(self, *args, **kwargs):
            try:
                stream = original_run_stream_method(self, *args, **kwargs)
                return _wrap_stream_context_manager(stream, self, sdk)
            except Exception as e:
                logging.warning("Error patching Agent.run_stream()", exc_info=e)
                return original_run_stream_method(self, *args, **kwargs)

        Agent.run_stream = run_stream__patched

    class TracedStreamContextManager:
        """Wrapper for run_stream() context manager that traces streaming events."""

        def __init__(self, stream, agent, sdk: VerseSDK):
            self._stream = stream
            self._agent = agent
            self._sdk = sdk
            self._inner_stream = None
            self._entered = False

        async def __aenter__(self):
            self._inner_stream = await self._stream.__aenter__()
            self._entered = True
            return self

        async def __aexit__(self, exc_type, exc_val, exc_tb):
            return await self._stream.__aexit__(exc_type, exc_val, exc_tb)

        def __aiter__(self):
            if not self._entered:
                raise RuntimeError(
                    "TracedStreamContextManager must be entered with 'async with' before iteration"
                )
            return _wrap_stream_iterator(self._inner_stream, self._agent, self._sdk)

        def stream_output(self):
            """Wrap stream_output() to trace events."""
            if not self._entered:
                raise RuntimeError(
                    "TracedStreamContextManager must be entered with 'async with' before calling stream_output()"
                )
            original_stream = self._inner_stream.stream_output()
            return _wrap_stream_iterator(original_stream, self._agent, self._sdk)

        def __getattr__(self, name):
            """Delegate attribute access to the underlying stream result."""
            if self._inner_stream is None:
                raise AttributeError(
                    f"'{type(self).__name__}' object has no attribute '{name}'. "
                    "Context manager must be entered first."
                )
            return getattr(self._inner_stream, name)

    def _wrap_stream_context_manager(stream, agent, sdk: VerseSDK):
        """Wrap the stream context manager to return a traced version."""
        return TracedStreamContextManager(stream, agent, sdk)

    async def _wrap_stream_iterator(stream_result, agent, sdk: VerseSDK):
        """Wrap the stream iterator to trace events as they come in."""
        generation_ctx = None
        first_event = True
        accumulated_content = []

        try:
            if hasattr(stream_result, "stream_output"):
                event_iter = stream_result.stream_output()
            elif hasattr(stream_result, "stream"):
                event_iter = stream_result.stream()
            elif hasattr(stream_result, "__aiter__"):
                event_iter = stream_result
            else:
                raise TypeError(f"Cannot iterate over {type(stream_result)}")

            async for event in event_iter:
                if first_event:
                    try:
                        op = "chat"
                        vendor, model_name = get_llm_info_from_model(agent.model)

                        generation_ctx = sdk.create_generation(
                            f"Pydantic AI: {op.capitalize()}",
                            model=model_name,
                            vendor=vendor,
                        )

                        generation_ctx.provider(vendor).operation(op)
                        first_event = False
                    except Exception as e:
                        logging.warning(
                            "Error initializing generation context for stream",
                            exc_info=e,
                        )

                if isinstance(event, PartDeltaEvent) and generation_ctx:
                    try:
                        if hasattr(event, "delta"):
                            delta = event.delta
                            if hasattr(delta, "content") and delta.content:
                                content_str = str(delta.content)
                                accumulated_content.append(content_str)
                                sdk.event(
                                    "gen_ai.stream.delta",
                                    metadata={"content": content_str[:512]},
                                    level="info",
                                    span=generation_ctx.span,
                                )
                    except Exception as e:
                        logging.warning("Error processing delta event", exc_info=e)

                yield event

            if generation_ctx:
                try:
                    if accumulated_content:
                        from ...contexts import CompletionChoice

                        full_content = "".join(accumulated_content)
                        generation_ctx.completions(
                            [
                                CompletionChoice(
                                    message={
                                        "role": "assistant",
                                        "content": full_content,
                                        "function_call": None,
                                        "name": None,
                                        "tool_calls": None,
                                    },
                                    delta=None,
                                    text=full_content,
                                )
                            ]
                        )
                except Exception as e:
                    logging.warning("Error finalizing stream span", exc_info=e)
                finally:
                    generation_ctx.span.end()

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

    patch()
