#  Copyright 2026 Synnax Labs, Inc.
#
#  Use of this software is governed by the Business Source License included in the file
#  licenses/BSL.txt.
#
#  As of the Change Date specified in that file, in accordance with the Business Source
#  License, use of this software will be governed by the Apache License, Version 2.0,
#  included in the file licenses/APL.txt.

from collections.abc import Iterator
from contextlib import contextmanager
from typing import Protocol

from opentelemetry.propagators.textmap import Setter, TextMapPropagator
from opentelemetry.trace import Span as OtelSpan
from opentelemetry.trace import (
    StatusCode,
)
from opentelemetry.trace import Tracer as OtelTracer
from opentelemetry.trace import TracerProvider as OtelTraceProvider

from alamos.environment import Environment, EnvironmentFilter, env_threshold_filter
from alamos.meta import InstrumentationMeta
from alamos.noop import Noop
from alamos.noop import noop as noopd


class Carrier(Protocol):
    def set(self, key: str, value: str) -> None: ...


class Span(Protocol):
    """A protocol for a Span that is part of a trace."""

    key: str
    """
    The key identifying the span. This is the name of the key passed into `trace`
    combined with the path of the instrumentation that started the span. For example,
    take instrumentation titled \"synnax\" and a call to trace with \"test\". The span key
    would be \"synnax.test\".
    """

    def record_exception(self, exc: Exception | None) -> None:
        """
        If exception is not none, records it on the span and sets the span's status to
        error.

        :param exc: An optionally defined exception to record.
        """
        ...


class _Span:
    """Base implementation of the Span protocol"""

    otel: OtelSpan
    key: str

    def __init__(self, otel: OtelSpan, key: str):
        self.otel = otel
        self.key = key

    def _(self) -> Span:
        return self

    def record_exception(self, exc: Exception | None) -> None:
        if exc is None:
            return
        self.otel.record_exception(exc)
        self.otel.set_status(StatusCode.ERROR)


class NoopSpan:
    """Span that does nothing."""

    key: str

    def _(self) -> Span:
        return self

    def record_exception(self, exc: Exception | None) -> None: ...


_NOOP_SPAN = NoopSpan()


class Tracer:
    """
    Tracer wraps an OpenTelemetry tracer to provide an opinionated interface for tracing
    within the Synnax stack.
    """

    noop: bool
    _meta: InstrumentationMeta
    _filter: EnvironmentFilter
    _otel_provider: OtelTraceProvider | None
    _otel_propagator: TextMapPropagator | None
    __otel_tracer: OtelTracer | None

    def _(self) -> Noop:
        return self

    def __init__(
        self,
        otel_provider: OtelTraceProvider | None = None,
        otel_propagator: TextMapPropagator | None = None,
        filter_: EnvironmentFilter = env_threshold_filter("debug"),
    ):
        self.noop = otel_provider is None and otel_propagator is None
        self._filter = filter_
        self._otel_provider = otel_provider
        self._otel_propagator = otel_propagator
        self.__otel_tracer = None

    @property
    def _otel_tracer(self) -> OtelTracer | None:
        if self.__otel_tracer is not None:
            return self.__otel_tracer
        if self._otel_provider is not None:
            self.__otel_tracer = self._otel_provider.get_tracer(self._meta.key)
            return self.__otel_tracer
        return None

    @contextmanager
    def trace(self, key: str, env: Environment) -> Iterator[Span]:
        """
        Starts a new span with the given key and environment. If a span already exists
        on the current context, the new span is made as its child.

        :param key: The key of the span.
        :param env: The environment to run the span under.
        :return: A span that tracks program execution. If the Tracer's environment
            filter rejects the provided env or the Tracer is noop, a no-op span is
            provided.
        """
        if self.noop or not self._filter(env) or self._otel_tracer is None:
            yield _NOOP_SPAN
            return
        with self._otel_tracer.start_as_current_span(key) as span:
            yield _Span(otel=span, key=self._meta.extend_path_(key))

    @contextmanager
    def debug(self, key: str) -> Iterator[Span]:
        """
        Starts a new span at the debug level. If a span already exists on the current
        context, the new span is made as its child.

        :param key: The key of the span.
        :return: A span that tracks program execution. If the Tracer's environment
            filter rejects the debug env or the Tracer is noop, a no-op span is
            provided.
        """
        with self.trace(key, "debug") as span:
            yield span

    @contextmanager
    def bench(self, key: str) -> Iterator[Span]:
        """
        Starts a new span at the bench level. If a span already exists on the current
        context, the new span is made as its child.

        :param key: The key of the span.
        :return: A span that tracks program execution. If the Tracer's environment
            filter rejects the bench env or the Tracer is noop, a no-op span is
            provided.
        """
        with self.trace(key, "bench") as span:
            yield span

    @contextmanager
    def prod(self, key: str) -> Iterator[Span]:
        """
        Starts a new span at the prod level. If a span already exists on the current
        context, the new span is made as its child.

        :param key: The key of the span.
        :return: A span that tracks program execution. If the Tracer's environment
            filter rejects the prod env or the Tracer is noop, a no-op span is provided.
        """
        with self.trace(key, "prod") as span:
            yield span

    class _Setter(Setter[Carrier]):
        @staticmethod
        def set(carrier: Carrier, key: str, value: str) -> None:
            print("setting", key, value)
            carrier.set(key, value)

    @noopd
    def propagate(self, carrier: Carrier) -> None:
        """
        Injects metadata about the current trace into the provided carrier. This
        metadata can be parsed on the other side of a network or IPC request using
        allowing the trace to propagate across services.

        :param carrier: The carrier to set the trace metadata on.
        """
        if self._otel_propagator is None:
            return
        self._otel_propagator.inject(carrier, setter=Tracer._Setter())

    def child_(self, meta: InstrumentationMeta) -> "Tracer":
        t = Tracer(
            otel_provider=self._otel_provider,
            otel_propagator=self._otel_propagator,
        )
        t._meta = meta
        return t


NOOP_TRACER = Tracer()
"""A tracer that does nothing."""
