from typing import Optional, TypeVar, Callable, Any, ParamSpec, Awaitable
import warnings

P = ParamSpec("P")
R = TypeVar("R")
F = TypeVar("F", bound=Callable[P, R | Awaitable[R]])


try:
    from opentelemetry.semconv_ai import TraceloopSpanKindValues
    from atatus.llmobs.decorators.base import (
        entity_class,
        entity_method,
    )
    OTEL_AVAILABLE = True
except ImportError:
    OTEL_AVAILABLE = False

    # Define dummy TraceloopSpanKindValues
    class TraceloopSpanKindValues:
        TASK = "task"
        WORKFLOW = "workflow"
        AGENT = "agent"
        TOOL = "tool"

    # Define no-op versions of entity_class and entity_method
    def entity_class(*args, **kwargs):
        def decorator(cls_or_func):
            return cls_or_func
        return decorator

    def entity_method(*args, **kwargs):
        def decorator(func):
            return func
        return decorator


def _make_decorator_factory(default_kind: str, is_task: bool):
    def decorator_factory(
        name: Optional[str] = None,
        version: Optional[int] = None,
        method_name: Optional[str] = None,
        tlp_span_kind: Optional[TraceloopSpanKindValues] = default_kind,
    ) -> Callable[[F], F]:
        if not OTEL_AVAILABLE:
            # return dummy decorator
            def dummy_decorator(func_or_cls):
                return func_or_cls
            return dummy_decorator

        if method_name is None:
            return entity_method(name=name, version=version, tlp_span_kind=tlp_span_kind)
        else:
            return entity_class(
                name=name,
                version=version,
                method_name=method_name,
                tlp_span_kind=tlp_span_kind,
            )

    return decorator_factory


task = _make_decorator_factory(TraceloopSpanKindValues.TASK, is_task=True)
workflow = _make_decorator_factory(TraceloopSpanKindValues.WORKFLOW, is_task=False)

def agent(
    name: Optional[str] = None,
    version: Optional[int] = None,
    method_name: Optional[str] = None,
) -> Callable[[F], F]:
    return workflow(
        name=name,
        version=version,
        method_name=method_name,
        tlp_span_kind=TraceloopSpanKindValues.AGENT,
    )


def tool(
    name: Optional[str] = None,
    version: Optional[int] = None,
    method_name: Optional[str] = None,
) -> Callable[[F], F]:
    return task(
        name=name,
        version=version,
        method_name=method_name,
        tlp_span_kind=TraceloopSpanKindValues.TOOL,
    )


def atask(
    name: Optional[str] = None,
    version: Optional[int] = None,
    method_name: Optional[str] = None,
    tlp_span_kind: Optional[TraceloopSpanKindValues] = TraceloopSpanKindValues.TASK,
) -> Callable[[F], F]:
    warnings.warn(
        "DeprecationWarning: The @atask decorator will be removed in a future version. "
        "Please migrate to @task for both sync and async operations.",
        DeprecationWarning,
        stacklevel=2,
    )
    return task(name=name, version=version, method_name=method_name, tlp_span_kind=tlp_span_kind)


def aworkflow(
    name: Optional[str] = None,
    version: Optional[int] = None,
    method_name: Optional[str] = None,
    tlp_span_kind: Optional[TraceloopSpanKindValues] = TraceloopSpanKindValues.WORKFLOW,
) -> Callable[[F], F]:
    warnings.warn(
        "DeprecationWarning: The @aworkflow decorator will be removed in a future version. "
        "Please migrate to @workflow for both sync and async operations.",
        DeprecationWarning,
        stacklevel=2,
    )
    return workflow(name=name, version=version, method_name=method_name, tlp_span_kind=tlp_span_kind)


def aagent(
    name: Optional[str] = None,
    version: Optional[int] = None,
    method_name: Optional[str] = None,
) -> Callable[[F], F]:
    warnings.warn(
        "DeprecationWarning: The @aagent decorator will be removed in a future version. "
        "Please migrate to @agent for both sync and async operations.",
        DeprecationWarning,
        stacklevel=2,
    )
    return agent(name=name, version=version, method_name=method_name)


def atool(
    name: Optional[str] = None,
    version: Optional[int] = None,
    method_name: Optional[str] = None,
) -> Callable[[F], F]:
    warnings.warn(
        "DeprecationWarning: The @atool decorator will be removed in a future version. "
        "Please migrate to @tool for both sync and async operations.",
        DeprecationWarning,
        stacklevel=2,
    )
    return tool(name=name, version=version, method_name=method_name)
