"""Tests for observability exporters and spans."""
from __future__ import annotations

import json
import gzip
from datetime import UTC, datetime

import httpx
import pytest

from logos_adk.agent import Agent
from logos_adk.observe import (
    Exporter,
    InMemoryExporter,
    Metric,
    NativeObservabilityExporter,
    NativeObservabilityStore,
    OpenTelemetryExporter,
    PostgreSQLObservabilityStore,
    RetentionPolicy,
    Span,
    get_observability_store,
    reset_observability_store,
)
from logos_adk.providers import provider_registry
from logos_adk.providers.mock import MockProvider
from logos_adk.runtime import _default_runtime_holder
from logos_adk.tools import tool
from logos_adk.types import Response, ToolCall, Usage


class RecordingExporter(Exporter):
    def __init__(self) -> None:
        self.spans: list[Span] = []
        self.metrics: list[Metric] = []
        self.flushed = 0

    async def on_span(self, span: Span) -> None:
        self.spans.append(span)

    async def on_metric(self, metric: Metric) -> None:
        self.metrics.append(metric)

    async def flush(self) -> None:
        self.flushed += 1


@pytest.fixture(autouse=True)
def _reset_runtime_and_registry():
    _default_runtime_holder.reset()
    provider_registry._entries.clear()
    reset_observability_store().reset()
    yield
    _default_runtime_holder.reset()
    provider_registry._entries.clear()
    reset_observability_store().reset()


def test_agent_builder_observe():
    exporter = RecordingExporter()
    agent = Agent("bot").observe(exporter)
    assert agent._config.observers == [exporter]


@pytest.mark.asyncio
async def test_in_memory_exporter_collects_spans_and_flushes():
    exporter = InMemoryExporter()
    provider = MockProvider(responses=["ok"])
    agent = Agent("bot").model(provider).observe(exporter)

    result = await agent.run("hi", session_id="s1")

    assert result.content == "ok"
    assert len(exporter.spans) >= 4
    assert exporter.flush_count == 1


@pytest.mark.asyncio
async def test_native_observability_exporter_records_spans_metrics_logs():
    exporter = NativeObservabilityExporter(get_observability_store())
    provider = MockProvider(responses=["ok"])
    agent = Agent("bot").model(provider).observe(exporter)

    await agent.run("hi", session_id="s1")

    summary = get_observability_store().summary()
    assert summary["totals"]["spans"] >= 1
    assert summary["totals"]["metrics"] >= 1


def test_observability_store_defaults_to_in_memory_without_dsn(monkeypatch):
    monkeypatch.delenv("LOGOS_ADK_OBSERVABILITY_DSN", raising=False)
    monkeypatch.delenv("LOGOS_ADK_DATABASE_URL", raising=False)
    monkeypatch.delenv("DATABASE_URL", raising=False)

    store = reset_observability_store()

    assert isinstance(store, NativeObservabilityStore)


def test_observability_store_prefers_postgres_with_dsn(monkeypatch):
    monkeypatch.setenv("LOGOS_ADK_OBSERVABILITY_DSN", "postgresql://obs-db/logos")

    class FakePostgresStore(PostgreSQLObservabilityStore):
        def _ensure_schema(self) -> None:
            return

    monkeypatch.setattr("logos_adk.observe.native.PostgreSQLObservabilityStore", FakePostgresStore)

    store = reset_observability_store()

    assert isinstance(store, FakePostgresStore)
    assert store.dsn == "postgresql://obs-db/logos"


def test_observability_store_applies_retention_from_env(monkeypatch):
    monkeypatch.setenv("LOGOS_ADK_OBSERVABILITY_DSN", "postgresql://obs-db/logos")
    monkeypatch.setenv("LOGOS_ADK_OBSERVABILITY_SPANS_TTL_DAYS", "3")
    monkeypatch.setenv("LOGOS_ADK_OBSERVABILITY_SPANS_MAX_ROWS", "123")

    class FakePostgresStore(PostgreSQLObservabilityStore):
        def _ensure_schema(self) -> None:
            return

    monkeypatch.setattr("logos_adk.observe.native.PostgreSQLObservabilityStore", FakePostgresStore)

    store = reset_observability_store()

    assert store.retention["spans"] == RetentionPolicy(ttl_days=3, max_rows=123)


class _FakeCursor:
    def __init__(self):
        self.executed: list[tuple[str, tuple | None]] = []

    def execute(self, query, params=None):
        self.executed.append((query, params))

    def fetchone(self):
        return None

    def fetchall(self):
        return []

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc, tb):
        return False


class _FakeConnection:
    def __init__(self, cursor):
        self.cursor_obj = cursor
        self.commits = 0

    def cursor(self):
        return self.cursor_obj

    def commit(self):
        self.commits += 1

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc, tb):
        return False


def test_postgres_observability_cleanup_applies_ttl_and_row_caps():
    cursor = _FakeCursor()
    connection = _FakeConnection(cursor)
    store = PostgreSQLObservabilityStore.__new__(PostgreSQLObservabilityStore)
    store.dsn = "postgresql://obs-db/logos"
    store.schema = "public"
    store.table_prefix = "adk_observability"
    store.rules = []
    store.retention = {
        "spans": RetentionPolicy(ttl_days=7, max_rows=100),
        "metrics": RetentionPolicy(ttl_days=14, max_rows=200),
        "logs": RetentionPolicy(ttl_days=21, max_rows=300),
        "alerts": RetentionPolicy(ttl_days=30, max_rows=400),
    }
    store.cleanup_interval_seconds = 300
    store._last_cleanup_at = None
    store._connect = lambda: connection

    PostgreSQLObservabilityStore.cleanup(store)

    delete_queries = [query for query, _ in cursor.executed if query.strip().startswith("DELETE FROM")]
    assert len(delete_queries) == 8
    assert any("adk_observability_spans" in query and "ended_at" in query for query in delete_queries)
    assert any("adk_observability_metrics" in query and "recorded_at" in query for query in delete_queries)
    assert any("adk_observability_logs" in query and "timestamp" in query for query in delete_queries)
    assert any("adk_observability_alerts" in query and "triggered_at" in query for query in delete_queries)


def test_postgres_observability_schema_creates_retention_and_dashboard_indexes():
    cursor = _FakeCursor()
    connection = _FakeConnection(cursor)
    store = PostgreSQLObservabilityStore.__new__(PostgreSQLObservabilityStore)
    store.dsn = "postgresql://obs-db/logos"
    store.schema = "public"
    store.table_prefix = "adk_observability"
    store.rules = []
    store.retention = {}
    store.cleanup_interval_seconds = 300
    store._last_cleanup_at = None
    store._connect = lambda: connection

    PostgreSQLObservabilityStore._ensure_schema(store)

    index_queries = [query for query, _ in cursor.executed if "CREATE INDEX IF NOT EXISTS" in query]
    assert any("spans_ended_at" in query for query in index_queries)
    assert any("spans_trace_id_started_at" in query for query in index_queries)
    assert any("spans_agent_status_ended_at" in query for query in index_queries)
    assert any("metrics_recorded_at" in query for query in index_queries)
    assert any("metrics_name_recorded_at" in query for query in index_queries)
    assert any("logs_timestamp" in query for query in index_queries)
    assert any("logs_level_timestamp" in query for query in index_queries)
    assert any("logs_trace_id_timestamp" in query for query in index_queries)
    assert any("alerts_triggered_at" in query for query in index_queries)
    assert any("alerts_rule_triggered_at" in query for query in index_queries)
    assert any("spans_attributes_gin" in query and "USING GIN (attributes)" in query for query in index_queries)
    assert any("spans_tenant_id" in query and "(attributes->>'tenant_id')" in query for query in index_queries)
    assert any("spans_tool_name" in query and "(attributes->>'tool_name')" in query for query in index_queries)
    assert any("spans_error_partial" in query and "WHERE status = 'error'" in query for query in index_queries)
    assert any("metrics_attributes_gin" in query and "USING GIN (attributes)" in query for query in index_queries)
    assert any("metrics_tenant_id" in query and "(attributes->>'tenant_id')" in query for query in index_queries)
    assert any("metrics_tool_name" in query and "(attributes->>'tool_name')" in query for query in index_queries)
    assert any("logs_attributes_gin" in query and "USING GIN (attributes)" in query for query in index_queries)
    assert any("logs_tenant_id" in query and "(attributes->>'tenant_id')" in query for query in index_queries)
    assert any("logs_tool_name" in query and "(attributes->>'tool_name')" in query for query in index_queries)
    assert any("logs_error_partial" in query and "WHERE level = 'error'" in query for query in index_queries)
    assert any("alerts_attributes_gin" in query and "USING GIN (attributes)" in query for query in index_queries)
    assert any("alerts_tenant_id" in query and "(attributes->>'tenant_id')" in query for query in index_queries)
    assert any("alerts_tool_name" in query and "(attributes->>'tool_name')" in query for query in index_queries)


def test_postgres_observability_query_helpers_build_attribute_filters():
    cursor = _FakeCursor()
    connection = _FakeConnection(cursor)
    store = PostgreSQLObservabilityStore.__new__(PostgreSQLObservabilityStore)
    store.dsn = "postgresql://obs-db/logos"
    store.schema = "public"
    store.table_prefix = "adk_observability"
    store.rules = []
    store.retention = {}
    store.cleanup_interval_seconds = 300
    store._last_cleanup_at = None
    store._connect = lambda: connection

    PostgreSQLObservabilityStore.query_spans(
        store,
        limit=25,
        trace_id="trace-1",
        agent_name="planner",
        status="error",
        attributes_filter={"tenant_id": "acme", "tool_name": "search"},
    )

    query, params = cursor.executed[-1]
    assert "trace_id = %s" in query
    assert "agent_name = %s" in query
    assert "status = %s" in query
    assert query.count("attributes @> %s::jsonb") == 2
    assert params[0:3] == ["trace-1", "planner", "error"]
    assert json.loads(params[3]) == {"tenant_id": "acme"}
    assert json.loads(params[4]) == {"tool_name": "search"}
    assert params[-1] == 25


def test_postgres_observability_query_logs_builds_partial_error_friendly_filters():
    cursor = _FakeCursor()
    connection = _FakeConnection(cursor)
    store = PostgreSQLObservabilityStore.__new__(PostgreSQLObservabilityStore)
    store.dsn = "postgresql://obs-db/logos"
    store.schema = "public"
    store.table_prefix = "adk_observability"
    store.rules = []
    store.retention = {}
    store.cleanup_interval_seconds = 300
    store._last_cleanup_at = None
    store._connect = lambda: connection

    PostgreSQLObservabilityStore.query_logs(
        store,
        limit=10,
        level="error",
        trace_id="trace-2",
        source="http",
        attributes_filter={"tenant_id": "tenant-x"},
    )

    query, params = cursor.executed[-1]
    assert "level = %s" in query
    assert "trace_id = %s" in query
    assert "source = %s" in query
    assert "ORDER BY timestamp DESC" in query
    assert json.loads(params[3]) == {"tenant_id": "tenant-x"}
    assert params[-1] == 10


@pytest.mark.asyncio
async def test_open_telemetry_exporter_flushes_to_http_endpoint():
    captured: list[dict] = []

    async def handler(request: httpx.Request) -> httpx.Response:
        captured.append({
            "url": str(request.url),
            "body": gzip.decompress(request.content).decode(),
            "authorization": request.headers.get("authorization"),
            "tenant": request.headers.get("x-tenant-id"),
            "encoding": request.headers.get("content-encoding"),
            "content_type": request.headers.get("content-type"),
        })
        return httpx.Response(200, json={"ok": True})

    exporter = OpenTelemetryExporter(
        "http://collector.test/v1/traces",
        headers={"x-tenant-id": "acme"},
        bearer_token="secret-token",
        compression="gzip",
        transport=httpx.MockTransport(handler),
    )
    await exporter.on_span(
        Span(
            name="Agent.run",
            trace_id="trace-1",
            span_id="span-1",
            parent_span_id=None,
            started_at=datetime.now(UTC),
            ended_at=datetime.now(UTC),
            duration_ms=1.5,
        )
    )
    await exporter.on_metric(Metric(name="tokens", value=42, unit="count"))
    await exporter.flush()

    assert len(captured) == 2
    assert captured[0]["url"] == "http://collector.test/v1/traces"
    assert captured[1]["url"] == "http://collector.test/v1/metrics"
    assert captured[0]["authorization"] == "Bearer secret-token"
    assert captured[1]["authorization"] == "Bearer secret-token"
    assert captured[0]["tenant"] == "acme"
    assert captured[1]["tenant"] == "acme"
    assert captured[0]["encoding"] == "gzip"
    assert captured[1]["encoding"] == "gzip"
    assert captured[0]["content_type"] == "application/json"
    assert "\"resourceSpans\"" in captured[0]["body"]
    assert "\"traceId\"" in captured[0]["body"]
    assert "\"resourceMetrics\"" in captured[1]["body"]
    assert "\"gauge\"" in captured[1]["body"]
    assert exporter.spans == []
    assert exporter.metrics == []


@pytest.mark.asyncio
async def test_open_telemetry_exporter_preserves_explicit_authorization_header():
    captured: list[str | None] = []

    async def handler(request: httpx.Request) -> httpx.Response:
        captured.append(request.headers.get("authorization"))
        return httpx.Response(200, json={"ok": True})

    exporter = OpenTelemetryExporter(
        "http://collector.test/v1/traces",
        headers={"Authorization": "Api-Key custom"},
        bearer_token="ignored-token",
        transport=httpx.MockTransport(handler),
    )
    await exporter.on_span(
        Span(
            name="Agent.run",
            trace_id="trace-1",
            span_id="span-1",
            parent_span_id=None,
            started_at=datetime.now(UTC),
            ended_at=datetime.now(UTC),
            duration_ms=1.5,
        )
    )

    await exporter.flush()

    assert captured == ["Api-Key custom"]


@pytest.mark.asyncio
async def test_open_telemetry_exporter_sends_plain_json_without_compression():
    captured: list[dict] = []

    async def handler(request: httpx.Request) -> httpx.Response:
        captured.append(
            {
                "encoding": request.headers.get("content-encoding"),
                "content_type": request.headers.get("content-type"),
                "body": request.content.decode(),
            }
        )
        return httpx.Response(200, json={"ok": True})

    exporter = OpenTelemetryExporter(
        "http://collector.test/v1/traces",
        compression="none",
        transport=httpx.MockTransport(handler),
    )
    await exporter.on_span(
        Span(
            name="Agent.run",
            trace_id="trace-1",
            span_id="span-1",
            parent_span_id=None,
            started_at=datetime.now(UTC),
            ended_at=datetime.now(UTC),
            duration_ms=1.5,
        )
    )

    await exporter.flush()

    assert len(captured) == 1
    assert captured[0]["encoding"] is None
    assert captured[0]["content_type"] == "application/json"
    assert "\"resourceSpans\"" in captured[0]["body"]


@pytest.mark.asyncio
async def test_open_telemetry_exporter_flushes_in_batches():
    captured: list[dict] = []

    async def handler(request: httpx.Request) -> httpx.Response:
        captured.append(json.loads(request.content.decode()))
        return httpx.Response(200, json={"ok": True})

    exporter = OpenTelemetryExporter(
        "http://collector.test/v1/traces",
        batch_size=2,
        transport=httpx.MockTransport(handler),
    )
    for idx in range(5):
        await exporter.on_span(
            Span(
                name=f"span-{idx}",
                trace_id="trace-1",
                span_id=f"span-{idx}",
                parent_span_id=None,
                started_at=datetime.now(UTC),
                ended_at=datetime.now(UTC),
                duration_ms=1.0,
            )
        )
    for idx in range(3):
        await exporter.on_metric(Metric(name=f"metric-{idx}", value=idx, unit="count"))

    await exporter.flush()

    trace_batches = [batch for batch in captured if "resourceSpans" in batch]
    metric_batches = [batch for batch in captured if "resourceMetrics" in batch]
    assert len(trace_batches) == 3
    assert len(metric_batches) == 2
    assert [
        len(batch["resourceSpans"][0]["scopeSpans"][0]["spans"])
        for batch in trace_batches
    ] == [2, 2, 1]
    assert [
        len(batch["resourceMetrics"][0]["scopeMetrics"][0]["metrics"])
        for batch in metric_batches
    ] == [2, 1]
    assert exporter.spans == []
    assert exporter.metrics == []


@pytest.mark.asyncio
async def test_open_telemetry_exporter_retries_with_backoff():
    attempts = 0
    sleeps: list[float] = []

    async def handler(request: httpx.Request) -> httpx.Response:
        nonlocal attempts
        attempts += 1
        if attempts < 3:
            return httpx.Response(503, json={"ok": False})
        return httpx.Response(200, json={"ok": True})

    async def fake_sleep(delay: float) -> None:
        sleeps.append(delay)

    exporter = OpenTelemetryExporter(
        "http://collector.test/v1/traces",
        max_retries=2,
        backoff_seconds=0.5,
        sleeper=fake_sleep,
        transport=httpx.MockTransport(handler),
    )
    await exporter.on_span(
        Span(
            name="Agent.run",
            trace_id="trace-1",
            span_id="span-1",
            parent_span_id=None,
            started_at=datetime.now(UTC),
            ended_at=datetime.now(UTC),
            duration_ms=1.5,
        )
    )

    await exporter.flush()

    assert attempts == 3
    assert sleeps == [0.5, 1.0]
    assert exporter.spans == []


@pytest.mark.asyncio
async def test_open_telemetry_exporter_retains_unsent_tail_on_batch_failure():
    attempts = 0

    async def handler(request: httpx.Request) -> httpx.Response:
        nonlocal attempts
        attempts += 1
        if attempts == 2:
            return httpx.Response(503, json={"ok": False})
        return httpx.Response(200, json={"ok": True})

    exporter = OpenTelemetryExporter(
        "http://collector.test/v1/traces",
        batch_size=2,
        max_retries=0,
        transport=httpx.MockTransport(handler),
    )
    for idx in range(3):
        await exporter.on_span(
            Span(
                name=f"span-{idx}",
                trace_id="trace-1",
                span_id=f"span-{idx}",
                parent_span_id=None,
                started_at=datetime.now(UTC),
                ended_at=datetime.now(UTC),
                duration_ms=1.0,
            )
        )

    with pytest.raises(httpx.HTTPStatusError):
        await exporter.flush()

    assert attempts == 2
    assert [span.name for span in exporter.spans] == ["span-2"]


@pytest.mark.asyncio
async def test_agent_run_emits_basic_spans():
    exporter = RecordingExporter()
    provider = MockProvider(responses=["ok"])
    agent = Agent("bot").model(provider).observe(exporter)

    result = await agent.run("hi", session_id="s1")

    assert result.content == "ok"
    names = [span.name for span in exporter.spans]
    assert "Memory.remember" in names
    assert "LLM.generate" in names
    assert "Memory.commit" in names
    assert "Agent.run" in names
    assert all(span.trace_id == exporter.spans[0].trace_id for span in exporter.spans)
    assert exporter.flushed == 1


@pytest.mark.asyncio
async def test_agent_run_emits_tool_span():
    exporter = RecordingExporter()

    @tool
    async def add(a: int, b: int) -> int:
        return a + b

    tc = ToolCall(id="tc1", name="add", arguments={"a": 2, "b": 3})
    provider = MockProvider(
        responses=[
            Response(content="", tool_calls=[tc], usage=Usage(1, 1), raw={}),
            "done",
        ]
    )
    agent = Agent("bot").model(provider).tools([add]).observe(exporter)

    result = await agent.run("compute")

    assert result.content == "done"
    tool_spans = [span for span in exporter.spans if span.name == "Tool.add"]
    assert len(tool_spans) == 1
    assert tool_spans[0].status == "ok"
    assert tool_spans[0].attributes["call_id"] == "tc1"


@pytest.mark.asyncio
async def test_agent_stream_emits_stream_spans():
    exporter = RecordingExporter()
    provider = MockProvider(stream_responses=[["Hel", "lo"]])
    agent = Agent("bot").model(provider).observe(exporter)

    chunks = [chunk async for chunk in agent.stream("hi", session_id="s1")]

    assert chunks[-1].response.content == "Hello"
    names = [span.name for span in exporter.spans]
    assert "Memory.remember" in names
    assert "LLM.stream" in names
    assert "Memory.commit" in names
    assert "Agent.stream" in names
    assert exporter.flushed == 1
