"""
Pydantic AI patches for Agent.run(), Agent.run_sync(), Agent.run_stream().

Wraps agent execution methods to create AGENT spans. Does NOT suppress
provider instrumentation — like CrewAI, we want both the agent-level span
(orchestration view) and the LLM call spans (individual model calls).

result.output is used for response data, result.usage() for token counts.
"""

from __future__ import annotations

import logging
import time
from typing import Any, Callable, Dict, TypeVar

import wrapt

from risicare.integrations._base import (
    create_framework_attributes,
    get_framework_version,
    get_tracer,
    is_tracing_enabled,
    record_error,
    safe_set_attribute,
    should_trace_content,
    truncate_content,
)

logger = logging.getLogger(__name__)

T = TypeVar("T")


def _extract_agent_info(instance: Any) -> Dict[str, Any]:
    """Extract agent info from Pydantic AI Agent instance (defensive)."""
    info: Dict[str, Any] = {}

    # Model name
    model = getattr(instance, "model", None)
    if model is not None:
        info["model"] = str(model)

    # Result type
    result_type = getattr(instance, "result_type", None)
    if result_type is not None:
        info["result_type"] = getattr(result_type, "__name__", str(result_type))

    # Tools
    tools = getattr(instance, "_tools", None)
    if tools is None:
        tools = getattr(instance, "tools", None)
    if tools:
        try:
            if isinstance(tools, dict):
                info["tools"] = list(tools.keys())[:20]
            elif isinstance(tools, (list, tuple)):
                info["tools"] = [
                    getattr(t, "name", getattr(t, "__name__", str(t)))
                    for t in tools[:20]
                ]
        except Exception:
            pass

    # System prompt length
    system_prompt = getattr(instance, "system_prompt", None)
    if system_prompt is not None:
        try:
            if callable(system_prompt):
                pass  # Can't measure dynamic prompts
            elif isinstance(system_prompt, str):
                info["system_prompt_length"] = len(system_prompt)
            elif isinstance(system_prompt, (list, tuple)):
                info["system_prompt_length"] = sum(
                    len(str(p)) for p in system_prompt
                )
        except Exception:
            pass

    # Agent name
    info["name"] = getattr(instance, "name", None) or getattr(instance, "__class__", type(instance)).__name__

    return info


def _extract_usage_from_result(result: Any) -> Dict[str, Any]:
    """Extract token usage from a Pydantic AI RunResult (defensive)."""
    usage_info: Dict[str, Any] = {}
    try:
        usage_fn = getattr(result, "usage", None)
        if usage_fn is None:
            return usage_info
        usage = usage_fn() if callable(usage_fn) else usage_fn
        if usage is not None:
            input_tokens = getattr(usage, "input_tokens", None)
            output_tokens = getattr(usage, "output_tokens", None)
            requests = getattr(usage, "requests", None)
            if input_tokens is not None:
                usage_info["prompt_tokens"] = input_tokens
            if output_tokens is not None:
                usage_info["completion_tokens"] = output_tokens
            if requests is not None:
                usage_info["requests"] = requests
    except Exception:
        pass
    return usage_info


def _create_agent_span_attrs(instance: Any, version: str | None) -> Dict[str, Any]:
    """Create standard attributes for an agent span."""
    agent_info = _extract_agent_info(instance)
    attrs = create_framework_attributes("pydantic_ai", version)
    attrs["framework.span_kind"] = "agent"

    if agent_info.get("model"):
        attrs["framework.pydantic_ai.model"] = agent_info["model"]
    if agent_info.get("result_type"):
        attrs["framework.pydantic_ai.result_type"] = agent_info["result_type"]
    if agent_info.get("tools"):
        attrs["framework.pydantic_ai.tools"] = agent_info["tools"]
    if agent_info.get("system_prompt_length"):
        attrs["framework.pydantic_ai.system_prompt_length"] = agent_info["system_prompt_length"]

    attrs["agent.type"] = "pydantic_ai"
    attrs["agent.name"] = agent_info.get("name", "Agent")

    return attrs, agent_info


def _wrap_run_sync(
    wrapped: Callable[..., T],
    instance: Any,
    args: tuple,
    kwargs: Dict[str, Any],
) -> T:
    """Wrapper for Agent.run_sync()."""
    if not is_tracing_enabled():
        return wrapped(*args, **kwargs)

    tracer = get_tracer()
    if tracer is None:
        return wrapped(*args, **kwargs)

    from risicare_core import SpanKind

    version = get_framework_version("pydantic-ai")
    attrs, agent_info = _create_agent_span_attrs(instance, version)
    attrs["framework.operation"] = "run_sync"
    agent_name = agent_info.get("name", "Agent")

    with tracer.start_span(
        name=f"pydantic_ai.agent.run/{agent_name}",
        kind=SpanKind.AGENT,
        attributes=attrs,
    ) as span:
        try:
            start_time = time.perf_counter()
            result = wrapped(*args, **kwargs)
            latency_ms = (time.perf_counter() - start_time) * 1000
            span.set_attribute("gen_ai.latency_ms", latency_ms)

            # Extract usage
            usage = _extract_usage_from_result(result)
            for key, value in usage.items():
                if key == "prompt_tokens":
                    safe_set_attribute(span, "gen_ai.usage.prompt_tokens", value)
                elif key == "completion_tokens":
                    safe_set_attribute(span, "gen_ai.usage.completion_tokens", value)

            # Capture output
            if should_trace_content():
                output = getattr(result, "output", None)
                if output is not None:
                    safe_set_attribute(
                        span,
                        "gen_ai.response.content",
                        truncate_content(str(output), 5000),
                    )

            return result
        except Exception as e:
            record_error(span, e)
            raise


async def _wrap_run(
    wrapped: Callable[..., T],
    instance: Any,
    args: tuple,
    kwargs: Dict[str, Any],
) -> T:
    """Wrapper for Agent.run() (async)."""
    if not is_tracing_enabled():
        return await wrapped(*args, **kwargs)

    tracer = get_tracer()
    if tracer is None:
        return await wrapped(*args, **kwargs)

    from risicare_core import SpanKind

    version = get_framework_version("pydantic-ai")
    attrs, agent_info = _create_agent_span_attrs(instance, version)
    attrs["framework.operation"] = "run"
    agent_name = agent_info.get("name", "Agent")

    with tracer.start_span(
        name=f"pydantic_ai.agent.run/{agent_name}",
        kind=SpanKind.AGENT,
        attributes=attrs,
    ) as span:
        try:
            start_time = time.perf_counter()
            result = await wrapped(*args, **kwargs)
            latency_ms = (time.perf_counter() - start_time) * 1000
            span.set_attribute("gen_ai.latency_ms", latency_ms)

            usage = _extract_usage_from_result(result)
            for key, value in usage.items():
                if key == "prompt_tokens":
                    safe_set_attribute(span, "gen_ai.usage.prompt_tokens", value)
                elif key == "completion_tokens":
                    safe_set_attribute(span, "gen_ai.usage.completion_tokens", value)

            if should_trace_content():
                output = getattr(result, "output", None)
                if output is not None:
                    safe_set_attribute(
                        span,
                        "gen_ai.response.content",
                        truncate_content(str(output), 5000),
                    )

            return result
        except Exception as e:
            record_error(span, e)
            raise


async def _wrap_run_stream(
    wrapped: Callable[..., T],
    instance: Any,
    args: tuple,
    kwargs: Dict[str, Any],
) -> T:
    """Wrapper for Agent.run_stream() (async context manager)."""
    if not is_tracing_enabled():
        return await wrapped(*args, **kwargs)

    tracer = get_tracer()
    if tracer is None:
        return await wrapped(*args, **kwargs)

    from risicare_core import SpanKind

    version = get_framework_version("pydantic-ai")
    attrs, agent_info = _create_agent_span_attrs(instance, version)
    attrs["framework.operation"] = "run_stream"
    agent_name = agent_info.get("name", "Agent")

    # For streaming, we start a span and return the stream result.
    # Uses try/finally to guarantee span_cm.__exit__ is always called,
    # even if record_error() or set_attribute() raises unexpectedly.
    span_cm = tracer.start_span(
        name=f"pydantic_ai.agent.run_stream/{agent_name}",
        kind=SpanKind.AGENT,
        attributes=attrs,
    )
    span = span_cm.__enter__()
    exc_info: tuple | None = None

    try:
        start_time = time.perf_counter()
        result = await wrapped(*args, **kwargs)
        latency_ms = (time.perf_counter() - start_time) * 1000
        span.set_attribute("gen_ai.latency_ms", latency_ms)
        return result
    except Exception as e:
        exc_info = (type(e), e, e.__traceback__)
        try:
            record_error(span, e)
        except Exception:
            pass
        raise
    finally:
        try:
            if exc_info is not None:
                span_cm.__exit__(*exc_info)
            else:
                span_cm.__exit__(None, None, None)
        except Exception:
            pass


def patch_pydantic_ai(module: Any) -> None:
    """Apply wrapt patches to Pydantic AI Agent class."""
    # Patch Agent.run_sync
    try:
        wrapt.wrap_function_wrapper(
            "pydantic_ai", "Agent.run_sync", _wrap_run_sync
        )
    except Exception as e:
        logger.debug(f"Failed to patch Agent.run_sync: {e}")

    # Patch Agent.run (async)
    try:
        wrapt.wrap_function_wrapper(
            "pydantic_ai", "Agent.run", _wrap_run
        )
    except Exception as e:
        logger.debug(f"Failed to patch Agent.run: {e}")

    # Patch Agent.run_stream (async)
    try:
        wrapt.wrap_function_wrapper(
            "pydantic_ai", "Agent.run_stream", _wrap_run_stream
        )
    except Exception as e:
        logger.debug(f"Failed to patch Agent.run_stream: {e}")
