"""Observability layer — structured tracing, metrics, and span export.

Extends the kernel tracing protocol with span collection, nested span
trees, metrics aggregation, and export-ready trace data. Plugs into
OTel-compatible exporters without requiring non-stdlib dependencies.

Stdlib-only (ADR-007 Tier-0 compliance).
"""

from __future__ import annotations

import json
import threading
import time
from collections import defaultdict
from dataclasses import dataclass, field
from typing import Any

from orxaq.kernel.tracing import SpanContext, TracingBackend


@dataclass
class SpanRecord:
    """Enriched span record with hierarchy and metadata."""

    span_id: str
    parent_id: str | None
    name: str
    start_time: float
    end_time: float
    duration_ms: float
    status: str
    error: str
    attributes: dict[str, Any] = field(default_factory=dict)
    children: list[str] = field(default_factory=list)

    def to_dict(self) -> dict:
        return {
            "span_id": self.span_id,
            "parent_id": self.parent_id,
            "name": self.name,
            "start_time": self.start_time,
            "end_time": self.end_time,
            "duration_ms": round(self.duration_ms, 3),
            "status": self.status,
            "error": self.error,
            "attributes": self.attributes,
            "children": self.children,
        }


@dataclass
class TraceMetrics:
    """Aggregated metrics from collected spans."""

    total_spans: int = 0
    error_spans: int = 0
    total_duration_ms: float = 0.0
    span_counts_by_name: dict[str, int] = field(default_factory=dict)
    avg_duration_by_name: dict[str, float] = field(default_factory=dict)
    p95_duration_by_name: dict[str, float] = field(default_factory=dict)
    error_rate: float = 0.0

    def to_dict(self) -> dict:
        return {
            "total_spans": self.total_spans,
            "error_spans": self.error_spans,
            "total_duration_ms": round(self.total_duration_ms, 3),
            "error_rate": round(self.error_rate, 4),
            "span_counts_by_name": dict(self.span_counts_by_name),
            "avg_duration_by_name": {
                k: round(v, 3) for k, v in self.avg_duration_by_name.items()
            },
            "p95_duration_by_name": {
                k: round(v, 3) for k, v in self.p95_duration_by_name.items()
            },
        }


class ObservabilityCollector(TracingBackend):
    """Collects spans with hierarchy tracking and metrics.

    Thread-safe span collector that implements TracingBackend protocol.
    Supports nested span trees, metrics aggregation, and export.

    Usage::

        from orxaq.kernel.tracing import set_tracer, trace_span
        collector = ObservabilityCollector(max_spans=10000)
        set_tracer(collector)

        with trace_span("my_operation", component="api"):
            do_work()

        metrics = collector.metrics()
        traces = collector.export_json()
    """

    def __init__(self, max_spans: int = 10000) -> None:
        self._max_spans = max_spans
        self._spans: list[SpanRecord] = []
        self._active_spans: dict[str, SpanContext] = {}
        self._span_counter = 0
        self._lock = threading.Lock()
        self._parent_stack: dict[int, list[str]] = {}
        self._children_map: dict[str, list[str]] = {}

    def start_span(
        self, name: str, attributes: dict[str, Any] | None = None
    ) -> SpanContext:
        """Start a new span, tracking parent via thread-local stack."""
        with self._lock:
            self._span_counter += 1
            span_id = f"span-{self._span_counter}"

            tid = threading.get_ident()
            parent_id = None
            if tid in self._parent_stack and self._parent_stack[tid]:
                parent_id = self._parent_stack[tid][-1]

            attrs = dict(attributes or {})
            attrs["_span_id"] = span_id
            attrs["_parent_id"] = parent_id

            span = SpanContext(
                name=name,
                start_time=time.time(),
                attributes=attrs,
            )
            self._active_spans[span_id] = span

            if tid not in self._parent_stack:
                self._parent_stack[tid] = []
            self._parent_stack[tid].append(span_id)

        return span

    def end_span(self, span: SpanContext) -> None:
        """End a span, recording it and updating parent's children."""
        span.end_time = time.time()
        span_id = span.attributes.get("_span_id", "")
        parent_id = span.attributes.get("_parent_id")

        record = SpanRecord(
            span_id=span_id,
            parent_id=parent_id,
            name=span.name,
            start_time=span.start_time,
            end_time=span.end_time,
            duration_ms=span.duration_ms,
            status=span.status,
            error=span.error,
            attributes={
                k: v for k, v in span.attributes.items()
                if not k.startswith("_")
            },
        )

        with self._lock:
            self._active_spans.pop(span_id, None)

            tid = threading.get_ident()
            if tid in self._parent_stack and self._parent_stack[tid]:
                self._parent_stack[tid].pop()
                if not self._parent_stack[tid]:
                    del self._parent_stack[tid]

            # Track parent->child relationship
            if parent_id:
                if parent_id not in self._children_map:
                    self._children_map[parent_id] = []
                self._children_map[parent_id].append(span_id)

            # Attach any children that ended before this span
            record.children = list(self._children_map.get(span_id, []))

            self._spans.append(record)
            if len(self._spans) > self._max_spans:
                self._spans = self._spans[-self._max_spans:]

    def record_exception(self, span: SpanContext, error: Exception) -> None:
        """Record an exception on a span."""
        span.status = "error"
        span.error = str(error)
        span.attributes["exception_type"] = type(error).__name__
        self.end_span(span)

    def spans(self, name: str | None = None) -> list[SpanRecord]:
        """Get collected spans, optionally filtered by name."""
        with self._lock:
            if name:
                return [s for s in self._spans if s.name == name]
            return list(self._spans)

    def root_spans(self) -> list[SpanRecord]:
        """Get only root spans (no parent)."""
        with self._lock:
            return [s for s in self._spans if s.parent_id is None]

    def span_tree(self, span_id: str) -> dict | None:
        """Build a nested tree from a root span."""
        with self._lock:
            span_map = {s.span_id: s for s in self._spans}

        root = span_map.get(span_id)
        if root is None:
            return None

        def build(sid: str) -> dict:
            s = span_map[sid]
            node = s.to_dict()
            node["children"] = [
                build(c) for c in s.children if c in span_map
            ]
            return node

        return build(span_id)

    def metrics(self) -> TraceMetrics:
        """Compute aggregate metrics from collected spans."""
        with self._lock:
            spans = list(self._spans)

        if not spans:
            return TraceMetrics()

        total = len(spans)
        errors = sum(1 for s in spans if s.status == "error")
        total_dur = sum(s.duration_ms for s in spans)

        durations_by_name: dict[str, list[float]] = defaultdict(list)
        counts_by_name: dict[str, int] = defaultdict(int)
        for s in spans:
            counts_by_name[s.name] += 1
            durations_by_name[s.name].append(s.duration_ms)

        avg_by_name = {}
        p95_by_name = {}
        for name, durs in durations_by_name.items():
            avg_by_name[name] = sum(durs) / len(durs) if durs else 0.0
            sorted_durs = sorted(durs)
            idx = int(len(sorted_durs) * 0.95)
            idx = min(idx, len(sorted_durs) - 1)
            p95_by_name[name] = sorted_durs[idx]

        return TraceMetrics(
            total_spans=total,
            error_spans=errors,
            total_duration_ms=total_dur,
            error_rate=errors / total if total > 0 else 0.0,
            span_counts_by_name=dict(counts_by_name),
            avg_duration_by_name=avg_by_name,
            p95_duration_by_name=p95_by_name,
        )

    def export_json(self, indent: int = 2) -> str:
        """Export all spans as JSON."""
        with self._lock:
            data = [s.to_dict() for s in self._spans]
        return json.dumps(data, indent=indent)

    def export_ndjson(self) -> str:
        """Export spans as newline-delimited JSON."""
        with self._lock:
            lines = [json.dumps(s.to_dict()) for s in self._spans]
        return "\n".join(lines)

    def clear(self) -> None:
        """Clear all collected spans."""
        with self._lock:
            self._spans.clear()
            self._active_spans.clear()
            self._parent_stack.clear()
            self._children_map.clear()

    def summary(self) -> dict:
        """High-level summary of observability state."""
        m = self.metrics()
        return {
            "total_spans": m.total_spans,
            "error_spans": m.error_spans,
            "error_rate": round(m.error_rate, 4),
            "unique_operations": len(m.span_counts_by_name),
            "active_spans": len(self._active_spans),
            "max_spans": self._max_spans,
        }


__all__ = [
    "ObservabilityCollector",
    "SpanRecord",
    "TraceMetrics",
]
