from __future__ import annotations

from dataclasses import dataclass
from functools import cached_property
from threading import Lock
from typing import Callable, Literal

from opentelemetry import context
from opentelemetry.sdk.trace import ReadableSpan, Span, SpanProcessor
from opentelemetry.sdk.trace.sampling import Sampler, TraceIdRatioBased
from typing_extensions import Self

from logfire._internal.constants import (
    ONE_SECOND_IN_NANOSECONDS,
    LevelName,
)
from logfire._internal.exporters.wrapper import WrapperSpanProcessor
from logfire._internal.utils import suppress_instrumentation
from logfire.types import SpanLevel


@dataclass
class TraceBuffer:
    """Arguments of `SpanProcessor.on_start` and `SpanProcessor.on_end` for spans in a single trace.

    These are stored until either the trace is included by tail sampling or it's completed and discarded.
    """

    started: list[tuple[Span, context.Context | None]]
    ended: list[ReadableSpan]
    first_span: Span

    @cached_property
    def trace_id(self) -> int:
        span_context = self.first_span.context
        assert span_context is not None
        return span_context.trace_id


@dataclass
class TailSamplingSpanInfo:
    """Argument passed to the [`SamplingOptions.tail`][logfire.sampling.SamplingOptions.tail] callback."""

    span: ReadableSpan
    """
    Raw span object being started or ended.
    """

    context: context.Context | None
    """
    Second argument of [`SpanProcessor.on_start`][opentelemetry.sdk.trace.SpanProcessor.on_start] or `None` for `SpanProcessor.on_end`.
    """

    event: Literal['start', 'end']
    """
    `'start'` if the span is being started, `'end'` if it's being ended.
    """

    # Intentionally undocumented for now.
    buffer: TraceBuffer

    @property
    def level(self) -> SpanLevel:
        """The log level of the span."""
        return SpanLevel.from_span(self.span)

    @property
    def duration(self) -> float:
        """The time in seconds between the start of the trace and the start/end of this span."""
        # span.end_time and span.start_time are in nanoseconds and can be None.
        return (
            (self.span.end_time or self.span.start_time or 0) - (self.buffer.first_span.start_time or float('inf'))
        ) / ONE_SECOND_IN_NANOSECONDS


@dataclass
class SamplingOptions:
    """Options for [`logfire.configure(sampling=...)`][logfire.configure(sampling)].

    See the [sampling guide](https://logfire.pydantic.dev/docs/guides/advanced/sampling/).
    """

    head: float | Sampler = 1.0
    """
    Head sampling options.

    If it's a float, it should be a number between 0.0 and 1.0.
    This is the probability that an entire trace will randomly included.

    Alternatively you can pass a custom
    [OpenTelemetry `Sampler`](https://opentelemetry-python.readthedocs.io/en/latest/sdk/trace.sampling.html).
    """

    tail: Callable[[TailSamplingSpanInfo], float] | None = None
    """
    An optional tail sampling callback which will be called for every span.

    It should return a number between 0.0 and 1.0, the probability that the entire trace will be included.
    Use [`SamplingOptions.level_or_duration`][logfire.sampling.SamplingOptions.level_or_duration]
    for a common use case.

    Every span in a trace will be stored in memory until either the trace is included by tail sampling
    or it's completed and discarded, so large traces may consume a lot of memory.
    """

    @classmethod
    def level_or_duration(
        cls,
        *,
        head: float | Sampler = 1.0,
        level_threshold: LevelName | None = 'notice',
        duration_threshold: float | None = 5.0,
        background_rate: float = 0.0,
    ) -> Self:
        """Returns a `SamplingOptions` instance that tail samples traces based on their log level and duration.

        If a trace has at least one span/log that has a log level greater than or equal to `level_threshold`,
        or if the duration of the whole trace is greater than `duration_threshold` seconds,
        then the whole trace will be included.
        Otherwise, the probability is `background_rate`.

        The `head` parameter is the same as in the `SamplingOptions` constructor.
        """
        head_sample_rate = head if isinstance(head, (float, int)) else 1.0

        if not (0.0 <= background_rate <= head_sample_rate <= 1.0):
            raise ValueError('Invalid sampling rates, must be 0.0 <= background_rate <= head <= 1.0')

        def get_tail_sample_rate(span_info: TailSamplingSpanInfo) -> float:
            if duration_threshold is not None and span_info.duration > duration_threshold:
                return 1.0

            if level_threshold is not None and span_info.level >= level_threshold:
                return 1.0

            return background_rate

        return cls(head=head, tail=get_tail_sample_rate)


def check_trace_id_ratio(trace_id: int, rate: float) -> bool:
    # Based on TraceIdRatioBased.should_sample.
    return (trace_id & TraceIdRatioBased.TRACE_ID_LIMIT) < TraceIdRatioBased.get_bound_for_rate(rate)


class TailSamplingProcessor(WrapperSpanProcessor):
    """Buffers spans until a span in the trace meets the sampling criteria.

    The wrapped `processor` receives `on_start` calls immediately (most processors just set attributes,
    which is safe before the sampling decision). Only `on_end` is buffered.

    The optional `deferred_processor` has both `on_start` and `on_end` buffered until the sampling
    decision. Use this for processors like `PendingSpanProcessor` that create new spans in `on_start`.
    Advanced users can mark their own processors for deferral by setting the
    `tail_sampling_defer_on_start` attribute to `True` on the processor class.
    """

    def __init__(
        self,
        processor: SpanProcessor,
        get_tail_sample_rate: Callable[[TailSamplingSpanInfo], float],
        deferred_processor: SpanProcessor | None = None,
    ) -> None:
        super().__init__(processor)
        self.get_tail_sample_rate = get_tail_sample_rate
        self.deferred_processor = deferred_processor

        # A TraceBuffer is typically created for each new trace.
        # If a span meets the sampling criteria, the buffer is dropped and all buffered spans
        # are pushed to the wrapped processor (and the deferred processor if present).
        # So when more spans arrive and there's no buffer, they get passed through immediately.
        self.traces: dict[int, TraceBuffer] = {}

        # Code that touches self.traces and its contents should be protected by this lock.
        self.lock = Lock()

    def on_start(self, span: Span, parent_context: context.Context | None = None) -> None:
        dropped = False
        buffer = None

        with self.lock:
            # span.context could supposedly be None, not sure how.
            if span.context:  # pragma: no branch
                trace_id = span.context.trace_id
                # If span.parent is None, it's the root span of a trace.
                if span.parent is None:
                    self.traces[trace_id] = TraceBuffer(started=[], ended=[], first_span=span)

                buffer = self.traces.get(trace_id)
                if buffer is not None:
                    # This trace's spans haven't met the criteria yet, so add this span to the buffer.
                    # Only track started spans if there's a deferred processor that needs replay.
                    if self.deferred_processor is not None:
                        buffer.started.append((span, parent_context))
                    dropped = self.check_span(TailSamplingSpanInfo(span, parent_context, 'start', buffer))

        # Always call the main processor's on_start immediately.
        # Most processors just set attributes, which is safe before the sampling decision.
        # This code may take longer since it calls processors which might do anything.
        # It shouldn't be inside the lock to avoid blocking other threads.
        # Since it's not in the lock, it shouldn't touch self.traces or its contents.
        super().on_start(span, parent_context)

        if buffer is None:
            # No buffer for this trace, meaning it was already sampled/discarded, or never tracked.
            # Deferred processor also gets on_start immediately.
            if self.deferred_processor is not None:
                self.deferred_processor.on_start(span, parent_context)
        elif dropped:
            self.push_buffer(buffer)

    def on_end(self, span: ReadableSpan) -> None:
        # Unlike on_start (where the main processor is called immediately),
        # on_end is buffered for both the main processor and the deferred processor.

        dropped = False
        buffer = None

        with self.lock:
            if span.context:  # pragma: no branch
                trace_id = span.context.trace_id
                buffer = self.traces.get(trace_id)
                if buffer is not None:
                    buffer.ended.append(span)
                    dropped = self.check_span(TailSamplingSpanInfo(span, None, 'end', buffer))
                    if span.parent is None:
                        # This is the root span, so the trace is hopefully complete.
                        # Delete the buffer to save memory.
                        self.traces.pop(trace_id, None)

        # This code may take longer since it calls processors which might do anything.
        # It shouldn't be inside the lock to avoid blocking other threads.
        # Since it's not in the lock, it shouldn't touch self.traces or its contents.
        if buffer is None:
            # No buffer for this trace, meaning it was already sampled/discarded, or never tracked.
            # Pass on_end through immediately to both processors.
            super().on_end(span)
            if self.deferred_processor is not None:
                self.deferred_processor.on_end(span)
        elif dropped:
            self.push_buffer(buffer)

    def check_span(self, span_info: TailSamplingSpanInfo) -> bool:
        """If the span meets the sampling criteria, drop the buffer and return True. Otherwise, return False.

        Must be called within self.lock since it modifies self.traces via drop_buffer.
        """
        sample_rate = self.get_tail_sample_rate(span_info)
        if sampled := check_trace_id_ratio(span_info.buffer.trace_id, sample_rate):
            self.drop_buffer(span_info.buffer)

        return sampled

    def drop_buffer(self, buffer: TraceBuffer) -> None:
        del self.traces[buffer.trace_id]

    def push_buffer(self, buffer: TraceBuffer) -> None:
        # For the deferred processor, replay both on_start and on_end.
        if self.deferred_processor is not None:
            for started in buffer.started:
                self.deferred_processor.on_start(*started)
            for span in buffer.ended:
                self.deferred_processor.on_end(span)

        # For the main processor, on_start was already called immediately.
        # Only replay on_end.
        for span in buffer.ended:
            super().on_end(span)

    def shutdown(self) -> None:
        super().shutdown()
        if self.deferred_processor is not None:
            # suppress_instrumentation prevents the deferred processor (e.g. PendingSpanProcessor)
            # from creating new spans during shutdown, which would re-enter this processor
            # and cause infinite recursion. The superclass already does this for the main processor.
            with suppress_instrumentation():
                self.deferred_processor.shutdown()

    def force_flush(self, timeout_millis: int = 30000) -> bool:
        result = super().force_flush(timeout_millis)
        if self.deferred_processor is not None:
            # Same reasoning as shutdown above.
            with suppress_instrumentation():
                result = self.deferred_processor.force_flush(timeout_millis) and result
        return result