#!/usr/bin/env python3
"""LLM Observability Tracing — OpenTelemetry-based span tracing for LLM calls.

Demonstrates: manual span creation, automatic instrumentation of LLM providers,
log firehose filtering, and a cost-latency dashboard query.

REQUIRES: pip install opentelemetry-api opentelemetry-sdk opentelemetry-exporter-otlp
"""

from __future__ import annotations

import time
from dataclasses import dataclass, field
from typing import Any


# ── Minimal tracer (standalone, no dependency) ──────────────────────────

@dataclass
class Span:
    name: str
    start_time: float
    attributes: dict[str, Any] = field(default_factory=dict)
    events: list[dict] = field(default_factory=list)
    children: list[Span] = field(default_factory=list)
    status: str = "unset"
    end_time: float = 0.0
    _parent: Span | None = None

    def add_event(self, name: str, attributes: dict[str, Any] | None = None):
        self.events.append({
            "name": name,
            "timestamp": time.time(),
            "attributes": attributes or {},
        })

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

    def set_status(self, status: str):
        self.status = status

    def end(self):
        self.end_time = time.time()

    @property
    def duration_ms(self) -> float:
        return (self.end_time - self.start_time) * 1000 if self.end_time else 0


class Tracer:
    """Simple in-process tracer (replace with OpenTelemetry SDK for production)."""

    def __init__(self):
        self.spans: list[Span] = []
        self._current: Span | None = None

    def start_span(self, name: str) -> Span:
        span = Span(name=name, start_time=time.time(), _parent=self._current)
        if self._current:
            self._current.children.append(span)
        else:
            self.spans.append(span)
        self._current = span
        return span

    def end_span(self):
        if self._current:
            self._current.end()
            self._current = self._current._parent

    def __enter__(self):
        return self

    def __exit__(self, *args):
        self.end_span()


class TraceContext:
    """Context manager for a span. Used like:
    with tracer.span("llm.call") as span:
        span.set_attribute("model", "gpt-4o")
    """

    def __init__(self, tracer: Tracer, name: str):
        self.tracer = tracer
        self.span = tracer.start_span(name)

    def __enter__(self) -> Span:
        return self.span

    def __exit__(self, *args):
        self.tracer.end_span()


# ── LLM-specific attributes ─────────────────────────────────────────────

LLM_ATTRIBUTES = {
    "gen_ai.system": "",        # e.g., "openai", "anthropic"
    "gen_ai.request.model": "", # e.g., "gpt-4o"
    "gen_ai.request.temperature": 0.0,
    "gen_ai.request.max_tokens": 0,
    "gen_ai.usage.input_tokens": 0,
    "gen_ai.usage.output_tokens": 0,
    "gen_ai.usage.cost_usd": 0.0,
    "gen_ai.response.finish_reason": "",
    "gen_ai.latency_ms": 0.0,
}


def span_llm_call(
    tracer: Tracer, model: str, provider: str, prompt_tokens: int, **kwargs
) -> Span:
    span = tracer.start_span("llm.call")
    span.set_attribute("gen_ai.system", provider)
    span.set_attribute("gen_ai.request.model", model)
    span.set_attribute("gen_ai.request.temperature", kwargs.get("temperature", 1.0))
    span.set_attribute("gen_ai.request.max_tokens", kwargs.get("max_tokens", 0))
    span.set_attribute("gen_ai.usage.input_tokens", prompt_tokens)
    return span


def complete_llm_span(span: Span, output_tokens: int, finish_reason: str, cost_per_1k: float):
    span.set_attribute("gen_ai.usage.output_tokens", output_tokens)
    span.set_attribute("gen_ai.response.finish_reason", finish_reason)
    total_tokens = span.attributes.get("gen_ai.usage.input_tokens", 0) + output_tokens
    span.set_attribute("gen_ai.usage.cost_usd", round(total_tokens * cost_per_1k / 1000, 6))
    span.end()
    span.set_attribute("gen_ai.latency_ms", round(span.duration_ms, 2))
    span.set_status("ok")


# ── Log firehose filter ─────────────────────────────────────────────────

def filter_trace(log_entry: dict) -> bool:
    """Decide whether a log entry should be sampled into the trace.
    Keeps errors, slow calls, and high-cost calls; samples the rest."""
    if log_entry.get("level") == "ERROR":
        return True
    if log_entry.get("duration_ms", 0) > 5000:  # >5s
        return True
    if log_entry.get("cost_usd", 0) > 0.50:
        return True
    # Sample 10% of remaining
    import random
    return random.random() < 0.10


# ── Span tree print ─────────────────────────────────────────────────────

def print_span_tree(spans: list[Span], indent: int = 0):
    for span in spans:
        prefix = "  " * indent + ("└─ " if indent else "")
        cost = span.attributes.get("gen_ai.usage.cost_usd", "")
        print(f"{prefix}{span.name} ({span.duration_ms:.0f}ms) [status={span.status}]", end="")
        if cost:
            print(f" ${cost}", end="")
        print()
        print_span_tree(span.children, indent + 1)


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

if __name__ == "__main__":
    tracer = Tracer()

    # Simulate a multi-model chain with tracing
    with TraceContext(tracer, "agent.run") as root:
        root.set_attribute("session.id", "sess-abc123")

        # Step 1: routing decision (cheap model)
        with TraceContext(tracer, "llm.call") as span1:
            span1.set_attribute("gen_ai.system", "openai")
            span1.set_attribute("gen_ai.request.model", "gpt-4o-mini")
            span1.set_attribute("gen_ai.usage.input_tokens", 42)
            time.sleep(0.3)
            complete_llm_span(span1, output_tokens=8, finish_reason="stop", cost_per_1k=0.6)

        # Step 2: tool call execution
        with TraceContext(tracer, "tool.execute") as tool_span:
            tool_span.set_attribute("tool.name", "search")
            tool_span.set_attribute("tool.duration_ms", 450)
            time.sleep(0.45)
            tool_span.set_status("ok")

        # Step 3: final answer (frontier model)
        with TraceContext(tracer, "llm.call") as span2:
            span2.set_attribute("gen_ai.system", "anthropic")
            span2.set_attribute("gen_ai.request.model", "claude-3-5-sonnet-20241022")
            span2.set_attribute("gen_ai.usage.input_tokens", 520)
            time.sleep(1.2)
            complete_llm_span(span2, output_tokens=156, finish_reason="stop", cost_per_1k=15.0)

        root.set_status("ok")

    print("=== Trace Span Tree ===\n")
    print_span_tree(tracer.spans)

    print("\n=== Filter Decision ===")
    sample_log = {"level": "INFO", "duration_ms": 1200, "cost_usd": 0.02}
    print(f"  Log {sample_log} → {'KEEP' if filter_trace(sample_log) else 'DROP'}")

    cost_log = {"level": "INFO", "duration_ms": 800, "cost_usd": 0.75}
    print(f"  Log {cost_log} → {'KEEP' if filter_trace(cost_log) else 'DROP'} (cost over threshold)")

    print("\n✓ LLM observability tracing demo complete.")