from __future__ import annotations

from pathlib import Path
from types import SimpleNamespace

import pytest

from agent_framework.observability.code_mapper import ObservabilityCodeMapper
from agent_framework.observability.observer import AgentObserver
from agent_framework.observability.telemetry import Telemetry
from app.observability.telemetry_observer import TelemetryBackedAgentObserver


class CaptureAnalytics:
    def __init__(self):
        self.calls = []

    async def publish(self, event_type, payload):
        self.calls.append((event_type, payload))


class CaptureBus:
    def __init__(self):
        self.calls = []

    async def publish(self, event_type, payload, **kwargs):
        self.calls.append((event_type, payload, kwargs))


def test_mapper_passthrough_and_contract_fields(tmp_path: Path):
    mapper = ObservabilityCodeMapper({"GRL.TOXOUT": "GRL.004"})
    code, payload, metadata = mapper.normalize_payload("GRL.TOXOUT", {"x": 1}, {})
    assert code == "GRL.004"
    assert payload["event_code_internal"] == "GRL.TOXOUT"
    assert metadata["event_code_internal"] == "GRL.TOXOUT"
    assert metadata["event_code_mapped"] == "GRL.004"
    assert mapper.map("GRL.NEW") == "GRL.NEW"


@pytest.mark.asyncio
async def test_agent_observer_maps_before_analytics_and_event_bus():
    analytics = CaptureAnalytics()
    bus = CaptureBus()
    observer = AgentObserver(
        analytics=analytics,
        event_bus=bus,
        code_mapper=ObservabilityCodeMapper({"GRL.TOXOUT": "GRL.004"}),
    )
    event = await observer.emit("GRL.TOXOUT", {"reason": "tox"})

    assert event["eventType"] == "GRL.004"
    assert event["payload"]["tag"] == "GRL.004"
    assert event["payload"]["event_code_internal"] == "GRL.TOXOUT"
    assert analytics.calls[0][0] == "GRL.004"
    assert bus.calls[0][0] == "GRL.004"


@pytest.mark.asyncio
async def test_telemetry_maps_direct_events(tmp_path: Path):
    mapping = tmp_path / "mapping.yaml"
    mapping.write_text("mappings:\n  GRL.TOXOUT: GRL.004\n", encoding="utf-8")
    settings = SimpleNamespace(
        OBSERVABILITY_CODE_MAPPING_ENABLED=True,
        OBSERVABILITY_CODE_MAPPING_PATH=str(mapping),
        ENABLE_LANGFUSE=False,
        ENABLE_OCI_STREAMING=False,
        ENABLE_OTEL=False,
        LANGFUSE_TRACE_MODE="verbose",
    )
    telemetry = Telemetry(settings)
    capture = []

    async def subscriber(evt):
        capture.append(evt)

    telemetry.event_bus.subscribe(subscriber)
    await telemetry.event("GRL.TOXOUT", {"a": 1}, kind="grl")

    assert capture
    assert capture[0].name == "GRL.004"
    assert capture[0].payload["event_code_internal"] == "GRL.TOXOUT"


@pytest.mark.asyncio
async def test_contas_telemetry_observer_returns_mapped_envelope(tmp_path: Path):
    mapping = tmp_path / "mapping.yaml"
    mapping.write_text("mappings:\n  GRL.TOXOUT: GRL.004\n", encoding="utf-8")
    settings = SimpleNamespace(
        OBSERVABILITY_CODE_MAPPING_ENABLED=True,
        OBSERVABILITY_CODE_MAPPING_PATH=str(mapping),
        ENABLE_LANGFUSE=False,
        ENABLE_OCI_STREAMING=False,
        ENABLE_OTEL=False,
        LANGFUSE_TRACE_MODE="verbose",
    )
    telemetry = Telemetry(settings)
    observer = TelemetryBackedAgentObserver(telemetry)
    event = await observer.emit("GRL.TOXOUT", {"rail": "TOXOUT"})

    assert event["eventType"] == "GRL.004"
    assert event["body"]["tag"] == "GRL.004"
    assert event["metadata"]["event_code_internal"] == "GRL.TOXOUT"

class _FakeObservation:
    def __init__(self):
        self.updates = []

    def update(self, **kwargs):
        self.updates.append(kwargs)


class _FakeObservationCM:
    def __init__(self, observation):
        self.observation = observation

    def __enter__(self):
        return self.observation

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


class _FakeLangfuse:
    def __init__(self):
        self.started = []

    def start_as_current_observation(self, **kwargs):
        self.started.append(dict(kwargs))
        return _FakeObservationCM(_FakeObservation())


@pytest.mark.asyncio
async def test_generation_observation_name_is_mapped_before_langfuse_and_event_bus(tmp_path: Path):
    mapping = tmp_path / "mapping.yaml"
    mapping.write_text(
        "mappings:\n  guardrail.dlex_in: GRL.004\n  guardrail.tox: GRL.005\n",
        encoding="utf-8",
    )
    settings = SimpleNamespace(
        OBSERVABILITY_CODE_MAPPING_ENABLED=True,
        OBSERVABILITY_CODE_MAPPING_PATH=str(mapping),
        ENABLE_LANGFUSE=False,
        ENABLE_OCI_STREAMING=False,
        ENABLE_OTEL=False,
        LANGFUSE_TRACE_MODE="verbose",
    )
    telemetry = Telemetry(settings)
    fake = _FakeLangfuse()
    telemetry.langfuse = fake
    telemetry.enabled = True
    capture = []

    async def subscriber(evt):
        capture.append(evt)

    telemetry.event_bus.subscribe(subscriber)
    async with telemetry.generation_span(
        "guardrail.dlex_in",
        "model-x",
        "input",
        metadata={"component": "guardrail.dlex_in"},
    ) as generation:
        generation.set_output("ok")

    assert fake.started[0]["name"] == "GRL.004"
    assert fake.started[0]["metadata"]["observability_name_internal"] == "guardrail.dlex_in"
    assert fake.started[0]["metadata"]["observability_name_mapped"] == "GRL.004"
    assert capture[0].name == "GRL.004"
    assert capture[0].payload["observability_name_internal"] == "guardrail.dlex_in"


@pytest.mark.asyncio
async def test_span_name_is_mapped_before_otel_event_bus_and_langfuse(tmp_path: Path):
    mapping = tmp_path / "mapping.yaml"
    mapping.write_text("mappings:\n  guardrail.tox: GRL.005\n", encoding="utf-8")
    settings = SimpleNamespace(
        OBSERVABILITY_CODE_MAPPING_ENABLED=True,
        OBSERVABILITY_CODE_MAPPING_PATH=str(mapping),
        ENABLE_LANGFUSE=False,
        ENABLE_OCI_STREAMING=False,
        ENABLE_OTEL=False,
        LANGFUSE_TRACE_MODE="verbose",
    )
    telemetry = Telemetry(settings)
    fake = _FakeLangfuse()
    telemetry.langfuse = fake
    telemetry.enabled = True
    capture = []

    async def subscriber(evt):
        capture.append(evt)

    telemetry.event_bus.subscribe(subscriber)
    async with telemetry.span("guardrail.tox", input={"text": "x"}):
        pass

    assert fake.started[0]["name"] == "GRL.005"
    assert fake.started[0]["metadata"]["observability_name_internal"] == "guardrail.tox"
    assert capture[0].name == "GRL.005.started"
    assert capture[-1].name == "GRL.005.completed"

@pytest.mark.asyncio
async def test_langfuse_analytics_direct_sdk_path_also_maps_observation_name(tmp_path: Path):
    from agent_framework.analytics.providers.langfuse import LangfuseAnalyticsPublisher

    mapping = tmp_path / "mapping.yaml"
    mapping.write_text("mappings:\n  guardrail.dlex_in: GRL.004\n", encoding="utf-8")
    settings = SimpleNamespace(
        OBSERVABILITY_CODE_MAPPING_ENABLED=True,
        OBSERVABILITY_CODE_MAPPING_PATH=str(mapping),
        LANGFUSE_PUBLIC_KEY="pk",
        LANGFUSE_SECRET_KEY="sk",
        LANGFUSE_HOST="http://localhost",
    )
    fake = _FakeLangfuse()
    publisher = LangfuseAnalyticsPublisher(settings=settings, langfuse=fake)

    # Give the technical event correlation so the provider does not intentionally skip it.
    from agent_framework.observability.context import set_observability_context, clear_observability_context
    clear_observability_context()
    set_observability_context(request_id="req-123", trace_id="req-123")
    try:
        await publisher.publish("guardrail.dlex_in", {"metadata": {"request_id": "req-123"}})
    finally:
        clear_observability_context()

    assert fake.started, "Langfuse analytics provider did not create an observation"
    assert fake.started[0]["name"] == "GRL.004"
    assert fake.started[0]["metadata"]["observability_name_internal"] == "guardrail.dlex_in"
    assert fake.started[0]["metadata"]["observability_name_mapped"] == "GRL.004"
