#!/usr/bin/env python3
"""the reference node usage report: rolls up the JSONL usage logs (voice/stt/image/music
generations) and the scraped vLLM token-counter snapshots into real numbers.

Usage:
  usage_report.py            # human-readable report, last 14 days
  usage_report.py --days 7   # narrower daily window
  usage_report.py --json     # machine-readable, same data
"""
import argparse
import datetime
import json
import os

USAGE_DIR = os.path.expanduser("~/lumbridge/compute/.compute/usage")

EVENT_LOGS = {
    "voice": ("voice.jsonl", "duration_s"),
    "stt": ("stt.jsonl", "audio_duration_s"),
    "image": ("image.jsonl", None),
    "music": ("music.jsonl", "duration_s"),
}

VLLM_MODELS = ("brain", "embed", "ocr")
VLLM_COUNTERS = ("prompt_tokens_total", "generation_tokens_total", "request_success_total")


def _read_jsonl(path):
    rows = []
    if not os.path.exists(path):
        return rows
    with open(path) as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            try:
                rows.append(json.loads(line))
            except json.JSONDecodeError:
                continue
    return rows


def _day(ts):
    return datetime.datetime.fromtimestamp(ts).date().isoformat()


def summarize_events(name, filename, duration_field, days):
    rows = _read_jsonl(os.path.join(USAGE_DIR, filename))
    cutoff = datetime.datetime.now().timestamp() - days * 86400
    by_day = {}
    total_count = len(rows)
    total_duration = 0.0
    for r in rows:
        ts = r.get("ts")
        if ts is None:
            continue
        dur = r.get(duration_field) or 0.0 if duration_field else 0.0
        total_duration += dur
        day = _day(ts)
        entry = by_day.setdefault(day, {"count": 0, "duration_s": 0.0})
        entry["count"] += 1
        entry["duration_s"] += dur
    recent_days = {d: v for d, v in by_day.items()
                   if datetime.datetime.fromisoformat(d).timestamp() >= cutoff - 86400}
    return {
        "service": name,
        "total_count": total_count,
        "total_duration_s": round(total_duration, 1),
        "total_duration_hours": round(total_duration / 3600, 2),
        "has_duration": duration_field is not None,
        "daily": dict(sorted(recent_days.items())),
    }


def summarize_vllm(days):
    rows = _read_jsonl(os.path.join(USAGE_DIR, "vllm_snapshots.jsonl"))
    cutoff = datetime.datetime.now().timestamp() - days * 86400
    by_model = {m: [] for m in VLLM_MODELS}
    for r in rows:
        model = r.get("model")
        if model in by_model:
            by_model[model].append(r)

    result = {}
    for model, snaps in by_model.items():
        snaps.sort(key=lambda r: r["ts"])
        current = snaps[-1] if snaps else None
        tracked = {c: 0.0 for c in VLLM_COUNTERS}
        prev = None
        for s in snaps:
            if s["ts"] < cutoff:
                prev = s
                continue
            if prev is not None:
                for c in VLLM_COUNTERS:
                    delta = s.get(c, 0) - prev.get(c, 0)
                    if delta > 0:
                        tracked[c] += delta
            prev = s
        result[model] = {
            "reachable_now": current is not None and current["ts"] >= cutoff,
            "snapshots_recorded": len(snaps),
            "current": {
                "prompt_tokens_total": current["prompt_tokens_total"],
                "generation_tokens_total": current["generation_tokens_total"],
                "request_success_total": current["request_success_total"],
            } if current else None,
            "tracked_since_monitoring_started": tracked,
        }
    return result


def render_text(events, vllm, days):
    lines = []
    lines.append(f"=== the reference node usage report (last {days} days) ===\n")

    for e in events:
        lines.append(f"-- {e['service']} --")
        if e["has_duration"]:
            lines.append(f"  total: {e['total_count']} generations, "
                         f"{e['total_duration_hours']}h ({e['total_duration_s']}s)")
        else:
            lines.append(f"  total: {e['total_count']} generations")
        if not e["daily"]:
            lines.append("  (no activity in this window)")
        else:
            for day, v in e["daily"].items():
                if e["has_duration"]:
                    lines.append(f"  {day}: {v['count']:>5} gens, {round(v['duration_s']/60, 1):>7} min")
                else:
                    lines.append(f"  {day}: {v['count']:>5} gens")
        lines.append("")

    lines.append("-- local LLM / vLLM token throughput --")
    for model, v in vllm.items():
        if not v["current"]:
            lines.append(f"  {model}: not reachable (no snapshot ever recorded)")
            continue
        status = "up" if v["reachable_now"] else f"down (last seen in a prior snapshot)"
        c = v["current"]
        t = v["tracked_since_monitoring_started"]
        lines.append(f"  {model}: {status}")
        lines.append(f"    current counters (since last process restart): "
                     f"{int(c['prompt_tokens_total']):,} prompt tok, "
                     f"{int(c['generation_tokens_total']):,} gen tok, "
                     f"{int(c['request_success_total']):,} requests")
        lines.append(f"    tracked in window ({v['snapshots_recorded']} scrapes): "
                     f"{int(t['prompt_tokens_total']):,} prompt tok, "
                     f"{int(t['generation_tokens_total']):,} gen tok, "
                     f"{int(t['request_success_total']):,} requests")
    lines.append("")
    return "\n".join(lines)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--days", type=int, default=14)
    parser.add_argument("--json", action="store_true")
    args = parser.parse_args()

    events = [summarize_events(name, filename, duration_field, args.days)
              for name, (filename, duration_field) in EVENT_LOGS.items()]
    vllm = summarize_vllm(args.days)

    if args.json:
        print(json.dumps({"events": events, "vllm": vllm, "days": args.days}, indent=2))
    else:
        print(render_text(events, vllm, args.days))


if __name__ == "__main__":
    main()