"""Tests for OTelMiddleware."""

from __future__ import annotations

from typing import Any
from unittest.mock import MagicMock, patch

import pytest

from agentbyte.context import AgentContext
from agentbyte.messages import Usage
from agentbyte.middleware.base import MiddlewareContext


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _make_context(
    operation: str = "model_call",
    agent_name: str = "test_agent",
    data: Any = None,
    metadata: dict | None = None,
) -> MiddlewareContext:
    """Build a minimal MiddlewareContext for testing."""
    return MiddlewareContext(
        operation=operation,
        agent_name=agent_name,
        agent_context=AgentContext(),
        data=data or {},
        metadata=metadata or {},
    )


def _make_usage(
    tokens_input: int = 100, tokens_output: int = 50, cost_estimate: float | None = None
) -> Usage:
    return Usage(
        tokens_input=tokens_input,
        tokens_output=tokens_output,
        cost_estimate=cost_estimate,
    )


# ---------------------------------------------------------------------------
# 1. Disabled by default
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_disabled_by_default(monkeypatch: pytest.MonkeyPatch) -> None:
    """Without AGENTBYTE_ENABLE_OTEL, middleware is a no-op pass-through."""
    monkeypatch.delenv("AGENTBYTE_ENABLE_OTEL", raising=False)

    from agentbyte.middleware.otel import OTelMiddleware

    mw = OTelMiddleware()
    assert mw._enabled is False

    ctx = _make_context()
    result_ctx = await mw.process_request(ctx)
    assert result_ctx is ctx

    result = object()
    assert await mw.process_response(ctx, result) is result


# ---------------------------------------------------------------------------
# 2. Enabled creates a span
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_enabled_creates_span() -> None:
    """With a live tracer mock, start_span is called for a model_call operation."""
    mock_span = MagicMock()
    mock_tracer = MagicMock()
    mock_tracer.start_span.return_value = mock_span

    from agentbyte.middleware.otel import OTelMiddleware

    # Bypass __init__ and wire mocks directly
    mw = OTelMiddleware.__new__(OTelMiddleware)
    mw._enabled = True
    mw._capture_content = False
    mw._tracer = mock_tracer
    mw._token_histogram = None
    mw._duration_histogram = None

    ctx = _make_context(operation="model_call", metadata={"model": "gpt-4.1-mini"})

    with (
        patch("agentbyte.middleware.otel.otel_context") as mock_otel_ctx,
        patch("agentbyte.middleware.otel.trace"),
    ):
        mock_otel_ctx.attach.return_value = "token"
        await mw.process_request(ctx)

    mock_tracer.start_span.assert_called_once()


# ---------------------------------------------------------------------------
# 3. Span name for model_call
# ---------------------------------------------------------------------------


def test_span_name_model_call() -> None:
    """model_call produces 'chat <model-name>' span name."""
    from agentbyte.middleware.otel import OTelMiddleware

    mw = OTelMiddleware.__new__(OTelMiddleware)
    ctx = _make_context(operation="model_call", metadata={"model": "gpt-4.1-mini"})
    assert mw._build_span_name(ctx) == "chat gpt-4.1-mini"


# ---------------------------------------------------------------------------
# 4. Span name for tool_call
# ---------------------------------------------------------------------------


def test_span_name_tool_call() -> None:
    """tool_call produces 'tool <tool-name>' span name."""
    from agentbyte.middleware.otel import OTelMiddleware

    mw = OTelMiddleware.__new__(OTelMiddleware)
    ctx = _make_context(
        operation="tool_call",
        data={"tool_name": "get_weather"},
    )
    assert mw._build_span_name(ctx) == "tool get_weather"


# ---------------------------------------------------------------------------
# 5. process_response records tokens_input / tokens_output (agentbyte field names)
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_process_response_records_tokens() -> None:
    """Token span attributes use agentbyte Usage field names (tokens_input/tokens_output)."""
    mock_span = MagicMock()
    mock_histogram = MagicMock()

    from agentbyte.middleware.otel import OTelMiddleware

    mw = OTelMiddleware.__new__(OTelMiddleware)
    mw._enabled = True
    mw._capture_content = False
    mw._token_histogram = mock_histogram
    mw._duration_histogram = None

    usage = _make_usage(tokens_input=120, tokens_output=40)
    result = MagicMock()
    result.usage = usage

    ctx = _make_context(
        operation="model_call",
        metadata={"_otel_span": mock_span, "_otel_start": 0.0, "_otel_token": "tok"},
    )

    with (
        patch("agentbyte.middleware.otel.otel_context"),
        patch("agentbyte.middleware.otel.StatusCode"),
        patch("agentbyte.middleware.otel.Status"),
    ):
        await mw.process_response(ctx, result)

    mock_span.set_attribute.assert_any_call("gen_ai.usage.input_tokens", 120)
    mock_span.set_attribute.assert_any_call("gen_ai.usage.output_tokens", 40)
    mock_span.set_attribute.assert_any_call("gen_ai.usage.total_tokens", 160)

    calls = [str(c) for c in mock_histogram.record.call_args_list]
    assert any("120" in c for c in calls), "input token histogram not recorded"
    assert any("40" in c for c in calls), "output token histogram not recorded"


@pytest.mark.asyncio
async def test_process_response_records_model_finish_reason() -> None:
    """Model finish_reason is captured on model_call spans when available."""
    mock_span = MagicMock()

    from agentbyte.middleware.otel import OTelMiddleware

    mw = OTelMiddleware.__new__(OTelMiddleware)
    mw._enabled = True
    mw._capture_content = False
    mw._token_histogram = None
    mw._duration_histogram = None

    usage = _make_usage(tokens_input=120, tokens_output=40, cost_estimate=0.00123)
    result = MagicMock()
    result.usage = usage
    result.finish_reason = "stop"

    ctx = _make_context(
        operation="model_call",
        metadata={"_otel_span": mock_span, "_otel_start": 0.0, "_otel_token": "tok"},
    )

    with (
        patch("agentbyte.middleware.otel.otel_context"),
        patch("agentbyte.middleware.otel.StatusCode"),
        patch("agentbyte.middleware.otel.Status"),
    ):
        await mw.process_response(ctx, result)

    mock_span.set_attribute.assert_any_call("gen_ai.response.finish_reason", "stop")


# ---------------------------------------------------------------------------
# 5b. process_response records cost_estimate when available
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_process_response_records_cost_estimate() -> None:
    """Cost estimate span attribute is set when Usage.cost_estimate is present."""
    mock_span = MagicMock()

    from agentbyte.middleware.otel import OTelMiddleware

    mw = OTelMiddleware.__new__(OTelMiddleware)
    mw._enabled = True
    mw._capture_content = False
    mw._token_histogram = None
    mw._duration_histogram = None

    usage = _make_usage(tokens_input=120, tokens_output=40, cost_estimate=0.00123)
    result = MagicMock()
    result.usage = usage

    ctx = _make_context(
        operation="model_call",
        metadata={"_otel_span": mock_span, "_otel_start": 0.0, "_otel_token": "tok"},
    )

    with (
        patch("agentbyte.middleware.otel.otel_context"),
        patch("agentbyte.middleware.otel.StatusCode"),
        patch("agentbyte.middleware.otel.Status"),
    ):
        await mw.process_response(ctx, result)

    mock_span.set_attribute.assert_any_call("gen_ai.usage.cost_estimate_usd", 0.00123)


@pytest.mark.asyncio
async def test_process_response_cost_estimate_none_not_recorded() -> None:
    """Cost estimate attribute is NOT set when Usage.cost_estimate is None."""
    mock_span = MagicMock()

    from agentbyte.middleware.otel import OTelMiddleware

    mw = OTelMiddleware.__new__(OTelMiddleware)
    mw._enabled = True
    mw._capture_content = False
    mw._token_histogram = None
    mw._duration_histogram = None

    usage = _make_usage(tokens_input=120, tokens_output=40, cost_estimate=None)
    result = MagicMock()
    result.usage = usage

    ctx = _make_context(
        operation="model_call",
        metadata={"_otel_span": mock_span, "_otel_start": 0.0, "_otel_token": "tok"},
    )

    with (
        patch("agentbyte.middleware.otel.otel_context"),
        patch("agentbyte.middleware.otel.StatusCode"),
        patch("agentbyte.middleware.otel.Status"),
    ):
        await mw.process_response(ctx, result)

    # Cost attribute should NOT be set if cost_estimate is None
    cost_calls = [
        c for c in mock_span.set_attribute.call_args_list
        if "cost_estimate" in str(c)
    ]
    assert len(cost_calls) == 0, f"cost_estimate should not be set when None, got {cost_calls}"


# ---------------------------------------------------------------------------
# 6. process_response ends span with OK status
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_process_response_ends_span_ok() -> None:
    """Successful response sets StatusCode.OK and calls span.end()."""
    mock_span = MagicMock()

    from agentbyte.middleware.otel import OTelMiddleware

    mw = OTelMiddleware.__new__(OTelMiddleware)
    mw._enabled = True
    mw._capture_content = False
    mw._token_histogram = None
    mw._duration_histogram = None

    ctx = _make_context(
        operation="tool_call",
        metadata={"_otel_span": mock_span, "_otel_start": 0.0, "_otel_token": "tok"},
    )
    result = MagicMock()
    result.success = True

    with (
        patch("agentbyte.middleware.otel.otel_context"),
        patch("agentbyte.middleware.otel.StatusCode"),
        patch("agentbyte.middleware.otel.Status"),
    ):
        await mw.process_response(ctx, result)

    mock_span.set_status.assert_called_once()
    mock_span.end.assert_called_once()


# ---------------------------------------------------------------------------
# 7. process_error sets ERROR status and re-raises
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_process_error_sets_error_status_and_reraises() -> None:
    """Error path sets StatusCode.ERROR and always re-raises the original exception."""
    mock_span = MagicMock()

    from agentbyte.middleware.otel import OTelMiddleware

    mw = OTelMiddleware.__new__(OTelMiddleware)
    mw._enabled = True

    ctx = _make_context(
        metadata={"_otel_span": mock_span, "_otel_start": 0.0, "_otel_token": "tok"},
    )
    error = ValueError("something broke")

    with (
        patch("agentbyte.middleware.otel.otel_context"),
        patch("agentbyte.middleware.otel.StatusCode"),
        patch("agentbyte.middleware.otel.Status"),
    ):
        with pytest.raises(ValueError, match="something broke"):
            await mw.process_error(ctx, error)

    mock_span.set_status.assert_called_once()
    mock_span.end.assert_called_once()


# ---------------------------------------------------------------------------
# 8. Content NOT captured by default
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_content_not_captured_by_default() -> None:
    """gen_ai.input.messages is NOT set when _capture_content is False."""
    mock_span = MagicMock()

    from agentbyte.middleware.otel import OTelMiddleware

    mw = OTelMiddleware.__new__(OTelMiddleware)
    mw._enabled = True
    mw._capture_content = False
    mw._token_histogram = None
    mw._duration_histogram = None

    ctx = _make_context(
        operation="model_call",
        metadata={"_otel_span": mock_span, "_otel_start": 0.0, "_otel_token": "tok"},
    )
    result = MagicMock()
    result.usage = _make_usage()

    with (
        patch("agentbyte.middleware.otel.otel_context"),
        patch("agentbyte.middleware.otel.StatusCode"),
        patch("agentbyte.middleware.otel.Status"),
    ):
        await mw.process_response(ctx, result)

    attribute_names = [call.args[0] for call in mock_span.set_attribute.call_args_list]
    assert "gen_ai.input.messages" not in attribute_names
    assert "gen_ai.output.messages" not in attribute_names


# ---------------------------------------------------------------------------
# 9. Content captured when opt-in
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_content_captured_opt_in(monkeypatch: pytest.MonkeyPatch) -> None:
    """gen_ai.input.messages is set on the span when _capture_content is True."""
    mock_span = MagicMock()

    from agentbyte.messages import UserMessage
    from agentbyte.middleware.otel import OTelMiddleware

    mw = OTelMiddleware.__new__(OTelMiddleware)
    mw._enabled = True
    mw._capture_content = True
    mw._token_histogram = None
    mw._duration_histogram = None
    mw._tracer = MagicMock()
    mw._tracer.start_span.return_value = mock_span

    messages = [UserMessage(content="hello", source="user")]

    ctx = _make_context(
        operation="model_call",
        data={"messages": messages},
        metadata={"model": "gpt-4.1-mini", "_otel_span": mock_span, "_otel_start": 0.0, "_otel_token": "tok"},
    )

    with (
        patch("agentbyte.middleware.otel.otel_context") as mock_otel_ctx,
        patch("agentbyte.middleware.otel.trace"),
        patch("agentbyte.middleware.otel.StatusCode"),
        patch("agentbyte.middleware.otel.Status"),
    ):
        mock_otel_ctx.attach.return_value = "token"
        await mw.process_request(ctx)

    attribute_names = [call.args[0] for call in mock_span.set_attribute.call_args_list]
    assert "gen_ai.input.messages" in attribute_names


# ---------------------------------------------------------------------------
# 10. auto_instrument prepends OTelMiddleware to Agent
# ---------------------------------------------------------------------------


def test_auto_instrument_prepends_middleware(monkeypatch: pytest.MonkeyPatch) -> None:
    """auto_instrument() patches Agent.__init__ so OTelMiddleware is first in chain."""
    monkeypatch.setenv("AGENTBYTE_ENABLE_OTEL", "true")

    from agentbyte.agents.agent import Agent
    from agentbyte.middleware.otel import auto_instrument

    original_init = getattr(Agent.__init__, "__agentbyte_original_init__", Agent.__init__)
    Agent.__init__ = original_init
    try:
        auto_instrument()
        assert Agent.__init__ is not original_init
    finally:
        Agent.__init__ = original_init


def test_auto_instrument_is_idempotent(monkeypatch: pytest.MonkeyPatch) -> None:
    """Repeated auto_instrument() calls should not stack-wrap Agent.__init__."""
    monkeypatch.setenv("AGENTBYTE_ENABLE_OTEL", "true")

    from agentbyte.agents.agent import Agent
    from agentbyte.middleware.otel import auto_instrument

    original_init = getattr(Agent.__init__, "__agentbyte_original_init__", Agent.__init__)
    Agent.__init__ = original_init

    try:
        auto_instrument()
        once_init = Agent.__init__

        auto_instrument()
        twice_init = Agent.__init__

        assert once_init is twice_init
        assert getattr(Agent.__init__, "__agentbyte_otel_instrumented__", False) is True
    finally:
        Agent.__init__ = original_init


def test_auto_instrument_does_not_patch_run(monkeypatch: pytest.MonkeyPatch) -> None:
    """auto_instrument() patches only Agent.__init__, not Agent.run."""
    monkeypatch.setenv("AGENTBYTE_ENABLE_OTEL", "true")

    from agentbyte.agents.agent import Agent
    from agentbyte.middleware.otel import auto_instrument

    original_init = getattr(Agent.__init__, "__agentbyte_original_init__", Agent.__init__)
    Agent.__init__ = original_init
    original_run = Agent.run

    try:
        auto_instrument()
        assert Agent.run is original_run, "Agent.run must not be monkey-patched"
    finally:
        Agent.__init__ = original_init


# ---------------------------------------------------------------------------
# 11. Graceful no-op when OTEL_AVAILABLE is False
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_graceful_degradation_no_sdk(monkeypatch: pytest.MonkeyPatch) -> None:
    """Middleware is a safe no-op when OTel SDK is not installed."""
    monkeypatch.setenv("AGENTBYTE_ENABLE_OTEL", "true")

    with patch("agentbyte.middleware.otel.OTEL_AVAILABLE", False):
        from agentbyte.middleware.otel import OTelMiddleware

        mw = OTelMiddleware()
        assert mw._enabled is False

    ctx = _make_context()
    # Should not raise
    result_ctx = await mw.process_request(ctx)
    assert result_ctx is ctx

    sentinel = object()
    assert await mw.process_response(ctx, sentinel) is sentinel


def test_setup_telemetry_is_process_singleton(monkeypatch: pytest.MonkeyPatch) -> None:
    """Telemetry provider setup should happen once even with multiple middleware instances."""
    monkeypatch.setenv("OTEL_EXPORTER_OTLP_ENDPOINT", "http://localhost:4318")
    monkeypatch.setenv("OTEL_SERVICE_NAME", "agentbyte-test")
    monkeypatch.delenv("OTEL_METRICS_ENABLED", raising=False)

    import agentbyte.middleware.otel as otel

    monkeypatch.setattr(otel, "_TELEMETRY_TRACER", None)
    monkeypatch.setattr(otel, "_TELEMETRY_METER", None)
    monkeypatch.setattr(otel, "_TELEMETRY_INITIALIZED", False)

    mock_trace_provider = MagicMock()
    mock_tracer = MagicMock()

    with (
        patch.object(otel, "OTEL_AVAILABLE", True),
        patch.object(otel, "Resource") as mock_resource,
        patch.object(otel, "TracerProvider", return_value=mock_trace_provider),
        patch.object(otel, "OTLPSpanExporter"),
        patch.object(otel, "BatchSpanProcessor", return_value=MagicMock()),
        patch.object(otel, "trace") as mock_trace,
    ):
        mock_resource.create.return_value = MagicMock()
        mock_trace.get_tracer.return_value = mock_tracer

        tracer1, meter1 = otel._setup_telemetry()
        tracer2, meter2 = otel._setup_telemetry()

    assert tracer1 is mock_tracer
    assert tracer2 is mock_tracer
    assert meter1 is None
    assert meter2 is None
    mock_trace.set_tracer_provider.assert_called_once_with(mock_trace_provider)
