#!/usr/bin/env python3
"""
Generate the OpenTelemetry trace waterfall shown in docs/get-started/monitor.qmd.

Runs the small travel-assistant example from the docs against a live model,
captures the real spans it emits (the chatlas agent/chat/tool spans, the tool's
own `fetch_forecast` span, and the underlying httpx HTTP spans), and renders
them into a compact static SVG at docs/images/otel-trace.svg.

The figure is real captured data, so re-run this whenever the example, span
shape, or attributes change.

Usage:
    uv run --with opentelemetry-sdk --with opentelemetry-instrumentation-httpx \\
        python scripts/gen_otel_trace_svg.py

Environment:
    OPENAI_API_KEY: Required — the example makes a couple of cheap gpt-4o calls.
"""

from __future__ import annotations

import time
from html import escape
from pathlib import Path
from typing import Any

from opentelemetry import trace
from opentelemetry.instrumentation.httpx import HTTPXClientInstrumentor
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import SimpleSpanProcessor
from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter

SVG_PATH = Path(__file__).parent.parent / "docs" / "images" / "otel-trace.svg"
MODEL = "gpt-4o"
PROMPT = "I'm headed to Tokyo this weekend — what should I pack?"


def main() -> None:
    spans = capture_spans()
    svg = render_svg(spans)
    SVG_PATH.write_text(svg)
    print(f"wrote {SVG_PATH.relative_to(Path.cwd())} ({len(svg)} bytes)")


def capture_spans() -> list[dict[str, Any]]:
    """Run the worked example and return its spans as plain dicts."""
    exporter = InMemorySpanExporter()
    provider = TracerProvider()
    provider.add_span_processor(SimpleSpanProcessor(exporter))
    trace.set_tracer_provider(provider)

    # So the model's HTTP call shows up nested under each `chat` span.
    HTTPXClientInstrumentor().instrument()

    from chatlas import ChatOpenAI

    tracer = trace.get_tracer("travel_assistant")

    def get_weather(city: str) -> str:
        "Get the current weather forecast for a city."
        with tracer.start_as_current_span("fetch_forecast"):
            time.sleep(0.3)  # stand-in for a real forecast lookup
            return f"{city}: 14°C, light rain, breezy this weekend."

    chat = ChatOpenAI(
        model=MODEL,
        system_prompt=(
            "You are a concise travel assistant. "
            "Always call get_weather before advising what to pack."
        ),
    )
    chat.register_tool(get_weather)
    chat.chat(PROMPT, echo="none")

    out: list[dict[str, Any]] = []
    for s in exporter.get_finished_spans():
        out.append(
            {
                "name": s.name,
                "span_id": format(s.context.span_id, "016x") if s.context else None,
                "parent_id": format(s.parent.span_id, "016x") if s.parent else None,
                "start_ns": s.start_time,
                "end_ns": s.end_time,
                "attributes": dict(s.attributes or {}),
            }
        )
    return out


def render_svg(spans: list[dict[str, Any]]) -> str:
    """Render captured spans into a dark waterfall SVG (see Option C in the docs)."""
    by_id = {s["span_id"]: s for s in spans}
    rows = sorted(spans, key=lambda s: s["start_ns"])
    t0 = min(s["start_ns"] for s in spans)
    total = max(s["end_ns"] for s in spans) - t0

    # geometry
    PAD, ROW, CHIP_W, NAME_W, DUR_W = 14, 26, 40, 196, 50
    HEAD, AXIS, W = 30, 18, 720
    track_x = PAD + CHIP_W + 8 + NAME_W
    track_w = W - track_x - DUR_W - PAD
    H = HEAD + AXIS + len(rows) * ROW + PAD

    p = [
        f"<!-- Generated by scripts/gen_otel_trace_svg.py — do not edit by hand. -->",
        f'<svg xmlns="http://www.w3.org/2000/svg" width="{W}" height="{H}" '
        f'font-family="ui-monospace,Menlo,monospace" font-size="12">',
        f'<rect width="{W}" height="{H}" rx="10" fill="#13151a"/>',
        f'<text x="{PAD}" y="20" fill="#f3f4f6" font-weight="700">'
        f'Trace · chat("…what should I pack?")</text>',
        f'<text x="{W-PAD}" y="20" fill="#7d8696" text-anchor="end">'
        f"total {total/1e9:.2f}s · {len(rows)} spans</text>",
        f'<line x1="{PAD}" y1="{HEAD-4}" x2="{W-PAD}" y2="{HEAD-4}" stroke="#262a33"/>',
    ]

    for i in range(int(total / 1e9) + 2):
        x = track_x + (i * 1e9) / total * track_w
        if x > track_x + track_w:
            break
        p.append(f'<line x1="{x:.0f}" y1="{HEAD+2}" x2="{x:.0f}" y2="{H-PAD}" stroke="#1f232b"/>')
        p.append(f'<text x="{x+4:.0f}" y="{HEAD+AXIS-4}" fill="#616b7a" font-size="10">{i}s</text>')

    y = HEAD + AXIS
    for s in rows:
        chip, color = span_kind(s)
        d = depth(s, by_id)
        cy = y + ROW / 2
        ind = '<tspan fill="#4b5564">└ </tspan>' if d >= 2 else ""
        bx = track_x + (s["start_ns"] - t0) / total * track_w
        bw = max(3, (s["end_ns"] - s["start_ns"]) / total * track_w)
        dur = (s["end_ns"] - s["start_ns"]) / 1e9
        p += [
            f'<rect x="{PAD}" y="{cy-8:.0f}" width="{CHIP_W}" height="16" rx="4" fill="{color}"/>',
            f'<text x="{PAD+CHIP_W/2:.0f}" y="{cy+4:.0f}" fill="#0b0b0b" font-weight="700" '
            f'font-size="9" text-anchor="middle">{chip}</text>',
            f'<text x="{PAD+CHIP_W+8+d*12}" y="{cy+4:.0f}" fill="#cdd3df">{ind}{escape(label(s))}</text>',
            f'<rect x="{track_x}" y="{cy-7:.0f}" width="{track_w}" height="14" rx="3" fill="#1d2128"/>',
            f'<rect x="{bx:.1f}" y="{cy-7:.0f}" width="{bw:.1f}" height="14" rx="3" fill="{color}"/>',
            f'<text x="{W-PAD}" y="{cy+4:.0f}" fill="#7d8696" text-anchor="end">{dur:.2f}s</text>',
        ]
        y += ROW

    p.append("</svg>")
    return "\n".join(p) + "\n"


def depth(span: dict[str, Any], by_id: dict[Any, dict[str, Any]]) -> int:
    d, parent = 0, span["parent_id"]
    while parent in by_id:
        d += 1
        parent = by_id[parent]["parent_id"]
    return d


# HTTP semantic conventions were renamed (http.method -> http.request.method,
# http.url -> url.full); accept both so the script works across instrumentor versions.
HTTP_METHOD_KEYS = ("http.request.method", "http.method")
HTTP_URL_KEYS = ("url.full", "http.url")


def span_kind(span: dict[str, Any]) -> tuple[str, str]:
    name, attrs = span["name"], span["attributes"]
    if name == "invoke_agent":
        return "AGENT", "#8b5cf6"
    if name.startswith("chat"):
        return "CHAT", "#3b82f6"
    if name.startswith("execute_tool"):
        return "TOOL", "#10b981"
    if any(k in attrs for k in HTTP_METHOD_KEYS):
        return "HTTP", "#9ca3af"
    return "SPAN", "#9ca3af"


def label(span: dict[str, Any]) -> str:
    attrs = span["attributes"]
    url = next((attrs[k] for k in HTTP_URL_KEYS if k in attrs), None)
    method = next((attrs[k] for k in HTTP_METHOD_KEYS if k in attrs), "POST")
    if url and "openai.com" in str(url):
        return f"{method} " + str(url).split("openai.com", 1)[-1]
    return span["name"]


if __name__ == "__main__":
    main()
