"""@trace() decorator — capture function inputs/outputs as Carrot traces."""

from __future__ import annotations

import asyncio
import functools
import logging
import time
import uuid
from typing import Any, Callable, TypeVar, overload

from carrot_ai._context import current_trace_id
from carrot_ai._wrap import _format_time, _safe_serialize

logger = logging.getLogger("carrot_ai")

F = TypeVar("F", bound=Callable[..., Any])


# ---------------------------------------------------------------------------
# Public API — supports @trace, @trace(), and @trace("name")
# ---------------------------------------------------------------------------


@overload
def trace(fn: F) -> F: ...


@overload
def trace(
    name: str | None = None,
    *,
    tags: list[str] | None = None,
    metadata: dict[str, Any] | None = None,
) -> Callable[[F], F]: ...


def trace(
    fn: F | str | None = None,
    *,
    name: str | None = None,
    tags: list[str] | None = None,
    metadata: dict[str, Any] | None = None,
) -> F | Callable[[F], F]:
    """Decorator that wraps a function as a Carrot trace.

    Usage::

        @trace
        def answer(question: str) -> str:
            ...

        @trace("my-pipeline")
        def answer(question: str) -> str:
            ...

        @trace(tags=["prod"], metadata={"version": "2"})
        def answer(question: str) -> str:
            ...

    LLM calls made through a ``wrap()``-ed client inside the decorated
    function are automatically linked via ``metadata.parent_trace_id``.
    """
    # @trace  (bare decorator, fn is the function)
    if callable(fn):
        return _wrap_function(fn, name=None, tags=tags, metadata=metadata)  # type: ignore[return-value]

    # @trace("name")  (fn is actually the name string)
    if isinstance(fn, str):
        trace_name = fn

        def decorator(func: F) -> F:
            return _wrap_function(func, name=trace_name, tags=tags, metadata=metadata)  # type: ignore[return-value]

        return decorator  # type: ignore[return-value]

    # @trace() or @trace(tags=[...])
    def decorator(func: F) -> F:
        return _wrap_function(func, name=name, tags=tags, metadata=metadata)  # type: ignore[return-value]

    return decorator  # type: ignore[return-value]


# ---------------------------------------------------------------------------
# Internals
# ---------------------------------------------------------------------------


def _wrap_function(
    fn: Callable[..., Any],
    *,
    name: str | None,
    tags: list[str] | None,
    metadata: dict[str, Any] | None,
) -> Any:
    trace_name = name or fn.__qualname__

    if asyncio.iscoroutinefunction(fn):

        @functools.wraps(fn)
        async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
            return await _run_async(fn, trace_name, tags, metadata, args, kwargs)

        return async_wrapper

    @functools.wraps(fn)
    def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
        return _run_sync(fn, trace_name, tags, metadata, args, kwargs)

    return sync_wrapper


def _run_sync(
    fn: Callable[..., Any],
    trace_name: str,
    tags: list[str] | None,
    extra_metadata: dict[str, Any] | None,
    args: tuple[Any, ...],
    kwargs: dict[str, Any],
) -> Any:
    trace_id = str(uuid.uuid4())
    token = current_trace_id.set(trace_id)
    started_at = time.time()

    trace_input = _build_input(fn, args, kwargs)

    try:
        result = fn(*args, **kwargs)
        _submit_trace(
            trace_id, trace_name, tags, extra_metadata,
            trace_input, _serialize_output(result),
            started_at, time.time(), "success",
        )
        return result
    except Exception as exc:
        _submit_trace(
            trace_id, trace_name, tags, extra_metadata,
            trace_input, {"error": str(exc), "error_type": type(exc).__name__},
            started_at, time.time(), "error",
        )
        raise
    finally:
        current_trace_id.reset(token)


async def _run_async(
    fn: Callable[..., Any],
    trace_name: str,
    tags: list[str] | None,
    extra_metadata: dict[str, Any] | None,
    args: tuple[Any, ...],
    kwargs: dict[str, Any],
) -> Any:
    trace_id = str(uuid.uuid4())
    token = current_trace_id.set(trace_id)
    started_at = time.time()

    trace_input = _build_input(fn, args, kwargs)

    try:
        result = await fn(*args, **kwargs)
        _submit_trace(
            trace_id, trace_name, tags, extra_metadata,
            trace_input, _serialize_output(result),
            started_at, time.time(), "success",
        )
        return result
    except Exception as exc:
        _submit_trace(
            trace_id, trace_name, tags, extra_metadata,
            trace_input, {"error": str(exc), "error_type": type(exc).__name__},
            started_at, time.time(), "error",
        )
        raise
    finally:
        current_trace_id.reset(token)


def _build_input(
    fn: Callable[..., Any],
    args: tuple[Any, ...],
    kwargs: dict[str, Any],
) -> dict[str, Any]:
    """Capture the function's positional and keyword args as a dict."""
    import inspect

    sig = inspect.signature(fn)
    params = list(sig.parameters.keys())

    input_data: dict[str, Any] = {}
    for i, val in enumerate(args):
        key = params[i] if i < len(params) else f"arg_{i}"
        input_data[key] = _safe_serialize(val)
    for key, val in kwargs.items():
        input_data[key] = _safe_serialize(val)

    return input_data


def _serialize_output(result: Any) -> dict[str, Any]:
    serialized = _safe_serialize(result)
    if isinstance(serialized, dict):
        return serialized
    return {"result": serialized}


def _submit_trace(
    trace_id: str,
    trace_name: str,
    tags: list[str] | None,
    extra_metadata: dict[str, Any] | None,
    trace_input: dict[str, Any],
    trace_output: dict[str, Any],
    started_at: float,
    ended_at: float,
    status: str,
) -> None:
    from carrot_ai import get_client

    meta: dict[str, Any] = {
        "source": "decorator",
        "latency_ms": round((ended_at - started_at) * 1000),
    }
    if extra_metadata:
        meta.update(extra_metadata)

    trace: dict[str, Any] = {
        "id": trace_id,
        "name": trace_name,
        "input": trace_input,
        "output": trace_output,
        "metadata": meta,
        "status": status,
        "started_at": _format_time(started_at),
        "ended_at": _format_time(ended_at),
    }
    if tags:
        trace["tags"] = list(tags)

    try:
        get_client().submit(trace)
    except Exception:
        logger.debug("carrot_ai: failed to submit decorator trace", exc_info=True)
