from __future__ import annotations

import logging
from uuid import uuid4

import agents.tracing.traces as trace_module
from agents.tracing import TracingConfig, set_tracing_disabled, trace
from agents.tracing.context import create_trace_for_run
from agents.tracing.scope import Scope
from agents.tracing.traces import (
    NoOpTrace,
    ReattachedTrace,
    TraceImpl,
    TraceState,
    _started_trace_ids,
    _started_trace_ids_lock,
)


def _new_trace_id() -> str:
    return f"trace_{uuid4().hex}"


def _clear_started_trace_ids() -> None:
    with _started_trace_ids_lock:
        _started_trace_ids.clear()


def _mark_trace_as_started(
    *,
    workflow_name: str = "workflow",
    group_id: str | None = "group-1",
    metadata: dict[str, str] | None = None,
    tracing_api_key: str | None = None,
) -> TraceState:
    metadata = metadata or {"key": "value"}
    trace_id = _new_trace_id()
    Scope.set_current_trace(None)
    set_tracing_disabled(False)

    original = trace(
        workflow_name=workflow_name,
        trace_id=trace_id,
        group_id=group_id,
        metadata=metadata,
        tracing={"api_key": tracing_api_key} if tracing_api_key is not None else None,
    )
    assert isinstance(original, TraceImpl)
    original.start()
    original.finish()

    trace_state = TraceState.from_trace(original)
    assert trace_state is not None
    return trace_state


def test_create_trace_for_run_reattaches_matching_started_trace() -> None:
    trace_state = _mark_trace_as_started(tracing_api_key="trace-key")

    created = create_trace_for_run(
        workflow_name="workflow",
        trace_id=trace_state.trace_id,
        group_id=trace_state.group_id,
        metadata=dict(trace_state.metadata or {}),
        tracing={"api_key": "trace-key"},
        disabled=False,
        trace_state=trace_state,
        reattach_resumed_trace=True,
    )

    assert isinstance(created, ReattachedTrace)
    assert created.trace_id == trace_state.trace_id


def test_create_trace_for_run_does_not_reattach_after_trace_state_reload() -> None:
    trace_state = _mark_trace_as_started()
    _clear_started_trace_ids()

    created = create_trace_for_run(
        workflow_name="workflow",
        trace_id=trace_state.trace_id,
        group_id=trace_state.group_id,
        metadata=dict(trace_state.metadata or {}),
        tracing=None,
        disabled=False,
        trace_state=trace_state,
        reattach_resumed_trace=True,
    )

    assert isinstance(created, TraceImpl)
    assert not isinstance(created, ReattachedTrace)


def test_create_trace_for_run_does_not_reattach_when_trace_id_differs() -> None:
    trace_state = _mark_trace_as_started()

    created = create_trace_for_run(
        workflow_name="workflow",
        trace_id=_new_trace_id(),
        group_id=trace_state.group_id,
        metadata=dict(trace_state.metadata or {}),
        tracing=None,
        disabled=False,
        trace_state=trace_state,
        reattach_resumed_trace=True,
    )

    assert isinstance(created, TraceImpl)
    assert not isinstance(created, ReattachedTrace)


def test_create_trace_for_run_reattaches_stripped_trace_key_with_matching_resume_key() -> None:
    trace_state = _mark_trace_as_started(tracing_api_key="trace-key")
    stripped_trace_state = TraceState.from_json(trace_state.to_json())
    assert stripped_trace_state is not None
    assert stripped_trace_state.tracing_api_key is None
    assert stripped_trace_state.tracing_api_key_hash == trace_state.tracing_api_key_hash

    created = create_trace_for_run(
        workflow_name="workflow",
        trace_id=stripped_trace_state.trace_id,
        group_id=stripped_trace_state.group_id,
        metadata=dict(stripped_trace_state.metadata or {}),
        tracing={"api_key": "trace-key"},
        disabled=False,
        trace_state=stripped_trace_state,
        reattach_resumed_trace=True,
    )

    assert isinstance(created, ReattachedTrace)
    assert created.tracing_api_key == "trace-key"


def test_create_trace_for_run_does_not_reattach_stripped_trace_key_with_mismatch() -> None:
    trace_state = _mark_trace_as_started(tracing_api_key="trace-key")
    stripped_trace_state = TraceState.from_json(trace_state.to_json())
    assert stripped_trace_state is not None

    created = create_trace_for_run(
        workflow_name="workflow",
        trace_id=stripped_trace_state.trace_id,
        group_id=stripped_trace_state.group_id,
        metadata=dict(stripped_trace_state.metadata or {}),
        tracing={"api_key": "other-trace-key"},
        disabled=False,
        trace_state=stripped_trace_state,
        reattach_resumed_trace=True,
    )

    assert isinstance(created, TraceImpl)
    assert not isinstance(created, ReattachedTrace)


def test_create_trace_for_run_does_not_reattach_when_settings_mismatch() -> None:
    trace_state = _mark_trace_as_started(tracing_api_key="trace-key")

    mismatch_cases: list[tuple[str, str | None, dict[str, str], TracingConfig]] = [
        (
            "workflow-override",
            trace_state.group_id,
            dict(trace_state.metadata or {}),
            {"api_key": "trace-key"},
        ),
        (
            "workflow",
            "group-override",
            dict(trace_state.metadata or {}),
            {"api_key": "trace-key"},
        ),
        (
            "workflow",
            trace_state.group_id,
            {"key": "override"},
            {"api_key": "trace-key"},
        ),
        (
            "workflow",
            trace_state.group_id,
            dict(trace_state.metadata or {}),
            {"api_key": "other-trace-key"},
        ),
    ]

    for workflow_name, group_id, metadata, tracing in mismatch_cases:
        Scope.set_current_trace(None)
        created = create_trace_for_run(
            workflow_name=workflow_name,
            trace_id=trace_state.trace_id,
            group_id=group_id,
            metadata=metadata,
            tracing=tracing,
            disabled=False,
            trace_state=trace_state,
            reattach_resumed_trace=True,
        )

        assert isinstance(created, TraceImpl)
        assert not isinstance(created, ReattachedTrace)


def test_create_trace_for_run_respects_disabled_flag_for_resume() -> None:
    trace_state = _mark_trace_as_started()

    created = create_trace_for_run(
        workflow_name="workflow",
        trace_id=trace_state.trace_id,
        group_id=trace_state.group_id,
        metadata=dict(trace_state.metadata or {}),
        tracing=None,
        disabled=True,
        trace_state=trace_state,
        reattach_resumed_trace=True,
    )

    assert isinstance(created, NoOpTrace)


def test_create_trace_for_run_uses_existing_current_trace() -> None:
    trace_state = _mark_trace_as_started()
    outer_trace = trace(workflow_name="outer", trace_id=_new_trace_id())
    assert isinstance(outer_trace, TraceImpl)

    with outer_trace:
        created = create_trace_for_run(
            workflow_name="workflow",
            trace_id=trace_state.trace_id,
            group_id=trace_state.group_id,
            metadata=dict(trace_state.metadata or {}),
            tracing=None,
            disabled=False,
            trace_state=trace_state,
            reattach_resumed_trace=True,
        )

        assert created is None


def test_trace_logs_warning_when_current_trace_exists(
    caplog,
) -> None:
    Scope.set_current_trace(None)
    outer_trace = trace(workflow_name="outer", trace_id=_new_trace_id())
    assert isinstance(outer_trace, TraceImpl)

    with outer_trace:
        with caplog.at_level(logging.WARNING, logger="openai.agents"):
            inner_trace = trace(workflow_name="inner", trace_id=_new_trace_id())

    assert isinstance(inner_trace, TraceImpl)
    assert "Trace already exists" in caplog.text


def test_started_trace_id_cache_is_bounded(monkeypatch) -> None:
    _clear_started_trace_ids()
    monkeypatch.setattr(trace_module, "_MAX_STARTED_TRACE_IDS", 2)

    first = _mark_trace_as_started(metadata={"key": "first"})
    second = _mark_trace_as_started(metadata={"key": "second"})
    third = _mark_trace_as_started(metadata={"key": "third"})

    assert len(_started_trace_ids) == 2
    assert list(_started_trace_ids) == [second.trace_id, third.trace_id]
    assert first.trace_id not in _started_trace_ids