"""
E2E OTLP round-trip test.

Validates the full GAP-03 pipeline:
    Risicare Span → OTLPExporter (Python) → OTLP JSON
    → risicare-parser (simulated via JSON structure)
    → enrich.rs (simulated via attribute checking)
    → verify risicare.span_kind survives

This test does not require a running gateway — it validates the
Python export side and verifies the JSON matches what the Rust
parser expects.
"""

from __future__ import annotations

import json
from datetime import datetime, timezone
from typing import Any, Dict

import pytest

from risicare_core.types.base import SpanKind, SpanStatus, SemanticPhase
from risicare_core.types.spans import (
    ExceptionInfo,
    LLMAttributes,
    Span,
    SpanEvent,
    SpanLink,
)
from risicare.exporters.otlp import OTLPExporter, _KIND_TO_OTLP, _STATUS_TO_OTLP


class TestOTLPRoundTrip:
    """End-to-end OTLP round-trip: Risicare → OTLP JSON → verify structure."""

    @pytest.fixture
    def exporter(self):
        return OTLPExporter(
            endpoint="http://localhost:4318/v1/traces",
            service_name="roundtrip-test",
            service_version="1.0.0",
            environment="test",
        )

    def _build_rich_span(self) -> Span:
        """Create a span with all fields populated for thorough testing."""
        return Span(
            trace_id="a" * 32,
            span_id="b" * 16,
            name="chat_completions",
            parent_span_id="c" * 16,
            kind=SpanKind.LLM_CALL,
            start_time=datetime(2024, 6, 15, 10, 30, 0, tzinfo=timezone.utc),
            end_time=datetime(2024, 6, 15, 10, 30, 2, 500000, tzinfo=timezone.utc),
            status=SpanStatus.OK,
            status_message=None,
            attributes={
                "custom.tag": "my-value",
                "request.retries": 3,
            },
            events=[
                SpanEvent(
                    name="token_stream_start",
                    timestamp=datetime(2024, 6, 15, 10, 30, 0, 100000, tzinfo=timezone.utc),
                    attributes={"stream.chunk_count": 0},
                ),
            ],
            links=[
                SpanLink(
                    trace_id="d" * 32,
                    span_id="e" * 16,
                    attributes={"link.type": "follows_from"},
                ),
            ],
            exceptions=[
                ExceptionInfo(
                    type="RateLimitError",
                    message="Too many requests",
                    stacktrace="File api.py, line 42",
                    timestamp=datetime(2024, 6, 15, 10, 30, 1, tzinfo=timezone.utc),
                ),
            ],
            session_id="sess-abc-123",
            agent_id="agent-xyz-789",
            semantic_phase=SemanticPhase.THINK,
            llm=LLMAttributes(
                model="gpt-4o",
                provider="openai",
                prompt_tokens=150,
                completion_tokens=75,
                total_tokens=225,
                cost_usd=0.0045,
                temperature=0.7,
                max_tokens=4096,
                stop_reason="stop",
            ),
        )

    def test_full_roundtrip_structure(self, exporter):
        """Build OTLP payload from a rich span and verify all fields survive."""
        span = self._build_rich_span()
        payload = exporter._build_otlp_payload([span])

        # ---- Top-level structure ----
        assert "resourceSpans" in payload
        rs = payload["resourceSpans"][0]
        assert "resource" in rs
        assert "scopeSpans" in rs

        # ---- Resource attributes ----
        resource_attrs = {a["key"]: a["value"] for a in rs["resource"]["attributes"]}
        assert resource_attrs["service.name"]["stringValue"] == "roundtrip-test"
        assert resource_attrs["telemetry.sdk.name"]["stringValue"] == "risicare-sdk"

        # ---- Scope ----
        scope_span = rs["scopeSpans"][0]
        assert scope_span["scope"]["name"] == "risicare-sdk"
        assert scope_span["scope"]["version"] == "0.1.0"

        # ---- Span ----
        otlp_span = scope_span["spans"][0]

        # Core IDs
        assert otlp_span["traceId"] == "a" * 32
        assert otlp_span["spanId"] == "b" * 16
        assert otlp_span["parentSpanId"] == "c" * 16
        assert otlp_span["name"] == "chat_completions"

        # Kind: LLM_CALL → CLIENT (3)
        assert otlp_span["kind"] == 3

        # Status: OK → 1
        assert otlp_span["status"]["code"] == 1

        # Timestamps
        assert otlp_span["startTimeUnixNano"] != "0"
        assert otlp_span["endTimeUnixNano"] != "0"

    def test_risicare_span_kind_attribute_survives(self, exporter):
        """risicare.span_kind should be set for non-standard kinds."""
        span = self._build_rich_span()  # LLM_CALL kind
        otlp_span = exporter._convert_span(span)

        attr_map = {a["key"]: a["value"] for a in otlp_span["attributes"]}

        # LLM_CALL is not a standard OTel kind, so risicare.span_kind should exist
        assert "risicare.span_kind" in attr_map
        assert attr_map["risicare.span_kind"]["stringValue"] == "llm_call"

    def test_standard_kind_no_risicare_span_kind(self, exporter):
        """Standard OTel kinds should NOT have risicare.span_kind."""
        span = Span(
            trace_id="a" * 32,
            span_id="b" * 16,
            name="http_handler",
            kind=SpanKind.SERVER,
            start_time=datetime(2024, 1, 1, tzinfo=timezone.utc),
            end_time=datetime(2024, 1, 1, 0, 0, 1, tzinfo=timezone.utc),
        )
        otlp_span = exporter._convert_span(span)
        attr_keys = {a["key"] for a in otlp_span["attributes"]}
        assert "risicare.span_kind" not in attr_keys

    def test_gen_ai_attributes_roundtrip(self, exporter):
        """LLM attributes should map to gen_ai.* and back."""
        span = self._build_rich_span()
        otlp_span = exporter._convert_span(span)

        attr_map = {a["key"]: a["value"] for a in otlp_span["attributes"]}

        # Forward mapping: Risicare LLM → gen_ai.*
        assert attr_map["gen_ai.system"]["stringValue"] == "openai"
        assert attr_map["gen_ai.response.model"]["stringValue"] == "gpt-4o"
        assert attr_map["gen_ai.usage.prompt_tokens"]["intValue"] == "150"
        assert attr_map["gen_ai.usage.completion_tokens"]["intValue"] == "75"
        assert attr_map["gen_ai.usage.total_tokens"]["intValue"] == "225"
        assert attr_map["gen_ai.usage.cost"]["doubleValue"] == 0.0045
        assert attr_map["gen_ai.request.temperature"]["doubleValue"] == 0.7
        assert attr_map["gen_ai.request.max_tokens"]["intValue"] == "4096"
        assert attr_map["gen_ai.response.finish_reason"]["stringValue"] == "stop"

    def test_risicare_metadata_attributes(self, exporter):
        """Session, agent, and semantic_phase should be in attributes."""
        span = self._build_rich_span()
        otlp_span = exporter._convert_span(span)

        attr_map = {a["key"]: a["value"] for a in otlp_span["attributes"]}

        assert attr_map["session.id"]["stringValue"] == "sess-abc-123"
        assert attr_map["agent.id"]["stringValue"] == "agent-xyz-789"
        assert attr_map["risicare.semantic_phase"]["stringValue"] == "think"

    def test_custom_attributes_preserved(self, exporter):
        """Custom span attributes should be preserved in OTLP output."""
        span = self._build_rich_span()
        otlp_span = exporter._convert_span(span)

        attr_map = {a["key"]: a["value"] for a in otlp_span["attributes"]}

        assert attr_map["custom.tag"]["stringValue"] == "my-value"
        assert attr_map["request.retries"]["intValue"] == "3"

    def test_events_preserved(self, exporter):
        """Span events should be preserved in OTLP output."""
        span = self._build_rich_span()
        events = exporter._convert_events(span)

        # Should have 1 custom event + 1 exception event
        assert len(events) == 2

        custom_events = [e for e in events if e["name"] != "exception"]
        assert len(custom_events) == 1
        assert custom_events[0]["name"] == "token_stream_start"
        assert custom_events[0]["timeUnixNano"] != "0"

        exc_events = [e for e in events if e["name"] == "exception"]
        assert len(exc_events) == 1

        exc_attrs = {a["key"]: a["value"] for a in exc_events[0]["attributes"]}
        assert exc_attrs["exception.type"]["stringValue"] == "RateLimitError"
        assert exc_attrs["exception.message"]["stringValue"] == "Too many requests"
        assert exc_attrs["exception.stacktrace"]["stringValue"] == "File api.py, line 42"

    def test_links_preserved(self, exporter):
        """Span links should be preserved in OTLP output."""
        span = self._build_rich_span()
        links = exporter._convert_links(span)

        assert len(links) == 1
        assert links[0]["traceId"] == "d" * 32
        assert links[0]["spanId"] == "e" * 16

        link_attrs = {a["key"]: a["value"] for a in links[0]["attributes"]}
        assert link_attrs["link.type"]["stringValue"] == "follows_from"

    def test_otlp_json_is_valid_for_parser(self, exporter):
        """The OTLP JSON should be parseable by the Rust parser's expected schema.

        We validate the JSON structure matches the OtlpExportRequest format:
        - resourceSpans[].scopeSpans[].spans[]
        - Each span has traceId, spanId, name, kind, startTimeUnixNano, etc.
        - Attributes use {key, value: {stringValue|intValue|...}} format
        """
        span = self._build_rich_span()
        payload = exporter._build_otlp_payload([span])

        # Serialize and re-parse to ensure valid JSON
        json_str = json.dumps(payload)
        reparsed = json.loads(json_str)

        # Validate structure expected by risicare_parser::otlp::OtlpExportRequest
        assert isinstance(reparsed["resourceSpans"], list)
        rs = reparsed["resourceSpans"][0]

        assert "resource" in rs
        assert isinstance(rs["resource"]["attributes"], list)

        assert isinstance(rs["scopeSpans"], list)
        ss = rs["scopeSpans"][0]

        assert "scope" in ss
        assert isinstance(ss["scope"]["name"], str)

        assert isinstance(ss["spans"], list)
        otlp_span = ss["spans"][0]

        # Required fields the Rust parser expects
        assert isinstance(otlp_span["traceId"], str)
        assert isinstance(otlp_span["spanId"], str)
        assert isinstance(otlp_span["name"], str)
        assert isinstance(otlp_span["kind"], int)
        assert isinstance(otlp_span["startTimeUnixNano"], str)
        assert isinstance(otlp_span["status"], dict)
        assert isinstance(otlp_span["attributes"], list)
        assert isinstance(otlp_span["events"], list)
        assert isinstance(otlp_span["links"], list)

        # Validate attribute format
        for attr in otlp_span["attributes"]:
            assert "key" in attr
            assert "value" in attr
            assert isinstance(attr["key"], str)
            assert isinstance(attr["value"], dict)

    def test_all_span_kinds_have_otlp_mapping(self):
        """Every SpanKind variant should have an OTLP integer mapping."""
        for kind in SpanKind:
            assert kind.value in _KIND_TO_OTLP, (
                f"SpanKind.{kind.name} ({kind.value}) has no OTLP mapping"
            )

    def test_all_span_statuses_have_otlp_mapping(self):
        """Every SpanStatus variant should have an OTLP integer mapping."""
        for status in SpanStatus:
            assert status.value in _STATUS_TO_OTLP, (
                f"SpanStatus.{status.name} ({status.value}) has no OTLP mapping"
            )


class TestOTLPRoundTripEnrichment:
    """Simulate the enrichment that enrich.rs would perform on OTLP spans.

    This validates that the attributes we emit match what enrich.rs expects.
    """

    @pytest.fixture
    def exporter(self):
        return OTLPExporter(endpoint="http://localhost:4318/v1/traces")

    def test_enrichment_attribute_keys_match(self, exporter):
        """The gen_ai.* attribute keys we emit must match what enrich.rs reads."""
        # These are the exact keys enrich.rs looks for
        expected_keys = {
            "gen_ai.system",
            "gen_ai.response.model",
            "gen_ai.request.model",
            "gen_ai.usage.prompt_tokens",
            "gen_ai.usage.completion_tokens",
            "gen_ai.usage.total_tokens",
            "gen_ai.usage.cost",
            "agent.id",
            "agent.name",
            "session.id",
            "risicare.semantic_phase",
            "risicare.span_kind",
            "tool.name",
        }

        span = Span(
            trace_id="a" * 32,
            span_id="b" * 16,
            name="test",
            kind=SpanKind.LLM_CALL,
            start_time=datetime(2024, 1, 1, tzinfo=timezone.utc),
            end_time=datetime(2024, 1, 1, 0, 0, 1, tzinfo=timezone.utc),
            llm=LLMAttributes(
                provider="openai",
                model="gpt-4o",
                prompt_tokens=100,
                completion_tokens=50,
                total_tokens=150,
                cost_usd=0.003,
            ),
            session_id="sess-1",
            agent_id="agent-1",
            semantic_phase=SemanticPhase.THINK,
        )

        attrs = exporter._convert_attributes(span)
        attr_keys = {a["key"] for a in attrs}

        # Verify that our emitted keys are a superset of what enrich.rs reads
        # (we should emit all the keys that enrichment expects)
        for key in ["gen_ai.system", "gen_ai.response.model",
                     "gen_ai.usage.prompt_tokens", "gen_ai.usage.completion_tokens",
                     "gen_ai.usage.total_tokens", "gen_ai.usage.cost",
                     "agent.id", "session.id", "risicare.semantic_phase",
                     "risicare.span_kind"]:
            assert key in attr_keys, f"Missing expected attribute: {key}"

    def test_span_kind_restoration_would_work(self, exporter):
        """Simulate what enrich.rs does: read risicare.span_kind and restore kind."""
        # Export an AGENT span
        span = Span(
            trace_id="a" * 32,
            span_id="b" * 16,
            name="my_agent",
            kind=SpanKind.AGENT,
            start_time=datetime(2024, 1, 1, tzinfo=timezone.utc),
            end_time=datetime(2024, 1, 1, 0, 0, 1, tzinfo=timezone.utc),
        )
        otlp_span = exporter._convert_span(span)

        # AGENT maps to OTLP kind 1 (INTERNAL)
        assert otlp_span["kind"] == 1

        # But risicare.span_kind should preserve the original
        attr_map = {a["key"]: a["value"] for a in otlp_span["attributes"]}
        assert attr_map["risicare.span_kind"]["stringValue"] == "agent"

        # Simulate enrichment: if risicare.span_kind present, restore it
        restored_kind = attr_map["risicare.span_kind"]["stringValue"]
        assert restored_kind == "agent"  # Would be set as span.kind in enrich.rs

    def test_llm_call_kind_inference_would_work(self, exporter):
        """Simulate what enrich.rs does for kind inference when no risicare.span_kind."""
        # Export an INTERNAL span with gen_ai.system
        span = Span(
            trace_id="a" * 32,
            span_id="b" * 16,
            name="completion",
            kind=SpanKind.INTERNAL,  # Standard kind, no risicare.span_kind emitted
            start_time=datetime(2024, 1, 1, tzinfo=timezone.utc),
            end_time=datetime(2024, 1, 1, 0, 0, 1, tzinfo=timezone.utc),
            llm=LLMAttributes(provider="openai", model="gpt-4o"),
        )
        otlp_span = exporter._convert_span(span)

        attr_map = {a["key"]: a["value"] for a in otlp_span["attributes"]}

        # No risicare.span_kind for standard kind
        assert "risicare.span_kind" not in attr_map

        # But gen_ai.system IS present
        assert "gen_ai.system" in attr_map

        # enrich.rs would see: gen_ai.system present + kind == internal → upgrade to llm_call
        assert otlp_span["kind"] == 1  # INTERNAL in OTLP
        assert attr_map["gen_ai.system"]["stringValue"] == "openai"
