"""Decorator API for plyra-trace."""

import functools
import inspect
import json
from collections.abc import Callable
from typing import Any

from opentelemetry import trace as trace_api
from opentelemetry.trace import Status, StatusCode

from plyra_trace import _tracer
from plyra_trace.models import GuardResult
from plyra_trace.semconv import SpanAttributes, SpanKind


def _serialize_value(value: Any) -> str:
    """Serialize a value to JSON string. Handle special cases."""
    try:
        # Handle GuardResult
        if isinstance(value, GuardResult):
            return value.model_dump_json()
        # Handle Pydantic models
        if hasattr(value, "model_dump_json"):
            return value.model_dump_json()
        # Handle standard types
        return json.dumps(value, default=str)
    except (TypeError, ValueError):
        return str(value)


def _set_span_attributes(
    span: trace_api.Span,
    kind: SpanKind,
    input_value: Any = None,
    output_value: Any = None,
    capture_input: bool = True,
    capture_output: bool = True,
    extra_attributes: dict[str, Any] | None = None,
) -> None:
    """Set common span attributes."""
    # Set span kind (both plyra and openinference)
    span.set_attribute(SpanAttributes.SPAN_KIND, kind.value)
    span.set_attribute(SpanAttributes.OPENINFERENCE_SPAN_KIND, kind.value)

    # Set input
    if capture_input and input_value is not None:
        span.set_attribute(SpanAttributes.INPUT_VALUE, _serialize_value(input_value))
        span.set_attribute(SpanAttributes.INPUT_MIME_TYPE, "application/json")

    # Set output
    if capture_output and output_value is not None:
        span.set_attribute(SpanAttributes.OUTPUT_VALUE, _serialize_value(output_value))
        span.set_attribute(SpanAttributes.OUTPUT_MIME_TYPE, "application/json")

    # Set extra attributes
    if extra_attributes:
        for key, value in extra_attributes.items():
            if value is not None:
                if isinstance(value, (dict, list)):
                    span.set_attribute(key, json.dumps(value))
                else:
                    span.set_attribute(key, value)


def trace(
    name: str | None = None,
    capture_input: bool = True,
    capture_output: bool = True,
) -> Callable:
    """
    Root span decorator. Creates a CHAIN span.
    Use on the top-level entry point of an agent run.

    Args:
        name: Span name (defaults to function name)
        capture_input: Capture function arguments as input.value
        capture_output: Capture return value as output.value

    Example:
        >>> @trace()
        ... def run_agent(query: str) -> str:
        ...     ...

        >>> @trace(name="my-pipeline", capture_output=False)
        ... async def process(data: dict) -> dict:
        ...     ...
    """

    def decorator(func: Callable) -> Callable:
        if inspect.iscoroutinefunction(func):

            @functools.wraps(func)
            async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
                tracer = _tracer.get_tracer()
                span_name = name or func.__name__

                with tracer.start_as_current_span(span_name) as span:
                    input_data = {"args": args, "kwargs": kwargs} if (args or kwargs) else None
                    _set_span_attributes(
                        span,
                        SpanKind.CHAIN,
                        input_value=input_data,
                        capture_input=capture_input,
                        capture_output=False,
                    )

                    try:
                        result = await func(*args, **kwargs)

                        if capture_output:
                            span.set_attribute(SpanAttributes.OUTPUT_VALUE, _serialize_value(result))
                            span.set_attribute(SpanAttributes.OUTPUT_MIME_TYPE, "application/json")

                        span.set_status(Status(StatusCode.OK))
                        return result

                    except Exception as e:
                        span.set_status(Status(StatusCode.ERROR, str(e)))
                        span.record_exception(e)
                        raise

            return async_wrapper
        else:

            @functools.wraps(func)
            def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
                tracer = _tracer.get_tracer()
                span_name = name or func.__name__

                with tracer.start_as_current_span(span_name) as span:
                    input_data = {"args": args, "kwargs": kwargs} if (args or kwargs) else None
                    _set_span_attributes(
                        span,
                        SpanKind.CHAIN,
                        input_value=input_data,
                        capture_input=capture_input,
                        capture_output=False,
                    )

                    try:
                        result = func(*args, **kwargs)

                        if capture_output:
                            span.set_attribute(SpanAttributes.OUTPUT_VALUE, _serialize_value(result))
                            span.set_attribute(SpanAttributes.OUTPUT_MIME_TYPE, "application/json")

                        span.set_status(Status(StatusCode.OK))
                        return result

                    except Exception as e:
                        span.set_status(Status(StatusCode.ERROR, str(e)))
                        span.record_exception(e)
                        raise

            return sync_wrapper

    return decorator


def agent(
    name: str | None = None,
    description: str | None = None,
    framework: str | None = None,
    pattern: str | None = None,
    capture_input: bool = True,
    capture_output: bool = True,
) -> Callable:
    """
    AGENT span decorator. Use on autonomous reasoning functions.

    Args:
        name: Agent name
        description: Agent description
        framework: Framework name (e.g., "langgraph", "crewai", "autogen")
        pattern: Agent pattern (e.g., "ReAct", "Chain-of-Thought")
        capture_input: Capture function arguments as input.value
        capture_output: Capture return value as output.value

    Example:
        >>> @agent(name="planner", framework="langgraph")
        ... async def planner_agent(state: dict) -> dict:
        ...     ...
    """

    def decorator(func: Callable) -> Callable:
        if inspect.iscoroutinefunction(func):

            @functools.wraps(func)
            async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
                tracer = _tracer.get_tracer()
                span_name = name or func.__name__

                with tracer.start_as_current_span(span_name) as span:
                    extra_attrs = {}
                    if name:
                        extra_attrs[SpanAttributes.AGENT_NAME] = name
                    if description:
                        extra_attrs[SpanAttributes.AGENT_DESCRIPTION] = description
                    if framework:
                        extra_attrs[SpanAttributes.AGENT_FRAMEWORK] = framework
                    if pattern:
                        extra_attrs[SpanAttributes.AGENT_PATTERN] = pattern

                    input_data = {"args": args, "kwargs": kwargs} if (args or kwargs) else None
                    _set_span_attributes(
                        span,
                        SpanKind.AGENT,
                        input_value=input_data,
                        capture_input=capture_input,
                        capture_output=False,
                        extra_attributes=extra_attrs,
                    )

                    try:
                        result = await func(*args, **kwargs)

                        if capture_output:
                            span.set_attribute(SpanAttributes.OUTPUT_VALUE, _serialize_value(result))
                            span.set_attribute(SpanAttributes.OUTPUT_MIME_TYPE, "application/json")

                        span.set_status(Status(StatusCode.OK))
                        return result

                    except Exception as e:
                        span.set_status(Status(StatusCode.ERROR, str(e)))
                        span.record_exception(e)
                        raise

            return async_wrapper
        else:

            @functools.wraps(func)
            def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
                tracer = _tracer.get_tracer()
                span_name = name or func.__name__

                with tracer.start_as_current_span(span_name) as span:
                    extra_attrs = {}
                    if name:
                        extra_attrs[SpanAttributes.AGENT_NAME] = name
                    if description:
                        extra_attrs[SpanAttributes.AGENT_DESCRIPTION] = description
                    if framework:
                        extra_attrs[SpanAttributes.AGENT_FRAMEWORK] = framework
                    if pattern:
                        extra_attrs[SpanAttributes.AGENT_PATTERN] = pattern

                    input_data = {"args": args, "kwargs": kwargs} if (args or kwargs) else None
                    _set_span_attributes(
                        span,
                        SpanKind.AGENT,
                        input_value=input_data,
                        capture_input=capture_input,
                        capture_output=False,
                        extra_attributes=extra_attrs,
                    )

                    try:
                        result = func(*args, **kwargs)

                        if capture_output:
                            span.set_attribute(SpanAttributes.OUTPUT_VALUE, _serialize_value(result))
                            span.set_attribute(SpanAttributes.OUTPUT_MIME_TYPE, "application/json")

                        span.set_status(Status(StatusCode.OK))
                        return result

                    except Exception as e:
                        span.set_status(Status(StatusCode.ERROR, str(e)))
                        span.record_exception(e)
                        raise

            return sync_wrapper

    return decorator


def tool(
    name: str | None = None,
    description: str | None = None,
    parameters: dict | None = None,
    capture_input: bool = True,
    capture_output: bool = True,
) -> Callable:
    """
    TOOL span decorator. Use on external function/API calls.

    Args:
        name: Tool name
        description: Tool description
        parameters: Tool parameters JSON schema
        capture_input: Capture function arguments as input.value
        capture_output: Capture return value as output.value

    Example:
        >>> @tool(name="web-search", description="Search the web")
        ... def web_search(query: str) -> list[dict]:
        ...     ...
    """

    def decorator(func: Callable) -> Callable:
        if inspect.iscoroutinefunction(func):

            @functools.wraps(func)
            async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
                tracer = _tracer.get_tracer()
                span_name = name or func.__name__

                with tracer.start_as_current_span(span_name) as span:
                    extra_attrs = {}
                    if name:
                        extra_attrs[SpanAttributes.TOOL_NAME] = name
                    if description:
                        extra_attrs[SpanAttributes.TOOL_DESCRIPTION] = description
                    if parameters:
                        extra_attrs[SpanAttributes.TOOL_PARAMETERS] = parameters

                    input_data = {"args": args, "kwargs": kwargs} if (args or kwargs) else None
                    _set_span_attributes(
                        span,
                        SpanKind.TOOL,
                        input_value=input_data,
                        capture_input=capture_input,
                        capture_output=False,
                        extra_attributes=extra_attrs,
                    )

                    try:
                        result = await func(*args, **kwargs)

                        if capture_output:
                            span.set_attribute(SpanAttributes.OUTPUT_VALUE, _serialize_value(result))
                            span.set_attribute(SpanAttributes.OUTPUT_MIME_TYPE, "application/json")

                        span.set_status(Status(StatusCode.OK))
                        return result

                    except Exception as e:
                        span.set_status(Status(StatusCode.ERROR, str(e)))
                        span.record_exception(e)
                        raise

            return async_wrapper
        else:

            @functools.wraps(func)
            def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
                tracer = _tracer.get_tracer()
                span_name = name or func.__name__

                with tracer.start_as_current_span(span_name) as span:
                    extra_attrs = {}
                    if name:
                        extra_attrs[SpanAttributes.TOOL_NAME] = name
                    if description:
                        extra_attrs[SpanAttributes.TOOL_DESCRIPTION] = description
                    if parameters:
                        extra_attrs[SpanAttributes.TOOL_PARAMETERS] = parameters

                    input_data = {"args": args, "kwargs": kwargs} if (args or kwargs) else None
                    _set_span_attributes(
                        span,
                        SpanKind.TOOL,
                        input_value=input_data,
                        capture_input=capture_input,
                        capture_output=False,
                        extra_attributes=extra_attrs,
                    )

                    try:
                        result = func(*args, **kwargs)

                        if capture_output:
                            span.set_attribute(SpanAttributes.OUTPUT_VALUE, _serialize_value(result))
                            span.set_attribute(SpanAttributes.OUTPUT_MIME_TYPE, "application/json")

                        span.set_status(Status(StatusCode.OK))
                        return result

                    except Exception as e:
                        span.set_status(Status(StatusCode.ERROR, str(e)))
                        span.record_exception(e)
                        raise

            return sync_wrapper

    return decorator


def guard(
    name: str | None = None,
    policy: str | None = None,
    provider: str = "plyra-guard",
    capture_input: bool = True,
    capture_output: bool = True,
) -> Callable:
    """
    GUARD span decorator. Use on policy/safety check functions.
    The decorated function should return a GuardResult.

    Args:
        name: Guard name
        policy: Policy name
        provider: Guard provider (e.g., "plyra-guard", "guardrails-ai")
        capture_input: Capture function arguments as input.value
        capture_output: Capture return value as output.value

    Example:
        >>> @guard(name="input-toxicity", policy="content-safety")
        ... async def check_input(text: str) -> GuardResult:
        ...     ...
    """

    def decorator(func: Callable) -> Callable:
        if inspect.iscoroutinefunction(func):

            @functools.wraps(func)
            async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
                tracer = _tracer.get_tracer()
                span_name = name or func.__name__

                with tracer.start_as_current_span(span_name) as span:
                    extra_attrs = {SpanAttributes.GUARD_PROVIDER: provider}
                    if policy:
                        extra_attrs[SpanAttributes.GUARD_POLICY] = policy

                    input_data = {"args": args, "kwargs": kwargs} if (args or kwargs) else None
                    _set_span_attributes(
                        span,
                        SpanKind.GUARD,
                        input_value=input_data,
                        capture_input=capture_input,
                        capture_output=False,
                        extra_attributes=extra_attrs,
                    )

                    try:
                        result = await func(*args, **kwargs)

                        if isinstance(result, GuardResult):
                            span.set_attribute(SpanAttributes.GUARD_ACTION, result.action)
                            span.set_attribute(SpanAttributes.GUARD_TRIGGERED, result.triggered)
                            if result.confidence is not None:
                                span.set_attribute(SpanAttributes.GUARD_CONFIDENCE, result.confidence)
                            if result.details is not None:
                                span.set_attribute(
                                    SpanAttributes.GUARD_DETAILS,
                                    json.dumps(result.details),
                                )

                        if capture_output:
                            span.set_attribute(SpanAttributes.OUTPUT_VALUE, _serialize_value(result))
                            span.set_attribute(SpanAttributes.OUTPUT_MIME_TYPE, "application/json")

                        span.set_status(Status(StatusCode.OK))
                        return result

                    except Exception as e:
                        span.set_status(Status(StatusCode.ERROR, str(e)))
                        span.record_exception(e)
                        raise

            return async_wrapper
        else:

            @functools.wraps(func)
            def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
                tracer = _tracer.get_tracer()
                span_name = name or func.__name__

                with tracer.start_as_current_span(span_name) as span:
                    extra_attrs = {SpanAttributes.GUARD_PROVIDER: provider}
                    if policy:
                        extra_attrs[SpanAttributes.GUARD_POLICY] = policy

                    input_data = {"args": args, "kwargs": kwargs} if (args or kwargs) else None
                    _set_span_attributes(
                        span,
                        SpanKind.GUARD,
                        input_value=input_data,
                        capture_input=capture_input,
                        capture_output=False,
                        extra_attributes=extra_attrs,
                    )

                    try:
                        result = func(*args, **kwargs)

                        if isinstance(result, GuardResult):
                            span.set_attribute(SpanAttributes.GUARD_ACTION, result.action)
                            span.set_attribute(SpanAttributes.GUARD_TRIGGERED, result.triggered)
                            if result.confidence is not None:
                                span.set_attribute(SpanAttributes.GUARD_CONFIDENCE, result.confidence)
                            if result.details is not None:
                                span.set_attribute(
                                    SpanAttributes.GUARD_DETAILS,
                                    json.dumps(result.details),
                                )

                        if capture_output:
                            span.set_attribute(SpanAttributes.OUTPUT_VALUE, _serialize_value(result))
                            span.set_attribute(SpanAttributes.OUTPUT_MIME_TYPE, "application/json")

                        span.set_status(Status(StatusCode.OK))
                        return result

                    except Exception as e:
                        span.set_status(Status(StatusCode.ERROR, str(e)))
                        span.record_exception(e)
                        raise

            return sync_wrapper

    return decorator


def llm(
    model: str | None = None,
    provider: str | None = None,
    capture_input: bool = True,
    capture_output: bool = True,
) -> Callable:
    """
    LLM span decorator. Rarely needed with auto-instrumentation,
    but useful for wrapping custom model calls.

    Args:
        model: Model name
        provider: Provider name (e.g., "openai", "anthropic")
        capture_input: Capture function arguments as input.value
        capture_output: Capture return value as output.value

    Example:
        >>> @llm(model="my-fine-tune", provider="custom")
        ... def call_model(prompt: str) -> str:
        ...     ...
    """

    def decorator(func: Callable) -> Callable:
        if inspect.iscoroutinefunction(func):

            @functools.wraps(func)
            async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
                tracer = _tracer.get_tracer()
                span_name = f"llm.{model or func.__name__}"

                with tracer.start_as_current_span(span_name) as span:
                    extra_attrs = {}
                    if model:
                        extra_attrs[SpanAttributes.LLM_MODEL_NAME] = model
                    if provider:
                        extra_attrs[SpanAttributes.LLM_PROVIDER] = provider
                        extra_attrs[SpanAttributes.LLM_SYSTEM] = provider

                    input_data = {"args": args, "kwargs": kwargs} if (args or kwargs) else None
                    _set_span_attributes(
                        span,
                        SpanKind.LLM,
                        input_value=input_data,
                        capture_input=capture_input,
                        capture_output=False,
                        extra_attributes=extra_attrs,
                    )

                    try:
                        result = await func(*args, **kwargs)

                        if capture_output:
                            span.set_attribute(SpanAttributes.OUTPUT_VALUE, _serialize_value(result))
                            span.set_attribute(SpanAttributes.OUTPUT_MIME_TYPE, "application/json")

                        span.set_status(Status(StatusCode.OK))
                        return result

                    except Exception as e:
                        span.set_status(Status(StatusCode.ERROR, str(e)))
                        span.record_exception(e)
                        raise

            return async_wrapper
        else:

            @functools.wraps(func)
            def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
                tracer = _tracer.get_tracer()
                span_name = f"llm.{model or func.__name__}"

                with tracer.start_as_current_span(span_name) as span:
                    extra_attrs = {}
                    if model:
                        extra_attrs[SpanAttributes.LLM_MODEL_NAME] = model
                    if provider:
                        extra_attrs[SpanAttributes.LLM_PROVIDER] = provider
                        extra_attrs[SpanAttributes.LLM_SYSTEM] = provider

                    input_data = {"args": args, "kwargs": kwargs} if (args or kwargs) else None
                    _set_span_attributes(
                        span,
                        SpanKind.LLM,
                        input_value=input_data,
                        capture_input=capture_input,
                        capture_output=False,
                        extra_attributes=extra_attrs,
                    )

                    try:
                        result = func(*args, **kwargs)

                        if capture_output:
                            span.set_attribute(SpanAttributes.OUTPUT_VALUE, _serialize_value(result))
                            span.set_attribute(SpanAttributes.OUTPUT_MIME_TYPE, "application/json")

                        span.set_status(Status(StatusCode.OK))
                        return result

                    except Exception as e:
                        span.set_status(Status(StatusCode.ERROR, str(e)))
                        span.record_exception(e)
                        raise

            return sync_wrapper

    return decorator
