"""Tests for observability layer."""

from __future__ import annotations

import json
import threading

from orxaq.kernel.observability import (
    ObservabilityCollector,
    SpanRecord,
    TraceMetrics,
)
from orxaq.kernel.tracing import set_tracer, trace_span

# ── SpanRecord ───────────────────────────────────────────


class TestSpanRecord:
    def test_to_dict(self):
        r = SpanRecord(
            span_id="s1",
            parent_id=None,
            name="test",
            start_time=1000.0,
            end_time=1001.0,
            duration_ms=1000.0,
            status="ok",
            error="",
            attributes={"key": "val"},
        )
        d = r.to_dict()
        assert d["span_id"] == "s1"
        assert d["duration_ms"] == 1000.0
        assert d["attributes"] == {"key": "val"}

    def test_with_children(self):
        r = SpanRecord(
            span_id="s1",
            parent_id=None,
            name="root",
            start_time=0,
            end_time=1,
            duration_ms=1000,
            status="ok",
            error="",
            children=["s2", "s3"],
        )
        d = r.to_dict()
        assert d["children"] == ["s2", "s3"]


# ── TraceMetrics ─────────────────────────────────────────


class TestTraceMetrics:
    def test_defaults(self):
        m = TraceMetrics()
        assert m.total_spans == 0
        assert m.error_rate == 0.0

    def test_to_dict(self):
        m = TraceMetrics(
            total_spans=10,
            error_spans=2,
            total_duration_ms=5000.0,
            error_rate=0.2,
            span_counts_by_name={"op": 10},
            avg_duration_by_name={"op": 500.0},
            p95_duration_by_name={"op": 900.0},
        )
        d = m.to_dict()
        assert d["total_spans"] == 10
        assert d["error_rate"] == 0.2
        assert d["span_counts_by_name"]["op"] == 10


# ── ObservabilityCollector basic ─────────────────────────


class TestCollectorBasic:
    def test_start_end_span(self):
        c = ObservabilityCollector()
        span = c.start_span("test_op", attributes={"key": "val"})
        c.end_span(span)
        spans = c.spans()
        assert len(spans) == 1
        assert spans[0].name == "test_op"
        assert spans[0].status == "ok"
        assert spans[0].duration_ms >= 0

    def test_error_span(self):
        c = ObservabilityCollector()
        span = c.start_span("failing_op")
        c.record_exception(span, ValueError("bad input"))
        spans = c.spans()
        assert len(spans) == 1
        assert spans[0].status == "error"
        assert spans[0].error == "bad input"
        assert spans[0].attributes.get("exception_type") == "ValueError"

    def test_filter_by_name(self):
        c = ObservabilityCollector()
        for name in ["op_a", "op_b", "op_a"]:
            span = c.start_span(name)
            c.end_span(span)
        assert len(c.spans("op_a")) == 2
        assert len(c.spans("op_b")) == 1

    def test_clear(self):
        c = ObservabilityCollector()
        span = c.start_span("test")
        c.end_span(span)
        assert len(c.spans()) == 1
        c.clear()
        assert len(c.spans()) == 0

    def test_max_spans_limit(self):
        c = ObservabilityCollector(max_spans=5)
        for i in range(10):
            span = c.start_span(f"op_{i}")
            c.end_span(span)
        assert len(c.spans()) == 5
        # Oldest spans should be evicted
        assert c.spans()[0].name == "op_5"


# ── Nested spans ─────────────────────────────────────────


class TestNestedSpans:
    def test_parent_child_tracking(self):
        c = ObservabilityCollector()
        parent = c.start_span("parent")
        child = c.start_span("child")
        c.end_span(child)
        c.end_span(parent)

        spans = c.spans()
        assert len(spans) == 2
        child_rec = next(s for s in spans if s.name == "child")
        parent_rec = next(s for s in spans if s.name == "parent")
        assert child_rec.parent_id == parent_rec.span_id
        assert child_rec.span_id in parent_rec.children

    def test_root_spans(self):
        c = ObservabilityCollector()
        root1 = c.start_span("root1")
        child1 = c.start_span("child1")
        c.end_span(child1)
        c.end_span(root1)
        root2 = c.start_span("root2")
        c.end_span(root2)

        roots = c.root_spans()
        assert len(roots) == 2
        assert all(r.parent_id is None for r in roots)

    def test_span_tree(self):
        c = ObservabilityCollector()
        root = c.start_span("root")
        root_id = root.attributes["_span_id"]
        c1 = c.start_span("child1")
        c.end_span(c1)
        c2 = c.start_span("child2")
        gc = c.start_span("grandchild")
        c.end_span(gc)
        c.end_span(c2)
        c.end_span(root)

        tree = c.span_tree(root_id)
        assert tree is not None
        assert tree["name"] == "root"
        assert len(tree["children"]) == 2
        child2_tree = next(ch for ch in tree["children"] if ch["name"] == "child2")
        assert len(child2_tree["children"]) == 1
        assert child2_tree["children"][0]["name"] == "grandchild"

    def test_span_tree_not_found(self):
        c = ObservabilityCollector()
        assert c.span_tree("nonexistent") is None


# ── Metrics ──────────────────────────────────────────────


class TestMetrics:
    def test_empty_metrics(self):
        c = ObservabilityCollector()
        m = c.metrics()
        assert m.total_spans == 0
        assert m.error_rate == 0.0

    def test_basic_metrics(self):
        c = ObservabilityCollector()
        for i in range(5):
            span = c.start_span("op_a")
            c.end_span(span)
        for i in range(3):
            span = c.start_span("op_b")
            if i == 0:
                c.record_exception(span, RuntimeError("fail"))
            else:
                c.end_span(span)

        m = c.metrics()
        assert m.total_spans == 8
        assert m.error_spans == 1
        assert m.error_rate > 0
        assert m.span_counts_by_name["op_a"] == 5
        assert m.span_counts_by_name["op_b"] == 3
        assert "op_a" in m.avg_duration_by_name
        assert "op_b" in m.p95_duration_by_name

    def test_p95_calculation(self):
        c = ObservabilityCollector()
        # Create 20 spans with varying durations (simulated)
        for _ in range(20):
            span = c.start_span("fast_op")
            c.end_span(span)
        m = c.metrics()
        assert m.p95_duration_by_name["fast_op"] >= 0


# ── Export ───────────────────────────────────────────────


class TestExport:
    def test_export_json(self):
        c = ObservabilityCollector()
        span = c.start_span("test")
        c.end_span(span)
        j = c.export_json()
        parsed = json.loads(j)
        assert len(parsed) == 1
        assert parsed[0]["name"] == "test"

    def test_export_ndjson(self):
        c = ObservabilityCollector()
        for i in range(3):
            span = c.start_span(f"op_{i}")
            c.end_span(span)
        ndjson = c.export_ndjson()
        lines = ndjson.strip().split("\n")
        assert len(lines) == 3
        for line in lines:
            parsed = json.loads(line)
            assert "span_id" in parsed


# ── Summary ──────────────────────────────────────────────


class TestSummary:
    def test_summary(self):
        c = ObservabilityCollector(max_spans=1000)
        for i in range(5):
            span = c.start_span(f"op_{i % 3}")
            c.end_span(span)
        s = c.summary()
        assert s["total_spans"] == 5
        assert s["unique_operations"] == 3
        assert s["max_spans"] == 1000


# ── Integration with trace_span ──────────────────────────


class TestTraceSpanIntegration:
    def test_with_trace_span_context_manager(self):
        c = ObservabilityCollector()
        set_tracer(c)
        with trace_span("my_operation", component="test"):
            pass
        spans = c.spans()
        assert len(spans) == 1
        assert spans[0].name == "my_operation"
        assert spans[0].attributes.get("component") == "test"
        # Restore null tracer
        from orxaq.kernel.tracing import NullTracer

        set_tracer(NullTracer())

    def test_nested_trace_spans(self):
        c = ObservabilityCollector()
        set_tracer(c)
        with trace_span("outer"):
            with trace_span("inner"):
                pass
        spans = c.spans()
        assert len(spans) == 2
        inner = next(s for s in spans if s.name == "inner")
        outer = next(s for s in spans if s.name == "outer")
        assert inner.parent_id == outer.span_id
        from orxaq.kernel.tracing import NullTracer

        set_tracer(NullTracer())

    def test_exception_in_trace_span(self):
        c = ObservabilityCollector()
        set_tracer(c)
        try:
            with trace_span("failing"):
                msg = "test error"
                raise ValueError(msg)
        except ValueError:
            pass
        spans = c.spans()
        assert len(spans) == 1
        assert spans[0].status == "error"
        assert "test error" in spans[0].error
        from orxaq.kernel.tracing import NullTracer

        set_tracer(NullTracer())


# ── Thread safety ────────────────────────────────────────


class TestThreadSafety:
    def test_concurrent_spans(self):
        c = ObservabilityCollector()
        errors: list[str] = []

        def worker(name: str, count: int):
            try:
                for _ in range(count):
                    span = c.start_span(f"thread_{name}")
                    c.end_span(span)
            except Exception as e:
                errors.append(str(e))

        threads = [threading.Thread(target=worker, args=(f"t{i}", 10)) for i in range(4)]
        for t in threads:
            t.start()
        for t in threads:
            t.join()

        assert len(errors) == 0
        assert len(c.spans()) == 40
