#!/usr/bin/env python3
"""Ask production's Loki or Tempo a question, from anywhere that can run docker on Tower.

    query.py logs    '<LogQL log query>'    [since]   lines, newest first
    query.py metric  '<LogQL metric query>' [since]   one value per series; $range is `since`
    query.py traces  '<TraceQL query>'      [since]   matching traces, slowest first
    query.py trace   <trace id>                       one trace's spans, as a tree

`since` is 1h, 24h, 7d and the like (default 24h). Loki and Tempo publish no query port on the
host, so each call runs a throwaway alpine container on the monitoring project's network.
"""
import json, subprocess, sys, time, urllib.parse

NET = "monitoring_default"


def fetch(url):
    out = subprocess.run(["docker", "run", "--rm", "--network", NET, "alpine", "wget", "-qO-", url],
                         capture_output=True, text=True, timeout=120)
    if out.returncode:
        sys.exit(f"query failed: {out.stderr.strip() or out.stdout.strip()}\n{url}")
    return json.loads(out.stdout)


def seconds(since):
    return int(since[:-1]) * {"m": 60, "h": 3600, "d": 86400}[since[-1]]


def main():
    if len(sys.argv) < 3:
        sys.exit(__doc__)
    kind, q = sys.argv[1], sys.argv[2]
    since = sys.argv[3] if len(sys.argv) > 3 else "24h"
    now = time.time()
    start = now - seconds(since)
    enc = urllib.parse.quote
    if kind == "logs":
        r = fetch(f"http://loki:3100/loki/api/v1/query_range?query={enc(q)}&limit=500"
                  f"&start={int(start * 1e9)}&end={int(now * 1e9)}&direction=backward")
        lines = [(ts, line) for s in r["data"]["result"] for ts, line in s["values"]]
        for ts, line in sorted(lines, reverse=True):
            print(line)
        print(f"-- {len(lines)} line(s){' (limit reached)' if len(lines) >= 500 else ''}", file=sys.stderr)
    elif kind == "metric":
        r = fetch(f"http://loki:3100/loki/api/v1/query?query={enc(q.replace('$range', since))}&time={int(now * 1e9)}")
        rows = sorted(r["data"]["result"], key=lambda s: -float(s["value"][1]))
        for s in rows:
            labels = {k: v for k, v in s["metric"].items() if k not in ("container", "compose_project", "service_name")}
            print(f"{s['value'][1]:>12}  {json.dumps(labels) if labels else ''}")
        print(f"-- {len(rows)} series", file=sys.stderr)
    elif kind == "traces":
        r = fetch(f"http://tempo:3200/api/search?q={enc(q)}&start={int(start)}&end={int(now)}&limit=100")
        traces = sorted(r.get("traces", []), key=lambda t: -t.get("durationMs", 0))
        for t in traces:
            when = time.strftime("%m-%d %H:%M:%S", time.gmtime(int(t["startTimeUnixNano"]) / 1e9))
            print(f"{t.get('durationMs', 0):>8}ms  {when}Z  {t['traceID']}  {t.get('rootTraceName', '')}")
        print(f"-- {len(traces)} trace(s)", file=sys.stderr)
    elif kind == "trace":
        r = fetch(f"http://tempo:3200/api/traces/{q}")
        spans = [sp for b in r.get("batches", r.get("resourceSpans", []))
                 for ss in b.get("scopeSpans", b.get("instrumentationLibrarySpans", [])) for sp in ss["spans"]]
        kids = {}
        for sp in spans:
            kids.setdefault(sp.get("parentSpanId", ""), []).append(sp)

        def show(sp, depth):
            ms = (int(sp["endTimeUnixNano"]) - int(sp["startTimeUnixNano"])) / 1e6
            attrs = {a["key"]: next(iter(a["value"].values()), None) for a in sp.get("attributes", [])
                     if not a["key"].startswith(("code.", "thread.")) and a["key"] not in ("busy_ns", "idle_ns", "target")}
            err = " ERROR" if sp.get("status", {}).get("code") in (2, "STATUS_CODE_ERROR") else ""
            print(f"{'  ' * depth}{sp['name']}  {ms:.0f}ms{err}  {json.dumps(attrs) if attrs else ''}")
            for e in sp.get("events", []):
                msg = next((a["value"].get("stringValue") for a in e.get("attributes", []) if a["key"] == "message"), e.get("name"))
                # A scan logs a skip for every feed not due; they bury what happened.
                if '"ev":"feed_skip"' in (msg or ""):
                    continue
                print(f"{'  ' * depth}  - {msg}")
            for k in sorted(kids.get(sp["spanId"], []), key=lambda s: int(s["startTimeUnixNano"])):
                show(k, depth + 1)

        ids = {sp["spanId"] for sp in spans}
        for root in [sp for sp in spans if sp.get("parentSpanId", "") not in ids]:
            show(root, 0)
    else:
        sys.exit(__doc__)


if __name__ == "__main__":
    main()