import random
import threading
from concurrent.futures import ThreadPoolExecutor
from unittest import mock

import pytest
from opentelemetry import trace
from opentelemetry.sdk.trace.id_generator import RandomIdGenerator

import mlflow
from mlflow.entities.trace_location import (
    MlflowExperimentLocation,
    UCSchemaLocation,
    UnityCatalog,
)
from mlflow.environment_variables import (
    MLFLOW_TRACE_ENABLE_OTLP_DUAL_EXPORT,
    MLFLOW_TRACE_SAMPLING_RATIO,
    MLFLOW_USE_DEFAULT_TRACER_PROVIDER,
)
from mlflow.exceptions import MlflowException, MlflowTracingException
from mlflow.tracing.destination import Databricks, MlflowExperiment
from mlflow.tracing.export.inference_table import (
    _TRACE_BUFFER,
    InferenceTableSpanExporter,
)
from mlflow.tracing.export.mlflow_v3 import MlflowV3SpanExporter
from mlflow.tracing.export.uc_table import DatabricksUCTableSpanExporter
from mlflow.tracing.fluent import start_span_no_context
from mlflow.tracing.processor.inference_table import InferenceTableSpanProcessor
from mlflow.tracing.processor.mlflow_v3 import MlflowV3SpanProcessor
from mlflow.tracing.processor.otel import OtelSpanProcessor
from mlflow.tracing.processor.uc_table import DatabricksUCTableSpanProcessor
from mlflow.tracing.provider import (
    _get_tracer,
    _initialize_tracer_provider,
    _IsolatedRandomIdGenerator,
    is_tracing_enabled,
    start_span_in_context,
    trace_disabled,
)
from mlflow.tracing.provider import (
    provider as _provider_wrapper,
)
from mlflow.tracing.utils import get_active_spans_table_name
from mlflow.utils.mlflow_tags import (
    MLFLOW_EXPERIMENT_DATABRICKS_TRACE_DESTINATION_PATH,
    MLFLOW_EXPERIMENT_DATABRICKS_TRACE_LOG_STORAGE_TABLE,
    MLFLOW_EXPERIMENT_DATABRICKS_TRACE_SPAN_STORAGE_TABLE,
)

from tests.tracing.helper import get_traces, purge_traces, skip_when_testing_trace_sdk


@pytest.fixture
def mock_setup_tracer_provider():
    # To count the number of times _initialize_tracer_provider is called
    with mock.patch(
        "mlflow.tracing.provider._initialize_tracer_provider",
        side_effect=_initialize_tracer_provider,
    ) as setup_mock:
        yield setup_mock


def test_tracer_provider_initialized_once(mock_setup_tracer_provider):
    assert mock_setup_tracer_provider.call_count == 0
    start_span_in_context("test1")
    assert mock_setup_tracer_provider.call_count == 1

    start_span_in_context("test_2")
    start_span_in_context("test_3")
    assert mock_setup_tracer_provider.call_count == 1

    # Thread safety
    with ThreadPoolExecutor(max_workers=2, thread_name_prefix="test-tracing-provider") as executor:
        executor.map(start_span_in_context, ["test_4", "test_5"])
    assert mock_setup_tracer_provider.call_count == 1


def test_reset_tracer_setup(mock_setup_tracer_provider):
    assert mock_setup_tracer_provider.call_count == 0

    start_span_in_context("test1")
    assert mock_setup_tracer_provider.call_count == 1

    mlflow.tracing.reset()
    assert mock_setup_tracer_provider.call_count == 2

    start_span_in_context("test2")
    assert mock_setup_tracer_provider.call_count == 3
    assert mock_setup_tracer_provider.mock_calls == ([
        mock.call(),
        mock.call(disabled=True),
        mock.call(),
    ])


def test_span_processor_and_exporter_model_serving(mock_databricks_serving_with_tracing_env):
    tracer = _get_tracer("test")
    processors = tracer.span_processor._span_processors
    assert len(processors) == 1
    assert isinstance(processors[0], InferenceTableSpanProcessor)
    assert isinstance(processors[0].span_exporter, InferenceTableSpanExporter)


def test_set_destination_mlflow_experiment(monkeypatch):
    mlflow.tracing.set_destination(destination=MlflowExperimentLocation(experiment_id="123"))

    tracer = _get_tracer("test")
    processors = tracer.span_processor._span_processors
    assert isinstance(processors[0], MlflowV3SpanProcessor)
    assert isinstance(processors[0].span_exporter, MlflowV3SpanExporter)


def test_set_destination_databricks(monkeypatch):
    monkeypatch.setenv("MLFLOW_TRACKING_URI", "databricks")
    mlflow.tracing.set_destination(destination=Databricks(experiment_id="123"))

    tracer = _get_tracer("test")
    processors = tracer.span_processor._span_processors
    assert len(processors) == 1
    assert isinstance(processors[0], MlflowV3SpanProcessor)
    assert isinstance(processors[0].span_exporter, MlflowV3SpanExporter)


def test_set_destination_databricks_uc():
    with mock.patch("mlflow.tracing.provider._logger.warning") as mock_warning:
        mlflow.tracing.set_destination(
            destination=UCSchemaLocation(
                catalog_name="catalog",
                schema_name="schema",
            )
        )

    mock_warning.assert_called_once()
    assert "Passing `UCSchemaLocation` to `mlflow.tracing.set_destination` is deprecated" in str(
        mock_warning.call_args.args[0]
    )

    tracer = _get_tracer("test")
    processors = tracer.span_processor._span_processors
    assert len(processors) == 1
    assert isinstance(processors[0], DatabricksUCTableSpanProcessor)
    assert isinstance(processors[0].span_exporter, DatabricksUCTableSpanExporter)
    assert get_active_spans_table_name() == "catalog.schema.mlflow_experiment_trace_otel_spans"


def test_set_destination_databricks_unity_catalog_rejected(monkeypatch):
    with pytest.raises(
        MlflowException,
        match=r"UnityCatalog table-prefix destinations are not supported by "
        r"`mlflow\.tracing\.set_destination`",
    ):
        mlflow.tracing.set_destination(
            destination=UnityCatalog(
                catalog_name="catalog",
                schema_name="schema",
                table_prefix="prefix",
            )
        )


def test_set_destination_databricks_uc_with_oltp_env_no_dual_export(monkeypatch):
    # set_destination is called but OLTP is also set w/o dual export mode enabled
    monkeypatch.setenv(MLFLOW_TRACE_ENABLE_OTLP_DUAL_EXPORT.name, "false")
    with (
        mock.patch("mlflow.tracing.provider.should_use_otlp_exporter", return_value=True),
        mock.patch("mlflow.tracing.provider.get_otlp_exporter") as mock_get_exporter,
    ):
        mock_get_exporter.return_value = mock.MagicMock()

        mlflow.tracing.reset()
        mlflow.tracing.set_destination(
            destination=UCSchemaLocation(
                catalog_name="catalog",
                schema_name="schema",
            )
        )
        tracer = _get_tracer("test")
        processors = tracer.span_processor._span_processors
        assert len(processors) == 1
        assert isinstance(processors[0], DatabricksUCTableSpanProcessor)
        assert isinstance(processors[0].span_exporter, DatabricksUCTableSpanExporter)
        assert get_active_spans_table_name() == "catalog.schema.mlflow_experiment_trace_otel_spans"


def test_set_destination_databricks_uc_with_oltp_env_with_dual_export(monkeypatch):
    # set_destination is called but OLTP is also set w/ dual export mode enabled
    monkeypatch.setenv(MLFLOW_TRACE_ENABLE_OTLP_DUAL_EXPORT.name, "true")
    with (
        mock.patch("mlflow.tracing.provider.should_use_otlp_exporter", return_value=True),
        mock.patch("mlflow.tracing.provider.get_otlp_exporter") as mock_get_exporter,
    ):
        mock_get_exporter.return_value = mock.MagicMock()

        mlflow.tracing.reset()
        mlflow.tracing.set_destination(
            destination=UCSchemaLocation(
                catalog_name="catalog",
                schema_name="schema",
            )
        )
        tracer = _get_tracer("test")
        processors = tracer.span_processor._span_processors
        assert len(processors) == 2
        assert isinstance(processors[0], DatabricksUCTableSpanProcessor)
        assert isinstance(processors[0].span_exporter, DatabricksUCTableSpanExporter)
        # OTLP processor needs to be there for dual export mode
        assert isinstance(processors[1], OtelSpanProcessor)
        assert get_active_spans_table_name() == "catalog.schema.mlflow_experiment_trace_otel_spans"


def test_set_destination_from_env_var_mlflow_experiment(monkeypatch):
    monkeypatch.setenv("MLFLOW_TRACING_DESTINATION", "123")

    tracer = _get_tracer("test")
    processors = tracer.span_processor._span_processors
    assert len(processors) == 1
    assert isinstance(processors[0], MlflowV3SpanProcessor)
    assert isinstance(processors[0].span_exporter, MlflowV3SpanExporter)


def test_set_destination_from_env_var_databricks_uc(monkeypatch):
    monkeypatch.setenv("MLFLOW_TRACING_DESTINATION", "catalog.schema")

    tracer = _get_tracer("test")
    processors = tracer.span_processor._span_processors
    assert len(processors) == 1
    assert isinstance(processors[0], DatabricksUCTableSpanProcessor)
    assert isinstance(processors[0].span_exporter, DatabricksUCTableSpanExporter)
    assert get_active_spans_table_name() == "catalog.schema.mlflow_experiment_trace_otel_spans"


def test_set_destination_in_model_serving(mock_databricks_serving_with_tracing_env, monkeypatch):
    monkeypatch.setenv("MLFLOW_TRACKING_URI", "databricks")
    monkeypatch.setenv("MLFLOW_TRACING_DESTINATION", "catalog.schema")

    tracer = _get_tracer("test")
    processors = tracer.span_processor._span_processors
    assert len(processors) == 1
    assert isinstance(processors[0], DatabricksUCTableSpanProcessor)
    assert isinstance(processors[0].span_exporter, DatabricksUCTableSpanExporter)
    assert get_active_spans_table_name() == "catalog.schema.mlflow_experiment_trace_otel_spans"


def test_set_destination_deprecated_classes():
    from mlflow.tracing.provider import _MLFLOW_TRACE_USER_DESTINATION

    with pytest.warns(FutureWarning, match="`mlflow.tracing.destination.MlflowExperiment``"):
        mlflow.tracing.set_destination(destination=MlflowExperiment(experiment_id="123"))

    destination = _MLFLOW_TRACE_USER_DESTINATION.get()
    assert isinstance(destination, MlflowExperimentLocation)
    assert destination.experiment_id == "123"

    with pytest.warns(FutureWarning, match="`mlflow.tracing.destination.Databricks`"):
        mlflow.tracing.set_destination(destination=Databricks(experiment_id="123"))

    destination = _MLFLOW_TRACE_USER_DESTINATION.get()
    assert isinstance(destination, MlflowExperimentLocation)
    assert destination.experiment_id == "123"


def test_disable_enable_tracing():
    @mlflow.trace
    def test_fn():
        pass

    test_fn()
    assert len(get_traces()) == 1
    assert isinstance(_get_tracer(__name__), trace.Tracer)
    purge_traces()

    mlflow.tracing.disable()
    test_fn()
    assert len(get_traces()) == 0
    assert isinstance(_get_tracer(__name__), trace.NoOpTracer)

    mlflow.tracing.enable()
    test_fn()
    assert len(get_traces()) == 1
    assert isinstance(_get_tracer(__name__), trace.Tracer)

    # enable() / disable() should only raise MlflowTracingException
    with mock.patch(
        "mlflow.tracing.provider.is_tracing_enabled", side_effect=ValueError("error")
    ) as is_enabled_mock:
        with pytest.raises(MlflowTracingException, match="error"):
            mlflow.tracing.disable()
        assert is_enabled_mock.call_count == 1

        with pytest.raises(MlflowTracingException, match="error"):
            mlflow.tracing.enable()
        assert is_enabled_mock.call_count == 2


@pytest.mark.parametrize("enabled_initially", [True, False])
def test_trace_disabled_decorator(enabled_initially):
    if not enabled_initially:
        mlflow.tracing.disable()
    assert is_tracing_enabled() == enabled_initially
    call_count = 0

    @trace_disabled
    def test_fn():
        with mlflow.start_span(name="test_span") as span:
            span.set_attribute("key", "value")
        nonlocal call_count
        call_count += 1
        return 0

    test_fn()
    assert len(get_traces()) == 0
    assert call_count == 1

    # Recover the initial state
    assert is_tracing_enabled() == enabled_initially

    # Tracing should be enabled back even if the function raises an exception
    @trace_disabled
    def test_fn_raise():
        nonlocal call_count
        call_count += 1
        raise ValueError("error")

    with pytest.raises(ValueError, match="error"):
        test_fn_raise()
    assert call_count == 2

    assert len(get_traces()) == 0
    assert is_tracing_enabled() == enabled_initially

    # @trace_disabled should not block the decorated function even if the
    # tracing machinery errors while checking/toggling the tracing state.
    with mock.patch(
        "mlflow.tracing.provider.is_tracing_enabled",
        side_effect=MlflowTracingException("error"),
    ) as is_enabled_mock:
        assert test_fn() == 0
        assert call_count == 3
        assert is_enabled_mock.call_count == 1


@pytest.mark.parametrize("enabled_initially", [True, False])
def test_trace_disabled_with_mlflow_trace_raising(enabled_initially):
    if not enabled_initially:
        mlflow.tracing.disable()

    @mlflow.trace
    def predict_fn(query):
        raise ValueError("Uhoh!")

    @trace_disabled
    def validate():
        predict_fn(query="What is MLflow?")

    with pytest.raises(ValueError, match="Uhoh!"):
        validate()

    assert is_tracing_enabled() == enabled_initially
    assert len(get_traces()) == 0


def test_disable_enable_tracing_not_mutate_otel_provider(monkeypatch):
    monkeypatch.setenv(MLFLOW_USE_DEFAULT_TRACER_PROVIDER.name, "true")

    # This test validates that disable/enable MLflow tracing does not mutate the OpenTelemetry's
    # global tracer provider instance.
    otel_tracer_provider = trace.get_tracer_provider()

    mlflow.tracing.disable()
    assert trace.get_tracer_provider() is otel_tracer_provider

    mlflow.tracing.enable()
    assert trace.get_tracer_provider() is otel_tracer_provider

    @trace_disabled
    def test_fn():
        assert trace.get_tracer_provider() is otel_tracer_provider

    test_fn()
    assert trace.get_tracer_provider() is otel_tracer_provider


def _count_batch_processor_threads() -> int:
    return sum("OtelBatchSpanRecordProcessor" in t.name for t in threading.enumerate())


@pytest.fixture
def batch_span_processor(monkeypatch):
    # Force the async BatchSpanProcessor path (which owns the leaked thread).
    # The backend is supplied by the autouse conftest fixtures, so don't override
    # the tracking URI here (sqlite:// breaks the SDK-only job, which has no store).
    monkeypatch.setenv("MLFLOW_USE_BATCH_SPAN_PROCESSOR", "true")
    monkeypatch.setenv("MLFLOW_ENABLE_ASYNC_TRACE_LOGGING", "true")


def test_disable_enable_does_not_leak_batch_processor_threads(batch_span_processor):
    @mlflow.trace
    def f():
        return 0

    # Prime a real BatchSpanProcessor (and its daemon thread).
    f()
    baseline = _count_batch_processor_threads()
    # Guard against a vacuous pass: the batch path must actually be active.
    assert baseline >= 1

    # Each enable() used to build a fresh provider + BatchSpanProcessor without
    # shutting down the old one, leaking one thread per cycle (issue #24209).
    for _ in range(10):
        mlflow.tracing.disable()
        mlflow.tracing.enable()

    assert _count_batch_processor_threads() <= baseline


def test_trace_disabled_does_not_leak_batch_processor_threads(batch_span_processor):
    @mlflow.trace
    def f():
        return 0

    @trace_disabled
    def wrapped():
        return 0

    f()
    baseline = _count_batch_processor_threads()
    assert baseline >= 1

    # trace_disabled wraps load_model/log_model; it must not create or destroy
    # the BatchSpanProcessor thread per call.
    for _ in range(20):
        wrapped()

    assert _count_batch_processor_threads() <= baseline
    # Tracing is fully restored and still records after the decorated call.
    purge_traces()
    f()
    assert len(get_traces()) == 1


def test_disable_enable_no_data_loss_on_retire(batch_span_processor):
    # A span emitted just before enable() rebuilds the provider must still be
    # exported: retiring the outgoing processor force_flushes before shutdown.
    @mlflow.trace
    def f():
        return 0

    f()
    purge_traces()
    f()
    # Rebuild the provider; the pending span must survive the retire (flush-then-shutdown).
    mlflow.tracing.disable()
    mlflow.tracing.enable()
    assert len(get_traces()) == 1


def test_nested_trace_disabled_restores_tracing(batch_span_processor):
    @mlflow.trace
    def f():
        return 0

    @trace_disabled
    def inner():
        return is_tracing_enabled()

    @trace_disabled
    def outer():
        assert not is_tracing_enabled()
        inner_state = inner()
        # The inner frame must NOT restore tracing on its own exit: we are still
        # inside the outer frame, so tracing must remain disabled.
        assert not is_tracing_enabled()
        return inner_state

    f()
    baseline = _count_batch_processor_threads()

    assert outer() is False
    # Only the outermost exit restores tracing.
    assert is_tracing_enabled()
    assert _count_batch_processor_threads() <= baseline
    purge_traces()
    f()
    assert len(get_traces()) == 1


def test_trace_disabled_under_concurrency_smoke(batch_span_processor):
    # Smoke check that trace_disabled is safe under concurrent use: after many
    # overlapping calls tracing is still enabled and no BSP thread leaked. This
    # asserts the happy end state; it does not deterministically force the rare
    # swap/restore interleaving the depth guard defends against.
    @mlflow.trace
    def f():
        return 0

    @trace_disabled
    def wrapped():
        return 0

    f()
    baseline = _count_batch_processor_threads()

    with ThreadPoolExecutor(max_workers=8, thread_name_prefix="trace-disabled-test") as executor:
        futures = [executor.submit(wrapped) for _ in range(200)]
        for future in futures:
            future.result()

    assert is_tracing_enabled()
    assert _count_batch_processor_threads() <= baseline
    purge_traces()
    f()
    assert len(get_traces()) == 1


def test_otlp_span_processor_is_retired_on_provider_replace(monkeypatch):
    # OtelSpanProcessor subclasses OTel's BatchSpanProcessor directly (not a
    # BaseMlflowSpanProcessor), so it must be retired on provider replace too.
    from mlflow.tracing.provider import _get_tracer

    monkeypatch.setenv(MLFLOW_TRACE_ENABLE_OTLP_DUAL_EXPORT.name, "false")
    with (
        mock.patch("mlflow.tracing.provider.should_use_otlp_exporter", return_value=True),
        mock.patch("mlflow.tracing.provider.get_otlp_exporter"),
    ):
        mlflow.tracing.reset()
        tracer = _get_tracer("test")
        (otel_processor,) = tracer.span_processor._span_processors
        assert isinstance(otel_processor, OtelSpanProcessor)

        with (
            mock.patch.object(otel_processor, "force_flush") as force_flush,
            mock.patch.object(otel_processor, "shutdown") as shutdown,
        ):
            # Replacing the provider must flush before it shuts the processor down.
            mlflow.tracing.disable()
            force_flush.assert_called_once()
            shutdown.assert_called_once()


def test_set_experiment_survives_tracing_state_error():
    # is_tracing_enabled() is raise_as_trace_exception-wrapped; a tracing error
    # must not break set_experiment (issue #24209 review).
    mlflow.set_experiment("first")

    @mlflow.trace
    def f():
        return 0

    f()  # ensure provider.once._done so the preserve-disabled branch runs
    with mock.patch(
        "mlflow.tracing.provider.is_tracing_enabled",
        side_effect=MlflowTracingException("boom"),
    ):
        # Should not raise despite the tracing-state check failing.
        mlflow.set_experiment("second")


def test_set_experiment_preserves_explicit_disable():
    mlflow.set_experiment("first")

    mlflow.tracing.disable()
    assert not is_tracing_enabled()

    # set_experiment resets the provider to re-derive the destination; it must
    # not silently re-enable tracing the user explicitly turned off (issue #24209).
    mlflow.set_experiment("second")
    assert not is_tracing_enabled()

    mlflow.tracing.enable()
    assert is_tracing_enabled()


def test_is_tracing_enabled():
    # Before doing anything -> tracing is considered as "on"
    assert is_tracing_enabled()

    # Generate a trace -> tracing is still "on"
    @mlflow.trace
    def foo():
        pass

    foo()
    assert is_tracing_enabled()

    # Disable tracing
    mlflow.tracing.disable()
    assert is_tracing_enabled() is False

    # Try to generate a trace -> tracing is still "off"
    foo()
    assert is_tracing_enabled() is False

    # Re-enable tracing
    mlflow.tracing.enable()
    assert is_tracing_enabled() is True

    # is_tracing_enabled() should only raise MlflowTracingException
    with mock.patch(
        "mlflow.tracing.provider._get_tracer", side_effect=ValueError("error")
    ) as get_tracer_mock:
        with pytest.raises(MlflowTracingException, match="error"):
            assert is_tracing_enabled() is False
        assert get_tracer_mock.call_count == 1


@pytest.mark.parametrize("enable_mlflow_tracing", [True, False, None])
def test_enable_mlflow_tracing_switch_in_serving_fluent(monkeypatch, enable_mlflow_tracing):
    if enable_mlflow_tracing is None:
        monkeypatch.delenv("ENABLE_MLFLOW_TRACING", raising=False)
    else:
        monkeypatch.setenv("ENABLE_MLFLOW_TRACING", str(enable_mlflow_tracing).lower())
    monkeypatch.setenv("IS_IN_DB_MODEL_SERVING_ENV", "true")

    @mlflow.trace
    def foo():
        return 1

    request_ids = ["id1", "id2", "id3"]
    with mock.patch(
        "mlflow.tracing.processor.inference_table.maybe_get_request_id", side_effect=request_ids
    ):
        for _ in range(3):
            foo()

    if enable_mlflow_tracing:
        assert sorted(_TRACE_BUFFER) == request_ids
    else:
        assert len(_TRACE_BUFFER) == 0


@pytest.mark.parametrize("enable_mlflow_tracing", [True, False])
def test_enable_mlflow_tracing_switch_in_serving_client(monkeypatch, enable_mlflow_tracing):
    monkeypatch.setenv("ENABLE_MLFLOW_TRACING", str(enable_mlflow_tracing).lower())
    monkeypatch.setenv("IS_IN_DB_MODEL_SERVING_ENV", "true")

    def foo():
        return bar()

    @mlflow.trace
    def bar():
        return 1

    request_ids = ["123", "234"]
    with mock.patch(
        "mlflow.tracing.processor.inference_table.maybe_get_request_id", side_effect=request_ids
    ):
        span = start_span_no_context("root")
        foo()
        if enable_mlflow_tracing:
            span.end()

    if enable_mlflow_tracing:
        assert sorted(_TRACE_BUFFER) == request_ids
    else:
        assert len(_TRACE_BUFFER) == 0


def test_sampling_ratio(monkeypatch):
    @mlflow.trace
    def test_function():
        return "test"

    # Test with 100% sampling (default)
    for _ in range(10):
        test_function()

    traces = get_traces()
    assert len(traces) == 10
    purge_traces()

    # Test with 0% sampling
    monkeypatch.setenv(MLFLOW_TRACE_SAMPLING_RATIO.name, "0.0")
    mlflow.tracing.reset()

    for _ in range(10):
        test_function()

    traces = get_traces()
    assert len(traces) == 0
    purge_traces()

    # With 50% sampling and 100 runs, we expect around 50 traces
    # but due to randomness, we check for a reasonable range
    monkeypatch.setenv(MLFLOW_TRACE_SAMPLING_RATIO.name, "0.5")
    mlflow.tracing.reset()

    for _ in range(100):
        test_function()

    traces = get_traces()
    assert 30 <= len(traces) <= 70, (
        f"Expected around 50 traces with 0.5 sampling, got {len(traces)}"
    )


def test_otlp_exclusive_vs_dual_export_with_no_set_location(monkeypatch):
    from mlflow.environment_variables import MLFLOW_TRACE_ENABLE_OTLP_DUAL_EXPORT
    from mlflow.tracing.processor.otel import OtelSpanProcessor
    from mlflow.tracing.provider import _get_tracer

    # Test 1: OTLP exclusive mode (dual export = false, default)
    monkeypatch.setenv(MLFLOW_TRACE_ENABLE_OTLP_DUAL_EXPORT.name, "false")
    with (
        mock.patch("mlflow.tracing.provider.should_use_otlp_exporter", return_value=True),
        mock.patch("mlflow.tracing.provider.get_otlp_exporter") as mock_get_exporter,
    ):
        mock_get_exporter.return_value = mock.MagicMock()

        mlflow.tracing.reset()
        tracer = _get_tracer("test")

        processors = tracer.span_processor._span_processors

        # Should have only OTLP processor as primary
        assert len(processors) == 1
        assert isinstance(processors[0], OtelSpanProcessor)

    # Test 2: Dual export mode (both MLflow and OTLP)
    monkeypatch.setenv(MLFLOW_TRACE_ENABLE_OTLP_DUAL_EXPORT.name, "true")
    with (
        mock.patch("mlflow.tracing.provider.should_use_otlp_exporter", return_value=True),
        mock.patch("mlflow.tracing.provider.get_otlp_exporter") as mock_get_exporter,
    ):
        mock_get_exporter.return_value = mock.MagicMock()

        mlflow.tracing.reset()
        tracer = _get_tracer("test")

        processors = tracer.span_processor._span_processors

        # Should have both processors
        assert len(processors) == 2
        assert isinstance(processors[0], OtelSpanProcessor)
        assert isinstance(processors[1], MlflowV3SpanProcessor)


@skip_when_testing_trace_sdk
@pytest.mark.parametrize("dual_export", [False, True])
def test_metrics_export_with_otlp_trace_export(monkeypatch, dual_export):
    monkeypatch.setenv("OTEL_EXPORTER_OTLP_TRACES_ENDPOINT", "http://localhost:4317")
    monkeypatch.setenv("OTEL_EXPORTER_OTLP_METRICS_ENDPOINT", "http://localhost:9090")

    if dual_export:
        monkeypatch.setenv(MLFLOW_TRACE_ENABLE_OTLP_DUAL_EXPORT.name, "true")

    mlflow.tracing.reset()
    tracer = _get_tracer("test")

    if dual_export:
        processors = tracer.span_processor._span_processors
        assert len(processors) == 2
        assert isinstance(processors[0], OtelSpanProcessor)
        assert isinstance(processors[1], MlflowV3SpanProcessor)

        # In dual export, MLflow processor exports metrics, OTLP doesn't
        assert processors[0]._export_metrics is False
        assert processors[1]._export_metrics is True
    else:
        processors = tracer.span_processor._span_processors
        assert len(processors) == 1
        assert isinstance(processors[0], OtelSpanProcessor)
        assert processors[0]._export_metrics is True


@skip_when_testing_trace_sdk
def test_metrics_export_without_otlp_trace_export(monkeypatch):
    monkeypatch.setenv("OTEL_EXPORTER_OTLP_METRICS_ENDPOINT", "http://localhost:9090")

    # No OTLP tracing endpoints set
    monkeypatch.delenv("OTEL_EXPORTER_OTLP_TRACES_ENDPOINT", raising=False)
    monkeypatch.delenv("OTEL_EXPORTER_OTLP_ENDPOINT", raising=False)

    mlflow.tracing.reset()
    tracer = _get_tracer("test")

    processors = tracer.span_processor._span_processors
    assert len(processors) == 1
    assert isinstance(processors[0], MlflowV3SpanProcessor)
    assert processors[0]._export_metrics is True


def test_otel_resource_attributes(monkeypatch):
    def resource_attributes(tracer):
        # opentelemetry-sdk 1.43.0+ auto-injects a random `service.instance.id` when it builds
        # the resource from env vars. It is not deterministic, so drop it before comparing.
        attributes = dict(tracer.resource.attributes)
        attributes.pop("service.instance.id", None)
        return attributes

    tracer = _get_tracer("test")
    # By default, only MLflow's SDK attributes are set on an empty resource
    assert resource_attributes(tracer) == {
        "telemetry.sdk.language": "python",
        "telemetry.sdk.name": "mlflow",
        "telemetry.sdk.version": mlflow.__version__,
    }

    mlflow.tracing.reset()
    # When otel attributes are set explicitly, they are merged into the resource
    # alongside MLflow's SDK attributes.
    monkeypatch.setenv("OTEL_RESOURCE_ATTRIBUTES", "favorite.fruit=apple,color=red")
    tracer = _get_tracer("test")
    assert resource_attributes(tracer) == {
        "favorite.fruit": "apple",
        "color": "red",
        "telemetry.sdk.language": "python",
        "telemetry.sdk.name": "mlflow",
        "telemetry.sdk.version": mlflow.__version__,
        "service.name": "unknown_service",
    }

    # Service name should be propagated from the env var
    mlflow.tracing.reset()
    monkeypatch.setenv("OTEL_SERVICE_NAME", "test-service")
    monkeypatch.delenv("OTEL_RESOURCE_ATTRIBUTES", raising=False)
    tracer = _get_tracer("test")
    assert resource_attributes(tracer) == {
        "service.name": "test-service",
        "telemetry.sdk.language": "python",
        "telemetry.sdk.name": "mlflow",
        "telemetry.sdk.version": mlflow.__version__,
    }

    # Invalid env var should be ignored and does not block the tracer provider initialization
    mlflow.tracing.reset()
    monkeypatch.setenv("OTEL_RESOURCE_ATTRIBUTES", "invalid")
    monkeypatch.delenv("OTEL_SERVICE_NAME", raising=False)
    tracer = _get_tracer("test")
    assert resource_attributes(tracer) == {
        "service.name": "unknown_service",
        "telemetry.sdk.language": "python",
        "telemetry.sdk.name": "mlflow",
        "telemetry.sdk.version": mlflow.__version__,
    }


def test_isolated_random_id_generator_not_affected_by_random_seed(monkeypatch):
    monkeypatch.setenv("MLFLOW_TRACE_USE_ISOLATED_RANDOM_ID_GENERATOR", "true")
    mlflow.tracing.reset()
    _initialize_tracer_provider()

    rng_state = random.getstate()
    try:
        random.seed(42)
        span1 = start_span_in_context("test1")
        trace_id_1 = span1.get_span_context().trace_id
        span_id_1 = span1.get_span_context().span_id
        span1.end()

        # Re-seeding with the same value would make RandomIdGenerator replay the exact same
        # ID sequence. _IsolatedRandomIdGenerator must be immune to this.
        random.seed(42)
        span2 = start_span_in_context("test2")
        trace_id_2 = span2.get_span_context().trace_id
        span_id_2 = span2.get_span_context().span_id
        span2.end()
    finally:
        random.setstate(rng_state)

    assert trace_id_1 != trace_id_2
    assert span_id_1 != span_id_2


def test_tracer_provider_uses_isolated_random_id_generator_when_env_var_set(monkeypatch):
    # Ensure env var is unset so the default id generator is used
    monkeypatch.delenv("MLFLOW_TRACE_USE_ISOLATED_RANDOM_ID_GENERATOR", raising=False)

    # Default: OTel's RandomIdGenerator is used
    mlflow.tracing.reset()
    _initialize_tracer_provider()
    tracer_provider = _provider_wrapper.get()
    assert isinstance(tracer_provider.id_generator, RandomIdGenerator)

    # Opt-in: _IsolatedRandomIdGenerator is used when the env var is set
    monkeypatch.setenv("MLFLOW_TRACE_USE_ISOLATED_RANDOM_ID_GENERATOR", "true")
    mlflow.tracing.reset()
    _initialize_tracer_provider()
    tracer_provider = _provider_wrapper.get()
    assert isinstance(tracer_provider.id_generator, _IsolatedRandomIdGenerator)


def test_set_destination_from_env_var_databricks_uc_with_table_prefix_rejected(monkeypatch):
    monkeypatch.setenv("MLFLOW_TRACING_DESTINATION", "catalog.schema.prefix")

    from mlflow.tracing.provider import _MLFLOW_TRACE_USER_DESTINATION

    with pytest.raises(
        MlflowException,
        match=r"Unity Catalog table-prefix destinations "
        r"\(<catalog_name>\.<schema_name>\.<table_prefix>\) are not supported in "
        r"MLFLOW_TRACING_DESTINATION.*Use `mlflow\.set_experiment",
    ):
        _MLFLOW_TRACE_USER_DESTINATION.get()


def test_set_destination_from_env_var_databricks_uc_with_table_prefix_rejected_on_init(
    monkeypatch,
):
    mlflow.tracing.reset()
    monkeypatch.setenv("MLFLOW_TRACING_DESTINATION", "catalog.schema.prefix")

    with pytest.raises(
        MlflowException,
        match=r"Unity Catalog table-prefix destinations "
        r"\(<catalog_name>\.<schema_name>\.<table_prefix>\) are not supported in "
        r"MLFLOW_TRACING_DESTINATION.*Use `mlflow\.set_experiment",
    ):
        _get_tracer("test")


def test_destination_resolution_precedence(monkeypatch):
    from mlflow.tracing.provider import _MLFLOW_TRACE_USER_DESTINATION

    _MLFLOW_TRACE_USER_DESTINATION.reset()
    monkeypatch.setenv("MLFLOW_TRACING_DESTINATION", "catalog.schema")

    # Env fallback is lowest priority.
    destination = _MLFLOW_TRACE_USER_DESTINATION.get()
    assert isinstance(destination, UCSchemaLocation)

    # Global slot wins over env.
    global_destination = UnityCatalog("catalog", "schema", table_prefix="global")
    _MLFLOW_TRACE_USER_DESTINATION.set(global_destination)
    assert _MLFLOW_TRACE_USER_DESTINATION.get().table_prefix == "global"

    # Context-local wins over global.
    local_destination = UnityCatalog("catalog", "schema", table_prefix="local")
    _MLFLOW_TRACE_USER_DESTINATION.set(local_destination, context_local=True)
    assert _MLFLOW_TRACE_USER_DESTINATION.get().table_prefix == "local"
    _MLFLOW_TRACE_USER_DESTINATION.reset()


def _experiment(tags=None):
    from mlflow.entities import Experiment
    from mlflow.entities.experiment_tag import ExperimentTag

    tag_entities = [ExperimentTag(k, v) for k, v in (tags or {}).items()] if tags else []
    return Experiment(
        experiment_id="123",
        name="test",
        artifact_location="file:/tmp",
        lifecycle_stage="active",
        tags=tag_entities,
    )


def test_resolve_uc_location_from_experiment_tag():
    from mlflow.tracing.provider import _resolve_experiment_uc_location

    mlflow.tracing.reset()

    with (
        mock.patch(
            "mlflow.tracing.provider.mlflow.get_tracking_uri",
            return_value="databricks",
        ),
        mock.patch("mlflow.tracking.fluent._get_experiment_id", return_value="123"),
        mock.patch("mlflow.tracking._tracking_service.utils._get_store") as mock_store_fn,
    ):
        mock_store_fn.return_value.get_experiment.return_value = _experiment(
            tags={MLFLOW_EXPERIMENT_DATABRICKS_TRACE_DESTINATION_PATH: "cat.sch.pfx"}
        )

        result = _resolve_experiment_uc_location()

        assert result == UnityCatalog("cat", "sch", table_prefix="pfx")

    mlflow.tracing.reset()


def test_resolve_uc_location_includes_table_names_from_tags():
    from mlflow.tracing.provider import _resolve_experiment_uc_location

    mlflow.tracing.reset()

    with (
        mock.patch(
            "mlflow.tracing.provider.mlflow.get_tracking_uri",
            return_value="databricks",
        ),
        mock.patch("mlflow.tracking.fluent._get_experiment_id", return_value="123"),
        mock.patch("mlflow.tracking._tracking_service.utils._get_store") as mock_store_fn,
    ):
        mock_store_fn.return_value.get_experiment.return_value = _experiment(
            tags={
                MLFLOW_EXPERIMENT_DATABRICKS_TRACE_DESTINATION_PATH: "cat.sch.pfx",
                MLFLOW_EXPERIMENT_DATABRICKS_TRACE_SPAN_STORAGE_TABLE: "cat.sch.pfx_otel_spans",
                MLFLOW_EXPERIMENT_DATABRICKS_TRACE_LOG_STORAGE_TABLE: "cat.sch.pfx_otel_logs",
            }
        )

        result = _resolve_experiment_uc_location()

        assert result == UnityCatalog("cat", "sch", table_prefix="pfx")
        assert result._otel_spans_table_name == "cat.sch.pfx_otel_spans"
        assert result._otel_logs_table_name == "cat.sch.pfx_otel_logs"

    mlflow.tracing.reset()


def test_resolve_uc_location_returns_none_for_2_part_path():
    from mlflow.tracing.provider import _resolve_experiment_uc_location

    mlflow.tracing.reset()

    with (
        mock.patch(
            "mlflow.tracing.provider.mlflow.get_tracking_uri",
            return_value="databricks",
        ),
        mock.patch("mlflow.tracking.fluent._get_experiment_id", return_value="123"),
        mock.patch("mlflow.tracking._tracking_service.utils._get_store") as mock_store_fn,
    ):
        mock_store_fn.return_value.get_experiment.return_value = _experiment(
            tags={MLFLOW_EXPERIMENT_DATABRICKS_TRACE_DESTINATION_PATH: "cat.sch"}
        )

        assert _resolve_experiment_uc_location() is None

    mlflow.tracing.reset()


def test_resolve_uc_location_returns_none_when_no_tag():
    from mlflow.tracing.provider import _resolve_experiment_uc_location

    mlflow.tracing.reset()

    with (
        mock.patch(
            "mlflow.tracing.provider.mlflow.get_tracking_uri",
            return_value="databricks",
        ),
        mock.patch("mlflow.tracking.fluent._get_experiment_id", return_value="123"),
        mock.patch("mlflow.tracking._tracking_service.utils._get_store") as mock_store_fn,
    ):
        mock_store_fn.return_value.get_experiment.return_value = _experiment()

        assert _resolve_experiment_uc_location() is None

    mlflow.tracing.reset()


def test_resolve_uc_location_returns_none_for_non_databricks():
    from mlflow.tracing.provider import _resolve_experiment_uc_location

    mlflow.tracing.reset()

    with mock.patch(
        "mlflow.tracing.provider.mlflow.get_tracking_uri",
        return_value="http://local",
    ):
        assert _resolve_experiment_uc_location() is None

    mlflow.tracing.reset()


def test_get_tracer_does_not_fail_when_experiment_id_resolution_fails():
    mlflow.tracing.reset()

    with (
        mock.patch(
            "mlflow.tracing.provider.mlflow.get_tracking_uri",
            return_value="databricks",
        ),
        mock.patch("mlflow.tracking.fluent._get_experiment_id", side_effect=RuntimeError("boom")),
    ):
        tracer = _get_tracer("test")

    assert tracer is not None
    mlflow.tracing.reset()