Part D: Observability

"""Observability for the support assistant: safe logs, traces, and a dashboard computed from them.

Two weeks of traffic are simulated through the Module 13 gateway with fake
providers (ScriptedLLM underneath, so replies are not model output). Latency
and failures are drawn from seeded distributions; on days 9 and 10 the primary
provider rate-limits hard, the way a real incident would look. Everything
below the simulation (redaction, tracing, the dashboard math) is real code you
can point at real logs.
"""
from __future__ import annotations

import contextvars
import hashlib
import hmac
import json
import os
import random
import re
import statistics
from collections import defaultdict
from contextlib import contextmanager
from dataclasses import dataclass, field
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Callable

from m13_gateway import FakeClock, FakeProvider, Gateway, Quota, Route, ticket_messages
from supportdesk.data import load_tickets
from supportdesk.kb_search import KBSearch

OUT = Path(__file__).resolve().parent / "out"
START = datetime(2026, 9, 1, tzinfo=timezone.utc).timestamp()

# Safe logging ----------------------------------------------------------------------

REDACTIONS = [
    (re.compile(r"[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Za-z]{2,}"), "[EMAIL]"),
    (re.compile(r"\b[A-Z]{2}\d{2}(?: ?[A-Z0-9]{4}){3,7}(?: ?[A-Z0-9]{1,3})?\b"), "[IBAN]"),
    (re.compile(r"\b(?:\d[ -]?){13,19}\b"), "[CARD]"),
    (re.compile(r"\+\d[\d ()-]{7,}\d"), "[PHONE]"),
]


def redact(text: str) -> str:
    for pattern, label in REDACTIONS:
        text = pattern.sub(label, text)
    return text


def hash_user(user_id: str) -> str:
    """Keyed hash: stable for joining logs, useless to anyone without the key (unlike plain sha256)."""
    key = os.environ.get("LOG_HASH_KEY", "dev-only-key-change-me").encode()
    return hmac.new(key, user_id.encode(), hashlib.sha256).hexdigest()[:16]


# Tracing -----------------------------------------------------------------------------

@dataclass
class Span:
    trace_id: str
    span_id: str
    parent_id: str | None
    name: str
    start: float
    attrs: dict[str, Any] = field(default_factory=dict)


_current: contextvars.ContextVar[Span | None] = contextvars.ContextVar("current_span", default=None)


class Tracer:
    """Writes one JSON line per finished span. Child spans inherit the trace id of their parent."""

    def __init__(self, path: Path, clock: Callable[[], float], seed: int = 0) -> None:
        self.path, self.clock, self.rng = path, clock, random.Random(seed)
        path.parent.mkdir(parents=True, exist_ok=True)
        self.file = path.open("w", encoding="utf-8")

    def _id(self, bits: int) -> str:
        return f"{self.rng.getrandbits(bits):0{bits // 4}x}"

    @contextmanager
    def span(self, name: str, **attrs: Any):
        parent = _current.get()
        span = Span(parent.trace_id if parent else self._id(128), self._id(64),
                    parent.span_id if parent else None, name, self.clock(), dict(attrs))
        token = _current.set(span)
        status = "ok"
        try:
            yield span
        except Exception as exc:
            status = "error"
            span.attrs["error"] = type(exc).__name__
            raise
        finally:
            _current.reset(token)
            record = {"trace_id": span.trace_id, "span_id": span.span_id, "parent_id": span.parent_id,
                      "name": name, "start": round(span.start, 3),
                      "duration_ms": round((self.clock() - span.start) * 1000, 1), "status": status, **span.attrs}
            self.file.write(json.dumps(record) + "\n")

    def close(self) -> None:
        self.file.close()


# Simulated traffic ---------------------------------------------------------------------

def simulate(path: Path, days: int = 14, per_day: int = 150, seed: int = 7) -> None:
    rng = random.Random(seed)
    clock = FakeClock(START)
    day = {"n": 0}

    def groq_outcome() -> str:
        if day["n"] in (9, 10) and rng.random() < 0.35:
            return "429:20"                     # incident: long Retry-After, so the gateway falls back
        return "500" if rng.random() < 0.01 else "ok"

    groq = FakeProvider("groq", clock=clock, latency_fn=lambda: rng.lognormvariate(6.0, 0.35), outcome_fn=groq_outcome)
    gemini = FakeProvider("gemini", clock=clock, latency_fn=lambda: rng.lognormvariate(6.8, 0.30),
                          outcome_fn=lambda: "500" if rng.random() < 0.02 else "ok")
    routes = {name: [Route("groq", groq, "openai/gpt-oss-120b"), Route("gemini", gemini, "gemini-3.5-flash")]
              for name in ("triage", "draft")}
    gw = Gateway(routes, quota=Quota(usd_per_day=1.0, requests_per_minute=30), clock=clock, sleep=clock.sleep, seed=seed)
    tracer = Tracer(path, clock, seed)
    search, tickets = KBSearch(), load_tickets()

    for d in range(1, days + 1):
        day["n"] = d
        for i in range(per_day):
            clock.now = START + (d - 1) * 86_400 + i * (86_400 / per_day)
            t = rng.choice(tickets)
            user = f"cust-{rng.randint(1, 400)}"
            body = t.body + (" Reach me at dana.lee@example.com or +49 30 1234 5678." if rng.random() < 0.2 else "")
            body += f" (ref {d}-{i})"          # real tickets are unique, so the exact cache rarely hits
            try:
                with tracer.span("handle_ticket", user=hash_user(user), ticket=t.id, day=d) as root:
                    with tracer.span("triage") as s:
                        r = gw.chat(ticket_messages(t.subject, body), route="triage", user=user)
                        s.attrs.update(provider=r.provider, cache=r.cache, usd=r.cost_usd, attempts=len(r.attempts))
                    with tracer.span("kb_search") as s:
                        hits = search.search(t.text, k=3)
                        clock.now += 0.004
                        s.attrs["top"] = hits[0].article_id if hits else None
                    with tracer.span("draft") as s:
                        r = gw.chat(ticket_messages(t.subject, body), route="draft", user=user)
                        s.attrs.update(provider=r.provider, cache=r.cache, usd=r.cost_usd, attempts=len(r.attempts),
                                       excerpt=redact(body)[:80])
                    root.attrs["category"] = t.gold["category"]
            except Exception:  # noqa: BLE001  (the span already recorded the error)
                pass
    tracer.close()


# Dashboard -----------------------------------------------------------------------------

def percentile(values: list[float], q: float) -> float:
    ordered = sorted(values)
    return ordered[min(int(q * len(ordered)), len(ordered) - 1)]


def dashboard(path: Path) -> dict[int, dict[str, float]]:
    spans = [json.loads(line) for line in path.open(encoding="utf-8")]
    by_trace: dict[str, list[dict]] = defaultdict(list)
    for s in spans:
        by_trace[s["trace_id"]].append(s)
    rows: dict[int, dict[str, list]] = defaultdict(lambda: defaultdict(list))
    for trace in by_trace.values():
        root = next(s for s in trace if s["parent_id"] is None)
        row = rows[root["day"]]
        row["latency"].append(root["duration_ms"])
        row["error"].append(root["status"] != "ok")
        row["usd"].append(sum(s.get("usd", 0.0) for s in trace))
        draft = next((s for s in trace if s["name"] == "draft" and s["status"] == "ok"), None)
        row["fallback"].append(bool(draft) and draft["provider"] != "groq")
        row["retries"].append(sum(s.get("attempts", 1) - 1 for s in trace if "attempts" in s))
    summary = {}
    print(f"{'day':>3} {'tasks':>6} {'p50 ms':>7} {'p95 ms':>7} {'errors':>7} {'fallback':>9} {'retries':>8} {'USD per 1k tasks':>17}")
    for d in sorted(rows):
        r = rows[d]
        ok_usd = [u for u, e in zip(r["usd"], r["error"]) if not e]
        summary[d] = {"p50": statistics.median(r["latency"]), "p95": percentile(r["latency"], 0.95),
                      "error_rate": sum(r["error"]) / len(r["error"]), "fallback": sum(r["fallback"]) / len(r["fallback"]),
                      "retries": statistics.mean(r["retries"]), "usd_per_1k": 1000 * sum(ok_usd) / max(len(ok_usd), 1)}
        s = summary[d]
        print(f"{d:>3} {len(r['latency']):>6} {s['p50']:>7.0f} {s['p95']:>7.0f} {s['error_rate']:>7.1%} "
              f"{s['fallback']:>9.0%} {s['retries']:>8.2f} {s['usd_per_1k']:>17.3f}")
    return summary


def show_slowest_trace(path: Path, day: int) -> None:
    spans = [json.loads(line) for line in path.open(encoding="utf-8")]
    roots = [s for s in spans if s["parent_id"] is None and s["day"] == day]
    slow = max(roots, key=lambda s: s["duration_ms"])
    print(f"\nslowest trace on day {day}: {slow['trace_id'][:12]}... {slow['duration_ms']:.0f} ms, user {slow['user']}")
    for s in sorted((s for s in spans if s["trace_id"] == slow["trace_id"]), key=lambda s: (s["start"], s["parent_id"] is not None)):
        detail = {k: s[k] for k in ("provider", "attempts", "cache", "usd", "excerpt") if k in s}
        indent = "  " if s["parent_id"] is None else "    "
        print(f"{indent}{s['name']:<14}{s['duration_ms']:>8.0f} ms  {s['status']:<5} {detail}")


def chart(summary: dict[int, dict[str, float]], path: Path) -> None:
    """Two panels, one measure each (never two y-scales on one axis); the incident days are shaded."""
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    ink, muted, surface = "#0b0b0b", "#52514e", "#fcfcfb"
    days = sorted(summary)
    fig, (a, b) = plt.subplots(1, 2, figsize=(10, 3.4), facecolor=surface)
    a.plot(days, [summary[d]["p50"] for d in days], color="#2a78d6", lw=2, marker="o", ms=5, label="p50")
    a.plot(days, [summary[d]["p95"] for d in days], color="#eb6834", lw=2, marker="o", ms=5, label="p95")
    a.set_title("Task latency (ms)", color=ink, loc="left")
    a.legend(frameon=False, labelcolor=muted, loc="upper left")
    b.bar(days, [summary[d]["usd_per_1k"] for d in days], color="#2a78d6", width=0.7)
    b.set_title("Cost per 1,000 tasks (USD)", color=ink, loc="left")
    for ax in (a, b):
        ax.set_facecolor(surface)
        ax.axvspan(8.5, 10.5, color="#f0efec", zorder=0)
        low, high = ax.get_ylim()
        ax.set_ylim(low, high + (high - low) * 0.12)                 # headroom so the label clears the data
        ax.text(9.5, high + (high - low) * 0.10, "incident", ha="center", va="top", color=muted, fontsize=8)
        ax.set_xlabel("day of September", color=muted)
        ax.tick_params(colors=muted)
        ax.grid(axis="y", color="#e5e4e0", lw=0.8)
        ax.set_axisbelow(True)
        for side in ("top", "right"):
            ax.spines[side].set_visible(False)
        for side in ("left", "bottom"):
            ax.spines[side].set_color("#c3c2b7")
    fig.tight_layout()
    fig.savefig(path, dpi=120, facecolor=surface)


if __name__ == "__main__":
    print(redact("Card 4111 1111 1111 1111, IBAN DE89 3704 0044 0532 0130 00, mail dana.lee@example.com, "
                 "call +49 30 1234 5678, invoice INV-2026-004512"))
    print("user cust-17 ->", hash_user("cust-17"))
    log = OUT / "m13_traces.jsonl"
    simulate(log)
    print(f"\nwrote {sum(1 for _ in log.open())} spans to {log.relative_to(Path.cwd())}\n")
    summary = dashboard(log)
    show_slowest_trace(log, 9)
    chart(summary, OUT / "m13_dashboard.png")
    print("\nchart saved to examples/out/m13_dashboard.png")

"""Drift detection: has this week's traffic (or this week's model output) moved away from the baseline?

Weekly ticket mixes are sampled from the real category proportions of the
72-ticket dataset (seeded), with a synthetic shift injected on purpose, so we
know the right answer and can check that the detectors find it.
"""
from __future__ import annotations

import math
import random
from collections import Counter

from scipy.stats import chi2_contingency

from supportdesk.data import CATEGORIES, load_tickets


def psi(expected: Counter, actual: Counter, keys, floor: float = 1e-4) -> float:
    """Population Stability Index: sum over bins of (a - e) * ln(a / e), using proportions."""
    e_total, a_total = sum(expected.values()), sum(actual.values())
    total = 0.0
    for k in keys:
        e = max(expected[k] / e_total, floor)
        a = max(actual[k] / a_total, floor)
        total += (a - e) * math.log(a / e)
    return total


def chi_square_p(expected: Counter, actual: Counter, keys) -> float:
    """p-value of a chi-square test that both weeks come from the same category distribution."""
    table = [[expected[k] for k in keys], [actual[k] for k in keys]]
    return chi2_contingency(table)[1]


def sample_week(weights: dict[str, float], n: int, rng: random.Random) -> Counter:
    return Counter(rng.choices(list(weights), weights=list(weights.values()), k=n))


def verdict(value: float) -> str:
    return "stable" if value < 0.1 else ("watch" if value < 0.25 else "ACT")


def main() -> None:
    base = Counter(t.gold["category"] for t in load_tickets())
    weights = {c: base[c] / sum(base.values()) for c in CATEGORIES}
    rng = random.Random(13)
    weeks = {"week 36 (baseline)": sample_week(weights, 1000, rng)}
    weeks["week 37 (no change)"] = sample_week(weights, 1000, rng)
    small = dict(weights, bug=weights["bug"] * 1.3)
    weeks["week 38 (bug +30%)"] = sample_week(small, 1000, rng)
    release = dict(weights, bug=weights["bug"] * 3, account_access=weights["account_access"] * 1.5)
    weeks["week 39 (bad release)"] = sample_week(release, 1000, rng)
    baseline = weeks["week 36 (baseline)"]

    print("share of tickets per category")
    print(f"{'week':<22}" + "".join(f"{c[:10]:>11}" for c in CATEGORIES))
    for name, counts in weeks.items():
        print(f"{name:<22}" + "".join(f"{counts[c] / sum(counts.values()):>11.1%}" for c in CATEGORIES))

    print(f"\n{'week vs baseline':<22}{'PSI':>7}{'verdict':>9}{'chi-square p':>14}")
    for name, counts in list(weeks.items())[1:]:
        value = psi(baseline, counts, CATEGORIES)
        print(f"{name:<22}{value:>7.3f}{verdict(value):>9}{chi_square_p(baseline, counts, CATEGORIES):>14.2g}")

    print("\nSame shift (bug +30%), different weekly volumes, 200 simulated weeks each:")
    for n in (100, 300, 1000, 3000):
        hits = Counter()
        for _ in range(200):
            a, b, same = sample_week(weights, n, rng), sample_week(small, n, rng), sample_week(weights, n, rng)
            hits["psi"] += psi(a, b, CATEGORIES) >= 0.1
            hits["psi_false"] += psi(a, same, CATEGORIES) >= 0.1
            hits["chi"] += chi_square_p(a, b, CATEGORIES) < 0.01
            hits["chi_false"] += chi_square_p(a, same, CATEGORIES) < 0.01
        print(f"  n={n:>5}/week: PSI>=0.1 flags {hits['psi'] / 2:>5.1f}% (false alarms {hits['psi_false'] / 2:>5.1f}%)"
              f"   chi-square p<0.01 flags {hits['chi'] / 2:>5.1f}% (false alarms {hits['chi_false'] / 2:.1f}%)")

    print("\nOutput drift: reply length in words, same inputs, model version changed silently")
    bins = ["<40", "40-79", "80-119", "120+"]

    def binned(lengths):
        return Counter(bins[min(int(x) // 40, 3)] for x in lengths)

    before = [rng.gauss(70, 18) for _ in range(800)]
    after = [rng.gauss(70 * 1.35, 25) for _ in range(800)]
    b, a = binned(before), binned(after)
    print("  bins:   " + "".join(f"{x:>8}" for x in bins))
    print("  before: " + "".join(f"{b[x]:>8}" for x in bins))
    print("  after:  " + "".join(f"{a[x]:>8}" for x in bins))
    value = psi(b, a, bins)
    print(f"  PSI {value:.3f} ({verdict(value)}), chi-square p {chi_square_p(b, a, bins):.2g}")


if __name__ == "__main__":
    main()