from typing import Any, Callable, Dict, Mapping, Optional, Union, cast

import wrapt
from opentelemetry.trace import Span
from opentelemetry.util.types import AttributeValue

from openinference.semconv.trace import (
    OpenInferenceSpanKindValues,
    SpanAttributes,
)

from ._attributes import (
    get_input_attributes,
    get_output_attributes,
    get_tool_attributes,
)
from ._genai_conversion import get_genai_attributes
from ._types import OpenInferenceMimeType
from .config import (
    TraceConfig,
    mask_without_externalization,
)

_IMPORTANT_ATTRIBUTES = [
    SpanAttributes.OPENINFERENCE_SPAN_KIND,
]


class OpenInferenceSpan(wrapt.ObjectProxy):  # type: ignore[misc,name-defined,type-arg,unused-ignore]
    def __init__(self, wrapped: Span, config: TraceConfig) -> None:
        super().__init__(wrapped)
        self._self_config = config
        self._self_important_attributes: Dict[str, AttributeValue] = {}
        self._self_ended = False

    def set_attributes(self, attributes: "Mapping[str, AttributeValue]") -> None:
        for k, v in attributes.items():
            self.set_attribute(k, v)

    def set_attribute(
        self,
        key: str,
        value: Union[AttributeValue, Callable[[], AttributeValue]],
    ) -> None:
        if key in _IMPORTANT_ATTRIBUTES:
            # Always bookkeep important attributes (e.g. the OpenInference
            # span kind) so helpers like set_input keep working even on
            # non-recording spans.
            masked_value = self._self_config.mask(key, value)
            if masked_value is not None:
                self._self_important_attributes[key] = masked_value
            return
        if not self.is_recording():
            # Setting attributes on a non-recording span is a no-op anyway;
            # returning early keeps masking side effects (notably blob
            # uploads) from running for spans nobody will see. For spans
            # that already ended, forward the (side-effect-free) masked
            # value so the SDK still surfaces its "Setting attribute on
            # ended span" warning — the value is dropped either way.
            if self._self_ended:
                masked_value = mask_without_externalization(self._self_config, key, value)
                if masked_value is not None:
                    cast(Span, self.__wrapped__).set_attribute(key, masked_value)
            return
        masked_value = self._self_config.mask(key, value)
        if masked_value is not None:
            span = cast(Span, self.__wrapped__)
            span.set_attribute(key, masked_value)

    def end(self, end_time: Optional[int] = None) -> None:
        self._self_ended = True
        span = cast(Span, self.__wrapped__)
        for k, v in reversed(self._self_important_attributes.items()):
            span.set_attribute(k, v)
        if self._self_config.enable_genai_semconv:
            source_attributes = span._attributes
            for key, value in get_genai_attributes(source_attributes).items():
                if key not in source_attributes:
                    span.set_attribute(key, value)
        span.end(end_time)

    def set_input(
        self,
        value: Any,
        *,
        mime_type: Optional[OpenInferenceMimeType] = None,
    ) -> None:
        if OPENINFERENCE_SPAN_KIND not in self._self_important_attributes:
            raise ValueError("Cannot set input attributes on a non-OpenInference span")
        self.set_attributes(get_input_attributes(value, mime_type=mime_type))

    def set_output(
        self,
        value: Any,
        *,
        mime_type: Optional[OpenInferenceMimeType] = None,
    ) -> None:
        if OPENINFERENCE_SPAN_KIND not in self._self_important_attributes:
            raise ValueError("Cannot set output attributes on a non-OpenInference span")
        self.set_attributes(get_output_attributes(value, mime_type=mime_type))

    def set_tool(
        self,
        *,
        name: str,
        description: Optional[str] = None,
        parameters: Union[str, Dict[str, Any]],
    ) -> None:
        if self._self_important_attributes.get(OPENINFERENCE_SPAN_KIND) != TOOL:
            raise ValueError("Cannot set tool attributes on a non-tool span")
        self.set_attributes(
            get_tool_attributes(
                name=name,
                description=description,
                parameters=parameters,
            )
        )


# span kinds
TOOL = OpenInferenceSpanKindValues.TOOL.value


# span attributes
OPENINFERENCE_SPAN_KIND = SpanAttributes.OPENINFERENCE_SPAN_KIND
