#!/usr/bin/env python3
"""OpenTelemetry instrumentation skeleton: spans, metrics, structured logs.

Demonstrates all three pillars wired together with trace_id correlation.
Requires: pip install opentelemetry-api opentelemetry-sdk opentelemetry-exporter-otlp
"""

from __future__ import annotations

import time
import uuid
from contextlib import contextmanager
from dataclasses import dataclass, field
from typing import Any, Optional


# ── Skeleton (swap for real OTel SDK in prod) ─────────────────────────────


class FakeTracer:
    """Stand-in for OpenTelemetry tracer. In production:
    from opentelemetry import trace
    tracer = trace.get_tracer(__name__)
    """

    def start_span(self, name: str, attributes: dict = None) -> "FakeSpan":
        return FakeSpan(name, attributes or {})


tracer = FakeTracer()


@dataclass
class FakeSpan:
    name: str
    attributes: dict[str, Any]
    start_time: float = field(default_factory=time.monotonic)
    trace_id: str = field(default_factory=lambda: uuid.uuid4().hex[:16])
    status: str = "OK"

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        duration_ms = (time.monotonic() - self.start_time) * 1000
        self.attributes["duration_ms"] = round(duration_ms, 1)
        if exc_type:
            self.status = "ERROR"
            self.attributes["error"] = str(exc_val)
        print(f"  [SPAN] {self.name} | {self.status} | {duration_ms:.1f}ms | trace={self.trace_id[:8]}")

    def set_attribute(self, key: str, value: Any) -> None:
        self.attributes[key] = value


# ── Structured log emitter (JSON lines) ───────────────────────────────────


def log_event(level: str, msg: str, trace_id: str = "", **fields) -> None:
    import json
    record = {"ts": time.strftime("%Y-%m-%dT%H:%M:%S"), "level": level, "msg": msg}
    if trace_id:
        record["trace_id"] = trace_id
    record.update(fields)
    print(json.dumps(record))


# ── Metrics collector (RED: Rate, Errors, Duration) ───────────────────────


class Metrics:
    def __init__(self):
        self.counters: dict[str, int] = {}
        self.durations: dict[str, list[float]] = {}

    def inc(self, name: str) -> None:
        self.counters[name] = self.counters.get(name, 0) + 1

    def record_duration(self, name: str, ms: float) -> None:
        self.durations.setdefault(name, []).append(ms)

    def summary(self) -> dict:
        result = {}
        for name, count in self.counters.items():
            result[f"{name}_total"] = count
        for name, vals in self.durations.items():
            if vals:
                sorted_vals = sorted(vals)
                result[f"{name}_p50"] = sorted_vals[len(sorted_vals) // 2]
                result[f"{name}_p95"] = sorted_vals[int(len(sorted_vals) * 0.95)]
        return result


metrics = Metrics()


# ── Service with all three pillars ────────────────────────────────────────


def handle_request(method: str, path: str) -> dict:
    """Simulated request handler wired with traces + logs + metrics."""
    with tracer.start_span(f"{method} {path}", {"http.method": method, "http.path": path}) as span:
        trace_id = span.trace_id

        # Log: request start (structured, correlated)
        log_event("INFO", "request_start", trace_id=trace_id, method=method, path=path)
        metrics.inc("request_count")

        try:
            # Simulate work (instrumented with child spans)
            with tracer.start_span("db_query") as child:
                child.set_attribute("db.statement", "SELECT ...")
                time.sleep(0.02)
                result = {"status": "ok", "count": 42}

            with tracer.start_span("llm_call") as child:
                child.set_attribute("llm.model", "demo-model")
                child.set_attribute("llm.tokens", 150)
                time.sleep(0.05)
                result["answer"] = "Here is the response."

        except Exception as e:
            log_event("ERROR", "request_failed", trace_id=trace_id, error=str(e))
            metrics.inc("error_count")
            raise

        duration_ms = (time.monotonic() - span.start_time) * 1000
        metrics.record_duration("request_latency", duration_ms)
        log_event("INFO", "request_end", trace_id=trace_id, status=200, duration_ms=round(duration_ms, 1))
        return result


# ── Demo ───────────────────────────────────────────────────────────────────


def main() -> None:
    print("=== OpenTelemetry demo (skeleton) ===")
    for i in range(3):
        print(f"\n--- Request {i+1} ---")
        result = handle_request("GET", f"/api/search?q=test{i}")

    print("\n=== Metrics summary (RED) ===")
    for k, v in metrics.summary().items():
        print(f"  {k}: {v}")

    print("\nProduction wiring (replace skeletons):")
    print("  from opentelemetry import trace, metrics")
    print("  from opentelemetry.exporter.otlp.proto.grpc import ...")
    print("  from opentelemetry.sdk.trace import TracerProvider")


if __name__ == "__main__":
    main()