"""Decorators for instrumenting agent code with trace collection.

This module provides decorators that wrap agent functions, tool calls,
and LLM interactions to automatically collect execution traces.

Example:
    from agent_debugger_sdk import trace_agent, trace_tool, trace_llm

    @trace_agent(name="my_agent", framework="pydantic_ai")
    async def my_agent(prompt: str) -> str:
        result = await call_llm([{"role": "user", "content": prompt}])
        return result

    @trace_tool(name="search_web")
    async def search_web(query: str) -> list[str]:
        return ["result1", "result2"]

    @trace_llm(model="gpt-4o")
    async def call_llm(messages: list) -> str:
        return "Hello!"
"""

from __future__ import annotations

import contextlib
import time
import traceback
from collections.abc import Awaitable
from collections.abc import Callable
from functools import wraps
from typing import Any
from typing import ParamSpec
from typing import TypeVar

from agent_debugger_sdk.core.context import TraceContext
from agent_debugger_sdk.core.context import get_current_context
from agent_debugger_sdk.core.events import ErrorEvent
from agent_debugger_sdk.core.events import EventType
from agent_debugger_sdk.core.events import LLMRequestEvent
from agent_debugger_sdk.core.events import LLMResponseEvent
from agent_debugger_sdk.core.events import ToolCallEvent
from agent_debugger_sdk.core.events import ToolResultEvent
from agent_debugger_sdk.core.events import TraceEvent

P = ParamSpec("P")
T = TypeVar("T")

__all__ = ["trace_agent", "trace_tool", "trace_llm"]


def trace_agent(
    name: str,
    framework: str = "unknown",
) -> Callable[[Callable[P, Awaitable[T]]], Callable[P, Awaitable[T]]]:
    """Decorator to trace an async agent function.

    Creates a new trace session on entry, records AGENT_START and AGENT_END
    events, and captures any exceptions as ERROR events.

    Args:
        name: Human-readable name for the agent.
        framework: The agent framework being used (pydantic_ai, langchain, autogen).

    Returns:
        A decorator function that wraps async agent functions.

    Example:
        @trace_agent(name="research_agent", framework="pydantic_ai")
        async def research_agent(query: str) -> str:
            results = await search_tool(query)
            return summarize(results)

    Note:
        The decorated function must be async. For sync functions, wrap them
        in an async wrapper or use asyncio.run() at the call site.
    """

    def decorator(func: Callable[P, Awaitable[T]]) -> Callable[P, Awaitable[T]]:
        @wraps(func)
        async def async_wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
            ctx = TraceContext(
                agent_name=name,
                framework=framework,
            )

            async with ctx:
                # Use the session start event's ID as the parent for child events
                if ctx._session_start_event:
                    ctx.set_parent(ctx._session_start_event.id)

                try:
                    result = await func(*args, **kwargs)
                    return result
                except Exception as e:
                    # Error is already recorded by TraceContext.__aexit__
                    raise

        return async_wrapper

    return decorator


def trace_tool(
    name: str,
) -> Callable[[Callable[P, Awaitable[T]]], Callable[P, Awaitable[T]]]:
    """Decorator to trace a tool function.

    Records TOOL_CALL before execution and TOOL_RESULT after execution,
    including duration and any errors.

    Args:
        name: Human-readable name for the tool.

    Returns:
        A decorator function that wraps async tool functions.

    Example:
        @trace_tool(name="web_search")
        async def web_search(query: str) -> list[str]:
            results = await search_api(query)
            return results

    Note:
        Works both inside a trace_agent context (uses existing session)
        or standalone (creates a temporary context).
    """

    def decorator(func: Callable[P, Awaitable[T]]) -> Callable[P, Awaitable[T]]:
        @wraps(func)
        async def async_wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
            ctx = get_current_context()
            own_context = ctx is None

            if own_context:
                ctx = TraceContext(
                    agent_name="tool_runner",
                    framework="unknown",
                )
                await ctx.__aenter__()

            if ctx is None:
                raise RuntimeError("TraceContext is None - this should not happen")

            tool_call_event = ToolCallEvent(
                session_id=ctx.session_id,
                parent_id=ctx.get_current_parent(),
                event_type=EventType.TOOL_CALL,
                name=f"{name}_call",
                tool_name=name,
                arguments=_sanitize_arguments(args, kwargs),
                importance=0.4,
            )
            await ctx._emit_event(tool_call_event)

            start_time = time.perf_counter()
            error: Exception | None = None
            result: T | None = None

            try:
                result = await func(*args, **kwargs)
                return result
            except Exception as e:
                error = e
                raise
            finally:
                end_time = time.perf_counter()
                duration_ms = (end_time - start_time) * 1000

                tool_result_event = ToolResultEvent(
                    session_id=ctx.session_id,
                    parent_id=tool_call_event.id,
                    event_type=EventType.TOOL_RESULT,
                    name=f"{name}_result",
                    tool_name=name,
                    result=_sanitize_result(result) if result else None,
                    error=str(error) if error else None,
                    duration_ms=duration_ms,
                    importance=0.9 if error else 0.5,
                )
                await ctx._emit_event(tool_result_event)

                if own_context:
                    # Pass actual exception info to __aexit__ when error occurred
                    if error is not None:
                        await ctx.__aexit__(type(error), error, error.__traceback__)
                    else:
                        await ctx.__aexit__(None, None, None)

        return async_wrapper

    return decorator


def trace_llm(
    model: str,
) -> Callable[[Callable[P, Awaitable[T]]], Callable[P, Awaitable[T]]]:
    """Decorator to trace an LLM call function.

    Records LLM_REQUEST before the call and LLM_RESPONSE after, including
    token usage, cost, and duration.

    Args:
        model: The model identifier (e.g., "gpt-4o", "claude-3-opus").

    Returns:
        A decorator function that wraps async LLM call functions.

    Example:
        @trace_llm(model="gpt-4o")
        async def call_gpt(messages: list[dict]) -> str:
            response = await openai.chat.completions.create(
                model="gpt-4o",
                messages=messages
            )
            return response.choices[0].message.content

    Note:
        The decorated function should return either:
        - A string (content only)
        - A dict with 'content', 'usage', 'cost_usd' keys
        - An object with these as attributes

        Works both inside a trace_agent context or standalone.
    """

    def decorator(func: Callable[P, Awaitable[T]]) -> Callable[P, Awaitable[T]]:
        @wraps(func)
        async def async_wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
            ctx = get_current_context()
            own_context = ctx is None

            if own_context:
                ctx = TraceContext(
                    agent_name="llm_runner",
                    framework="unknown",
                )
                await ctx.__aenter__()

            if ctx is None:
                raise RuntimeError("TraceContext is None - this should not happen")

            messages = _extract_messages(args, kwargs)

            llm_request_event = LLMRequestEvent(
                session_id=ctx.session_id,
                parent_id=ctx.get_current_parent(),
                event_type=EventType.LLM_REQUEST,
                name=f"llm_call_{model}",
                model=model,
                messages=messages,
                tools=_extract_tools(args, kwargs),
                settings=_extract_settings(args, kwargs),
                importance=0.3,
            )
            await ctx._emit_event(llm_request_event)

            start_time = time.perf_counter()
            error: Exception | None = None
            result: T | None = None

            try:
                result = await func(*args, **kwargs)
                return result
            except Exception as e:
                error = e
                raise
            finally:
                end_time = time.perf_counter()
                duration_ms = (end_time - start_time) * 1000

                content = ""
                usage: dict[str, int] = {"input_tokens": 0, "output_tokens": 0}
                cost_usd = 0.0
                tool_calls: list[dict[str, Any]] = []

                if result is not None and error is None:
                    content, usage, cost_usd, tool_calls = _extract_llm_response(result)

                llm_response_event = LLMResponseEvent(
                    session_id=ctx.session_id,
                    parent_id=llm_request_event.id,
                    event_type=EventType.LLM_RESPONSE,
                    name=f"llm_response_{model}",
                    model=model,
                    content=content,
                    tool_calls=tool_calls,
                    usage=usage,
                    cost_usd=cost_usd,
                    duration_ms=duration_ms,
                    importance=0.9 if error else 0.5,
                )
                await ctx._emit_event(llm_response_event)

                if own_context:
                    # Pass actual exception info to __aexit__ when error occurred
                    if error is not None:
                        await ctx.__aexit__(type(error), error, error.__traceback__)
                    else:
                        await ctx.__aexit__(None, None, None)

        return async_wrapper

    return decorator


def _sanitize_arguments(args: tuple, kwargs: dict) -> dict[str, Any]:
    """Sanitize function arguments for trace storage.

    Converts positional and keyword arguments to a dictionary,
    truncating large values.

    Args:
        args: Positional arguments.
        kwargs: Keyword arguments.

    Returns:
        A dictionary of sanitized arguments.
    """
    sanitized: dict[str, Any] = {}

    for i, arg in enumerate(args):
        key = f"arg_{i}"
        sanitized[key] = _truncate_value(arg)

    for key, value in kwargs.items():
        sanitized[key] = _truncate_value(value)

    return sanitized


def _truncate_value(value: Any, max_length: int = 1000) -> Any:
    """Truncate a value if it's too large for trace storage.

    Args:
        value: The value to potentially truncate.
        max_length: Maximum string length before truncation.

    Returns:
        The value, possibly truncated.
    """
    if isinstance(value, str):
        if len(value) > max_length:
            return value[:max_length] + "...[truncated]"
        return value

    if isinstance(value, list | tuple):
        if len(value) > 100:
            return [_truncate_value(v, max_length) for v in value[:10]] + [f"...[{len(value) - 10} more items]"]
        return [_truncate_value(v, max_length) for v in value]

    if isinstance(value, dict):
        if len(value) > 50:
            truncated = {}
            for i, (k, v) in enumerate(value.items()):
                if i >= 20:
                    truncated["__truncated__"] = f"{len(value) - 20} more keys"
                    break
                truncated[k] = _truncate_value(v, max_length)
            return truncated
        return {k: _truncate_value(v, max_length) for k, v in value.items()}

    return value


def _sanitize_result(result: Any) -> Any:
    """Sanitize a function result for trace storage.

    Args:
        result: The result to sanitize.

    Returns:
        A sanitized version of the result.
    """
    return _truncate_value(result, max_length=5000)


def _extract_messages(args: tuple, kwargs: dict) -> list[dict[str, Any]]:
    """Extract messages from LLM call arguments.

    Looks for 'messages' in kwargs or first list argument.

    Args:
        args: Positional arguments.
        kwargs: Keyword arguments.

    Returns:
        A list of message dictionaries.
    """
    if "messages" in kwargs:
        messages = kwargs["messages"]
        if isinstance(messages, list):
            return _truncate_value(messages)
        return [{"role": "unknown", "content": str(messages)}]

    for arg in args:
        if (
            isinstance(arg, list)
            and len(arg) > 0
            and isinstance(arg[0], dict)
            and ("role" in arg[0] or "content" in arg[0])
        ):
            return _truncate_value(arg)

    return []


def _extract_tools(args: tuple, kwargs: dict) -> list[dict[str, Any]]:
    """Extract tool definitions from LLM call arguments.

    Args:
        args: Positional arguments.
        kwargs: Keyword arguments.

    Returns:
        A list of tool definition dictionaries.
    """
    if "tools" in kwargs:
        return _truncate_value(kwargs["tools"])
    return []


def _extract_settings(args: tuple, kwargs: dict) -> dict[str, Any]:
    """Extract model settings from LLM call arguments.

    Args:
        args: Positional arguments.
        kwargs: Keyword arguments.

    Returns:
        A dictionary of settings.
    """
    settings: dict[str, Any] = {}

    for key in ["temperature", "max_tokens", "top_p", "frequency_penalty", "presence_penalty"]:
        if key in kwargs:
            settings[key] = kwargs[key]

    return settings


def _extract_llm_response(result: Any) -> tuple[str, dict[str, int], float, list[dict[str, Any]]]:
    """Extract content, usage, cost, and tool calls from an LLM response.

    Handles various response formats (dict, object with attributes).

    Args:
        result: The LLM response to extract from.

    Returns:
        A tuple of (content, usage, cost_usd, tool_calls).
    """
    content = ""
    usage: dict[str, int] = {"input_tokens": 0, "output_tokens": 0}
    cost_usd = 0.0
    tool_calls: list[dict[str, Any]] = []

    if isinstance(result, str):
        content = result
    elif isinstance(result, dict):
        content = result.get("content", "")
        if "usage" in result:
            usage = result["usage"]
        if "cost_usd" in result:
            cost_usd = result["cost_usd"]
        if "tool_calls" in result:
            tool_calls = result["tool_calls"]
    else:
        if hasattr(result, "content"):
            content = str(result.content)
        elif hasattr(result, "choices"):
            try:
                content = result.choices[0].message.content
            except (AttributeError, IndexError, KeyError):
                content = str(result)

        if hasattr(result, "usage"):
            with contextlib.suppress(AttributeError):
                usage = {
                    "input_tokens": getattr(result.usage, "prompt_tokens", 0)
                    or getattr(result.usage, "input_tokens", 0),
                    "output_tokens": getattr(result.usage, "completion_tokens", 0)
                    or getattr(result.usage, "output_tokens", 0),
                }

        if hasattr(result, "tool_calls"):
            with contextlib.suppress(AttributeError):
                for tc in result.tool_calls or []:
                    tool_calls.append(
                        {
                            "id": getattr(tc, "id", ""),
                            "name": getattr(tc.function, "name", "")
                            if hasattr(tc, "function")
                            else getattr(tc, "name", ""),
                            "arguments": getattr(tc.function, "arguments", "")
                            if hasattr(tc, "function")
                            else getattr(tc, "arguments", {}),
                        }
                    )

    content = _truncate_value(content, max_length=5000) if isinstance(content, str) else str(content)

    return content, usage, cost_usd, tool_calls
