from collections.abc import Generator
from contextlib import contextmanager
from enum import Enum
from typing import Any, Optional

from opentelemetry import trace
from opentelemetry.trace import Span, SpanKind

from .._agent_instantiation import AgentInstantiationContext

# OpenTelemetry semantic convention constants for GenAI operations
# Copied from opentelemetry-semantic-conventions to avoid dependency

# GenAI Agent attributes
GEN_AI_AGENT_DESCRIPTION = "gen_ai.agent.description"
GEN_AI_AGENT_ID = "gen_ai.agent.id"
GEN_AI_AGENT_NAME = "gen_ai.agent.name"

# GenAI Operation attributes
GEN_AI_OPERATION_NAME = "gen_ai.operation.name"
GEN_AI_SYSTEM = "gen_ai.system"

# GenAI Tool attributes
GEN_AI_TOOL_CALL_ID = "gen_ai.tool.call.id"
GEN_AI_TOOL_DESCRIPTION = "gen_ai.tool.description"
GEN_AI_TOOL_NAME = "gen_ai.tool.name"

# Error attributes
ERROR_TYPE = "error.type"


class GenAiOperationNameValues(Enum):
    """Enum for GenAI operation name values."""

    CHAT = "chat"
    CREATE_AGENT = "create_agent"
    EMBEDDINGS = "embeddings"
    EXECUTE_TOOL = "execute_tool"
    GENERATE_CONTENT = "generate_content"
    INVOKE_AGENT = "invoke_agent"
    TEXT_COMPLETION = "text_completion"


# Constant for system name
GENAI_SYSTEM_AUTOGEN = "autogen"


@contextmanager
def trace_tool_span(
    tool_name: str,
    *,
    tracer: Optional[trace.Tracer] = None,
    parent: Optional[Span] = None,
    tool_description: Optional[str] = None,
    tool_call_id: Optional[str] = None,
) -> Generator[Span, Any, None]:
    """Context manager to create a span for tool execution following the
    OpenTelemetry Semantic conventions for generative AI systems.

    See the GenAI semantic conventions documentation:
    `OpenTelemetry GenAI Semantic Conventions <https://opentelemetry.io/docs/specs/semconv/gen-ai/>`__

    .. warning::

        The GenAI Semantic Conventions are still in incubation and
        subject to changes in future releases.


    Args:
        tool_name (str): The name of the tool being executed.
        tracer (Optional[trace.Tracer]): The tracer to use for creating the span.
        parent (Optional[Span]): The parent span to link this span to.
        tool_description (Optional[str]): A description of the tool.
        tool_call_id (Optional[str]): A unique identifier for the tool call.
    """
    if tracer is None:
        tracer = trace.get_tracer("autogen-core")
    span_attributes = {
        GEN_AI_OPERATION_NAME: GenAiOperationNameValues.EXECUTE_TOOL.value,
        GEN_AI_SYSTEM: GENAI_SYSTEM_AUTOGEN,
        GEN_AI_TOOL_NAME: tool_name,
    }
    if tool_description is not None:
        span_attributes[GEN_AI_TOOL_DESCRIPTION] = tool_description
    if tool_call_id is not None:
        span_attributes[GEN_AI_TOOL_CALL_ID] = tool_call_id
    with tracer.start_as_current_span(
        f"{GenAiOperationNameValues.EXECUTE_TOOL.value} {tool_name}",
        kind=SpanKind.INTERNAL,
        context=trace.set_span_in_context(parent) if parent else None,
        attributes=span_attributes,
    ) as span:
        try:
            yield span
        except Exception as e:
            # Set the exception details on the span if an error occurs
            span.record_exception(e)
            span.set_status(trace.Status(trace.StatusCode.ERROR, str(e)))
            span.set_attribute(ERROR_TYPE, type(e).__name__)
            raise


@contextmanager
def trace_create_agent_span(
    agent_name: str,
    *,
    tracer: Optional[trace.Tracer] = None,
    parent: Optional[Span] = None,
    agent_id: Optional[str] = None,
    agent_description: Optional[str] = None,
) -> Generator[Span, Any, None]:
    """Context manager to create a span for agent creation following the
    OpenTelemetry Semantic conventions for generative AI systems.

    See the GenAI semantic conventions documentation:
    `OpenTelemetry GenAI Semantic Conventions <https://opentelemetry.io/docs/specs/semconv/gen-ai/>`__

    .. warning::

        The GenAI Semantic Conventions are still in incubation and
        subject to changes in future releases.

    Args:
        agent_name (str): The name of the agent being created.
        tracer (Optional[trace.Tracer]): The tracer to use for creating the span.
        parent (Optional[Span]): The parent span to link this span to.
        agent_id (Optional[str]): The unique identifier for the agent.
        agent_description (Optional[str]): A description of the agent.
    """
    if tracer is None:
        tracer = trace.get_tracer("autogen-core")
    span_attributes = {
        GEN_AI_OPERATION_NAME: GenAiOperationNameValues.CREATE_AGENT.value,
        GEN_AI_SYSTEM: GENAI_SYSTEM_AUTOGEN,
        GEN_AI_AGENT_NAME: agent_name,
    }
    if agent_id is None:
        # Try to see if we can get the agent ID from the current context
        try:
            agent_id = str(AgentInstantiationContext.current_agent_id())
        except RuntimeError:
            agent_id = None
    if agent_id is not None:
        span_attributes[GEN_AI_AGENT_ID] = agent_id
    if agent_description is not None:
        span_attributes[GEN_AI_AGENT_DESCRIPTION] = agent_description
    with tracer.start_as_current_span(
        f"{GenAiOperationNameValues.CREATE_AGENT.value} {agent_name}",
        kind=SpanKind.CLIENT,
        context=trace.set_span_in_context(parent) if parent else None,
        attributes=span_attributes,
    ) as span:
        try:
            yield span
        except Exception as e:
            # Set the exception details on the span if an error occurs
            span.record_exception(e)
            span.set_status(trace.Status(trace.StatusCode.ERROR, str(e)))
            span.set_attribute(ERROR_TYPE, type(e).__name__)
            raise


@contextmanager
def trace_invoke_agent_span(
    agent_name: str,
    *,
    tracer: Optional[trace.Tracer] = None,
    parent: Optional[Span] = None,
    agent_id: Optional[str] = None,
    agent_description: Optional[str] = None,
) -> Generator[Span, Any, None]:
    """Context manager to create a span for invoking an agent following the
    OpenTelemetry Semantic conventions for generative AI systems.

    See the GenAI semantic conventions documentation:
    `OpenTelemetry GenAI Semantic Conventions <https://opentelemetry.io/docs/specs/semconv/gen-ai/>`__

    .. warning::

        The GenAI Semantic Conventions are still in incubation and
        subject to changes in future releases.

    Args:
        agent_name (str): The name of the agent being invoked.
        tracer (Optional[trace.Tracer]): The tracer to use for creating the span.
        parent (Optional[Span]): The parent span to link this span to.
        agent_id (Optional[str]): The unique identifier for the agent.
        agent_description (Optional[str]): A description of the agent.
    """
    if tracer is None:
        tracer = trace.get_tracer("autogen-core")
    span_attributes = {
        GEN_AI_OPERATION_NAME: GenAiOperationNameValues.INVOKE_AGENT.value,
        GEN_AI_SYSTEM: GENAI_SYSTEM_AUTOGEN,
        GEN_AI_AGENT_NAME: agent_name,
    }
    if agent_id is not None:
        span_attributes[GEN_AI_AGENT_ID] = agent_id
    if agent_description is not None:
        span_attributes[GEN_AI_AGENT_DESCRIPTION] = agent_description
    with tracer.start_as_current_span(
        f"{GenAiOperationNameValues.INVOKE_AGENT.value} {agent_name}",
        kind=SpanKind.CLIENT,
        context=trace.set_span_in_context(parent) if parent else None,
        attributes=span_attributes,
    ) as span:
        try:
            yield span
        except Exception as e:
            # Set the exception details on the span if an error occurs
            span.record_exception(e)
            span.set_status(trace.Status(trace.StatusCode.ERROR, str(e)))
            span.set_attribute(ERROR_TYPE, type(e).__name__)
            raise
