from __future__ import annotations

import functools
import inspect
import json
import logging
from collections.abc import Callable
from typing import Any, TypeVar

from .client import get_tracer
from .types import KEY_AGENT_NAME, KEY_RESULT, META_TAGS, META_TOOL_NAME, SpanKind, SpanStatus

logger = logging.getLogger("staso")
F = TypeVar("F", bound=Callable[..., Any])


def _safe_json(obj: Any) -> str:
    try:
        return json.dumps(obj, default=str)
    except (TypeError, ValueError, OverflowError):
        return "{}"


def _capture_args(fn: Callable, args: tuple, kwargs: dict) -> dict:
    """Best-effort capture of function arguments as a serialisable dict."""
    try:
        sig = inspect.signature(fn)
        params = list(sig.parameters.keys())
        captured: dict[str, Any] = {}
        for i, val in enumerate(args):
            key = params[i] if i < len(params) else f"arg{i}"
            captured[key] = val
        captured.update(kwargs)
        return captured
    except Exception:
        return {}


def agent(
    name: str | None = None,
    *,
    tags: list[str] | None = None,
    metadata: dict | None = None,
    capture_input: bool = True,
    capture_output: bool = True,
) -> Callable[[F], F]:
    """Decorator for agent entry-points."""
    _tags = list(tags or [])
    _metadata = dict(metadata) if metadata else None
    _capture_input = capture_input
    _capture_output = capture_output

    def decorator(fn: F) -> F:
        _name = name or fn.__name__

        if inspect.iscoroutinefunction(fn):

            @functools.wraps(fn)
            async def _async(*args: Any, **kwargs: Any) -> Any:
                tracer = get_tracer()
                captured = _capture_args(fn, args, kwargs)
                with tracer.start_as_current_span(_name, kind=SpanKind.AGENT) as span:
                    if _capture_input:
                        span.input = {KEY_AGENT_NAME: _name, **captured}
                    if _tags:
                        span.metadata[META_TAGS] = _tags
                    if _metadata:
                        span.metadata.update(_metadata)
                    try:
                        result = await fn(*args, **kwargs)
                        if _capture_output:
                            span.output = {KEY_RESULT: result}
                        span.status = SpanStatus.OK
                        return result
                    except Exception as exc:
                        span.record_exception(exc)
                        raise

            return _async  # type: ignore[return-value]

        else:

            @functools.wraps(fn)
            def _sync(*args: Any, **kwargs: Any) -> Any:
                tracer = get_tracer()
                captured = _capture_args(fn, args, kwargs)
                with tracer.start_as_current_span(_name, kind=SpanKind.AGENT) as span:
                    if _capture_input:
                        span.input = {KEY_AGENT_NAME: _name, **captured}
                    if _tags:
                        span.metadata[META_TAGS] = _tags
                    if _metadata:
                        span.metadata.update(_metadata)
                    try:
                        result = fn(*args, **kwargs)
                        if _capture_output:
                            span.output = {KEY_RESULT: result}
                        span.status = SpanStatus.OK
                        return result
                    except Exception as exc:
                        span.record_exception(exc)
                        raise

            return _sync  # type: ignore[return-value]

    return decorator


def tool(
    name: str | None = None,
    *,
    tags: list[str] | None = None,
    metadata: dict | None = None,
    capture_input: bool = True,
    capture_output: bool = True,
) -> Callable[[F], F]:
    """Decorator for tool functions."""
    _tags = list(tags or [])
    _metadata = dict(metadata) if metadata else None
    _capture_input = capture_input
    _capture_output = capture_output

    def decorator(fn: F) -> F:
        _name = name or fn.__name__

        if inspect.iscoroutinefunction(fn):

            @functools.wraps(fn)
            async def _async(*args: Any, **kwargs: Any) -> Any:
                tracer = get_tracer()
                captured = _capture_args(fn, args, kwargs)
                with tracer.start_as_current_span(_name, kind=SpanKind.TOOL) as span:
                    span.metadata[META_TOOL_NAME] = _name
                    if _capture_input:
                        span.input = captured
                    if _tags:
                        span.metadata[META_TAGS] = _tags
                    if _metadata:
                        span.metadata.update(_metadata)
                    try:
                        result = await fn(*args, **kwargs)
                        if _capture_output:
                            span.output = {KEY_RESULT: result}
                        span.status = SpanStatus.OK
                        return result
                    except Exception as exc:
                        span.record_exception(exc)
                        raise

            return _async  # type: ignore[return-value]

        else:

            @functools.wraps(fn)
            def _sync(*args: Any, **kwargs: Any) -> Any:
                tracer = get_tracer()
                captured = _capture_args(fn, args, kwargs)
                with tracer.start_as_current_span(_name, kind=SpanKind.TOOL) as span:
                    span.metadata[META_TOOL_NAME] = _name
                    if _capture_input:
                        span.input = captured
                    if _tags:
                        span.metadata[META_TAGS] = _tags
                    if _metadata:
                        span.metadata.update(_metadata)
                    try:
                        result = fn(*args, **kwargs)
                        if _capture_output:
                            span.output = {KEY_RESULT: result}
                        span.status = SpanStatus.OK
                        return result
                    except Exception as exc:
                        span.record_exception(exc)
                        raise

            return _sync  # type: ignore[return-value]

    return decorator


def trace_fn(
    name: str | None = None,
    *,
    kind: str | SpanKind = SpanKind.CHAIN,
    tags: list[str] | None = None,
    metadata: dict | None = None,
    capture_input: bool = True,
    capture_output: bool = True,
) -> Callable[[F], F]:
    """Generic span decorator for arbitrary functions."""
    _tags = list(tags or [])
    _metadata = dict(metadata) if metadata else None
    _capture_input = capture_input
    _capture_output = capture_output
    # Normalize uppercase strings like "CHAIN", "RETRIEVER" for backward compat
    _kind = SpanKind(kind.lower()) if isinstance(kind, str) else kind

    def decorator(fn: F) -> F:
        _name = name or fn.__name__

        if inspect.iscoroutinefunction(fn):

            @functools.wraps(fn)
            async def _async(*args: Any, **kwargs: Any) -> Any:
                tracer = get_tracer()
                captured = _capture_args(fn, args, kwargs)
                with tracer.start_as_current_span(_name, kind=_kind) as span:
                    if _capture_input:
                        span.input = captured
                    if _tags:
                        span.metadata[META_TAGS] = _tags
                    if _metadata:
                        span.metadata.update(_metadata)
                    try:
                        result = await fn(*args, **kwargs)
                        if _capture_output:
                            span.output = {KEY_RESULT: result}
                        span.status = SpanStatus.OK
                        return result
                    except Exception as exc:
                        span.record_exception(exc)
                        raise

            return _async  # type: ignore[return-value]

        else:

            @functools.wraps(fn)
            def _sync(*args: Any, **kwargs: Any) -> Any:
                tracer = get_tracer()
                captured = _capture_args(fn, args, kwargs)
                with tracer.start_as_current_span(_name, kind=_kind) as span:
                    if _capture_input:
                        span.input = captured
                    if _tags:
                        span.metadata[META_TAGS] = _tags
                    if _metadata:
                        span.metadata.update(_metadata)
                    try:
                        result = fn(*args, **kwargs)
                        if _capture_output:
                            span.output = {KEY_RESULT: result}
                        span.status = SpanStatus.OK
                        return result
                    except Exception as exc:
                        span.record_exception(exc)
                        raise

            return _sync  # type: ignore[return-value]

    return decorator
