"""#143 parity analysis: after-capture Jaeger dump vs the #137 baseline.

Writes span-inventory.json (baseline schema + nesting columns) and prints a
before/after diff plus targeted nesting / propagation checks.
"""
import json
import sys
from collections import defaultdict

W = sys.argv[1]
after = json.load(open(f"{W}/traces-after.json"))["data"]
before_inv = {r["operation"]: r for r in json.load(open(f"{W}/inventory-before.json"))}

# Attributes tracing-opentelemetry adds to every span (location / thread /
# timing metadata). Reported separately so the diff shows semantic changes.
TRACING_EXTRAS = {
    "busy_ns", "idle_ns", "target", "thread.id", "thread.name",
    "code.file.path", "code.line.number", "code.module.name",
}
# Spans produced by my auth probes, not by the baseline request set.
PROBE_NOISE = set()

spans = {}
by_trace = defaultdict(list)
for t in after:
    for s in t["spans"]:
        s["_tags"] = {kv["key"]: kv["value"] for kv in s["tags"]}
        spans[s["spanID"]] = s
        by_trace[s["traceID"]].append(s)


def parent(s):
    for r in s.get("references", []):
        if r["refType"] == "CHILD_OF":
            return r["spanID"], r["traceID"]
    return None, None


def op(sid):
    return spans[sid]["operationName"] if sid in spans else None


# Inventory in the baseline's schema, plus nesting columns.
inv = defaultdict(lambda: {"span_count": 0, "traces": set(), "kinds": set(), "scopes": set(),
                           "statuses": set(), "keys": set(), "parents": defaultdict(int),
                           "orphans": 0})
for s in spans.values():
    r = inv[s["operationName"]]
    r["span_count"] += 1
    r["traces"].add(s["traceID"])
    r["kinds"].add(s["_tags"].get("span.kind", "internal"))
    r["scopes"].add(s["_tags"].get("otel.scope.name", "?"))
    r["statuses"].add(s["_tags"].get("otel.status_code", "(none)"))
    r["keys"].update(s["_tags"])
    pid, _ = parent(s)
    if pid is None:
        r["orphans"] += 1
    else:
        r["parents"][op(pid) or "(remote)"] += 1

inventory = []
for name, r in sorted(inv.items(), key=lambda kv: -kv[1]["span_count"]):
    inventory.append({
        "operation": name,
        "span_count": r["span_count"],
        "trace_count": len(r["traces"]),
        "kind": "/".join(sorted(r["kinds"])),
        "scope": "/".join(sorted(r["scopes"])),
        "status": "/".join(sorted(r["statuses"])),
        "attribute_keys": sorted(r["keys"] - TRACING_EXTRAS),
        "parents": dict(r["parents"]),
        "root_count": r["orphans"],
    })
json.dump(inventory, open(f"{W}/span-inventory.json", "w"), indent=2)

print("== Operations: before -> after ==")
before_ops, after_ops = set(before_inv), set(inv)
print("only before:", sorted(before_ops - after_ops))
print("only after: ", sorted(after_ops - before_ops))
print()
print("== Common operations: kind / scope / attribute-key changes ==")
for name in sorted(before_ops & after_ops):
    b, a = before_inv[name], next(x for x in inventory if x["operation"] == name)
    bk, ak = set(b["attribute_keys"]), set(a["attribute_keys"])
    changes = []
    if b["kind"] != a["kind"]:
        changes.append(f"kind {b['kind']}->{a['kind']}")
    if b["scope"] != a["scope"]:
        changes.append(f"scope {b['scope']}->{a['scope']}")
    if bk - ak:
        changes.append(f"LOST {sorted(bk - ak)}")
    if ak - bk:
        changes.append(f"added {sorted(ak - bk)}")
    nest = f"traces {b['trace_count']}/{b['span_count']} -> {a['trace_count']}/{a['span_count']}, roots={a['root_count']}"
    print(f"{name}: {'; '.join(changes) or 'same'} | {nest}")

print()
print("== Roots in the after capture (span with no parent) ==")
for s in spans.values():
    if parent(s)[0] is None:
        print(f"  {s['operationName']}  trace={s['traceID'][:8]} spans_in_trace={len(by_trace[s['traceID']])}")


def chain(s):
    out = [s["operationName"]]
    pid, _ = parent(s)
    while pid:
        if pid not in spans:
            out.append(f"(remote {pid})")
            break
        out.append(spans[pid]["operationName"])
        pid, _ = parent(spans[pid])
    return " <- ".join(out)


print()
print("== Ancestry of key spans ==")
seen = set()
for s in spans.values():
    n = s["operationName"]
    key = (n, s["_tags"].get("llm.operation"))
    if n in ("llm.request", "ai.agentic.loop", "ai.insight.generate_agentic",
             "db.insert.store_insight", "db.update.update_training_messages",
             "ai.chat.turn.stream", "list_photos", "list_memories",
             "generate_video_thumbnail") and key not in seen:
        seen.add(key)
        extra = f" [llm.operation={key[1]}]" if key[1] else ""
        print(f"  {chain(s)}{extra}")

print()
print("== traceparent probe ==")
probe = [s for s in spans.values() if s["traceID"].startswith("11111111")]
for s in sorted(probe, key=lambda s: s["startTime"]):
    print(f"  {chain(s)}  trace={s['traceID']}")

print()
print("== llm.request attribute presence ==")
for s in spans.values():
    if s["operationName"] == "llm.request":
        t = s["_tags"]
        print(f"  op={t.get('llm.operation')} model={t.get('llm.model')} provider={t.get('llm.provider')} "
              f"status={t.get('http.status_code')} tokens_in={t.get('llm.tokens.input')} "
              f"parent={op(parent(s)[0])}")
print()
print("== turn / insight spans: llm.backend + llm.model ==")
for s in spans.values():
    if s["operationName"] in ("ai.chat.turn.stream", "ai.insight.generate_agentic"):
        t = s["_tags"]
        print(f"  {s['operationName']}: llm.backend={t.get('llm.backend')} llm.model={t.get('llm.model')} backend={t.get('backend')} model={t.get('model')}")