from typing import Optional

from opentelemetry.semconv.ai import TraceloopSpanKindValues

from traceloop.sdk.decorators.base import (
    aentity_class,
    aentity_method,
    entity_class,
    entity_method,
)


def task(
    name: Optional[str] = None,
    method_name: Optional[str] = None,
    tlp_span_kind: Optional[TraceloopSpanKindValues] = TraceloopSpanKindValues.TASK,
):
    if method_name is None:
        return entity_method(name=name, tlp_span_kind=tlp_span_kind)
    else:
        return entity_class(
            name=name, method_name=method_name, tlp_span_kind=tlp_span_kind
        )


def workflow(
    name: Optional[str] = None,
    method_name: Optional[str] = None,
    tlp_span_kind: Optional[TraceloopSpanKindValues] = TraceloopSpanKindValues.WORKFLOW,
):
    if method_name is None:
        return entity_method(name=name, tlp_span_kind=tlp_span_kind)
    else:
        return entity_class(
            name=name,
            method_name=method_name,
            tlp_span_kind=tlp_span_kind,
        )


def agent(name: Optional[str] = None, method_name: Optional[str] = None):
    return workflow(
        name=name, method_name=method_name, tlp_span_kind=TraceloopSpanKindValues.AGENT
    )


def tool(name: Optional[str] = None, method_name: Optional[str] = None):
    return task(
        name=name, method_name=method_name, tlp_span_kind=TraceloopSpanKindValues.TOOL
    )


# Async Decorators
def atask(
    name: Optional[str] = None,
    method_name: Optional[str] = None,
    tlp_span_kind: Optional[TraceloopSpanKindValues] = TraceloopSpanKindValues.TASK,
):
    if method_name is None:
        return aentity_method(name=name, tlp_span_kind=tlp_span_kind)
    else:
        return aentity_class(
            name=name, method_name=method_name, tlp_span_kind=tlp_span_kind
        )


def aworkflow(
    name: Optional[str] = None,
    method_name: Optional[str] = None,
    tlp_span_kind: Optional[TraceloopSpanKindValues] = TraceloopSpanKindValues.WORKFLOW,
):
    if method_name is None:
        return aentity_method(name=name, tlp_span_kind=tlp_span_kind)
    else:
        return aentity_class(
            name=name, method_name=method_name, tlp_span_kind=tlp_span_kind
        )


def aagent(name: Optional[str] = None, method_name: Optional[str] = None):
    return atask(
        name=name, method_name=method_name, tlp_span_kind=TraceloopSpanKindValues.AGENT
    )


def atool(name: Optional[str] = None, method_name: Optional[str] = None):
    return atask(
        name=name, method_name=method_name, tlp_span_kind=TraceloopSpanKindValues.TOOL
    )
