from __future__ import annotations

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

from openinference.semconv.trace import OpenInferenceSpanKindValues
from openinference.semconv.trace import SpanAttributes as OI
from opentelemetry import trace
from opentelemetry.trace import StatusCode

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 _get_tracer() -> trace.Tracer:
    return trace.get_tracer("staso")


def agent(
    name: str | None = None,
    *,
    session_id: str = "",
    tags: list[str] | None = None,
) -> Callable[[F], F]:
    """Decorator for agent entry-points."""
    _tags = list(tags or [])

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

        if asyncio.iscoroutinefunction(fn):

            @functools.wraps(fn)
            async def _async(*args: Any, **kwargs: Any) -> Any:
                tracer = _get_tracer()
                with tracer.start_as_current_span(
                    _name,
                    attributes={
                        OI.OPENINFERENCE_SPAN_KIND: OpenInferenceSpanKindValues.AGENT.value,
                        OI.INPUT_VALUE: _safe_json({"agent_name": _name}),
                    },
                ) as span:
                    if _tags:
                        span.set_attribute(OI.TAG_TAGS, _tags)
                    try:
                        result = await fn(*args, **kwargs)
                        span.set_status(StatusCode.OK)
                        return result
                    except Exception as exc:
                        span.set_status(StatusCode.ERROR, str(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()
                with tracer.start_as_current_span(
                    _name,
                    attributes={
                        OI.OPENINFERENCE_SPAN_KIND: OpenInferenceSpanKindValues.AGENT.value,
                        OI.INPUT_VALUE: _safe_json({"agent_name": _name}),
                    },
                ) as span:
                    if _tags:
                        span.set_attribute(OI.TAG_TAGS, _tags)
                    try:
                        result = fn(*args, **kwargs)
                        span.set_status(StatusCode.OK)
                        return result
                    except Exception as exc:
                        span.set_status(StatusCode.ERROR, str(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,
) -> Callable[[F], F]:
    """Decorator for tool functions."""
    _tags = list(tags or [])

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

        if asyncio.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,
                    attributes={
                        OI.OPENINFERENCE_SPAN_KIND: OpenInferenceSpanKindValues.TOOL.value,
                        OI.TOOL_NAME: _name,
                        OI.INPUT_VALUE: _safe_json(captured),
                    },
                ) as span:
                    if _tags:
                        span.set_attribute(OI.TAG_TAGS, _tags)
                    try:
                        result = await fn(*args, **kwargs)
                        span.set_attribute(OI.OUTPUT_VALUE, _safe_json(result))
                        span.set_status(StatusCode.OK)
                        return result
                    except Exception as exc:
                        span.set_status(StatusCode.ERROR, str(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,
                    attributes={
                        OI.OPENINFERENCE_SPAN_KIND: OpenInferenceSpanKindValues.TOOL.value,
                        OI.TOOL_NAME: _name,
                        OI.INPUT_VALUE: _safe_json(captured),
                    },
                ) as span:
                    if _tags:
                        span.set_attribute(OI.TAG_TAGS, _tags)
                    try:
                        result = fn(*args, **kwargs)
                        span.set_attribute(OI.OUTPUT_VALUE, _safe_json(result))
                        span.set_status(StatusCode.OK)
                        return result
                    except Exception as exc:
                        span.set_status(StatusCode.ERROR, str(exc))
                        span.record_exception(exc)
                        raise

            return _sync  # type: ignore[return-value]

    return decorator


def trace_fn(
    name: str | None = None,
    *,
    kind: str = "CHAIN",
    tags: list[str] | None = None,
) -> Callable[[F], F]:
    """Generic span decorator for arbitrary functions."""
    _tags = list(tags or [])

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

        if asyncio.iscoroutinefunction(fn):

            @functools.wraps(fn)
            async def _async(*args: Any, **kwargs: Any) -> Any:
                tracer = _get_tracer()
                with tracer.start_as_current_span(
                    _name,
                    attributes={OI.OPENINFERENCE_SPAN_KIND: kind},
                ) as span:
                    if _tags:
                        span.set_attribute(OI.TAG_TAGS, _tags)
                    try:
                        result = await fn(*args, **kwargs)
                        span.set_status(StatusCode.OK)
                        return result
                    except Exception as exc:
                        span.set_status(StatusCode.ERROR, str(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()
                with tracer.start_as_current_span(
                    _name,
                    attributes={OI.OPENINFERENCE_SPAN_KIND: kind},
                ) as span:
                    if _tags:
                        span.set_attribute(OI.TAG_TAGS, _tags)
                    try:
                        result = fn(*args, **kwargs)
                        span.set_status(StatusCode.OK)
                        return result
                    except Exception as exc:
                        span.set_status(StatusCode.ERROR, str(exc))
                        span.record_exception(exc)
                        raise

            return _sync  # type: ignore[return-value]

    return decorator
