#!/usr/bin/env python3
"""omni-cswap usage reporter: snapshot `cswap list --json` on an interval and persist it.

Writes two local sinks under ~/.omni-cswap/usage/:
  · usage.jsonl  — append-only, one JSON object per account per tick (raw audit trail)
  · usage.db     — SQLite, table `usage_snapshots`, for querying without parsing JSONL

Each row records: ts, machine, account number, email, org name/uuid, whether it was the
active default, usageStatus, and the three utilization windows cswap exposes (5h %, 7d %,
per-model scoped %) plus their reset timestamps.

IMPORTANT — what cswap can and cannot report:
  cswap reports PERCENTAGES of each rate-limit window, not token counts and not cost. There is
  no per-session or per-prompt attribution in this data. If you need token/cost accounting, that
  lives in Claude's own usage reporting, not here.

Remote push: fill in push_to_remote() below (see REPORTING.md). It's a no-op by default, so
this runs standalone with no credentials.

Env knobs:
  OMNI_REPORT_INTERVAL   seconds between snapshots (default 900 = 15min)
  OMNI_REPORT_ONCE=1     take one snapshot and exit (cron-friendly)
  OMNI_MACHINE_ID        label for this host (default: hostname)
  OMNI_REPORT_DIR        override output dir (default ~/.omni-cswap/usage)
"""
import json, os, socket, sqlite3, subprocess, sys, time
from pathlib import Path

SELF = Path(__file__).resolve().parent
OUT = Path(os.environ.get("OMNI_REPORT_DIR", SELF / "usage"))
INTERVAL = float(os.environ.get("OMNI_REPORT_INTERVAL", "900"))
MACHINE = os.environ.get("OMNI_MACHINE_ID") or socket.gethostname()

DDL = """
CREATE TABLE IF NOT EXISTS usage_snapshots (
  ts            TEXT NOT NULL,
  machine       TEXT NOT NULL,
  account       TEXT NOT NULL,
  email         TEXT,
  org_name      TEXT,
  org_uuid      TEXT,
  is_active     INTEGER,
  usage_status  TEXT,
  five_hour_pct REAL,
  five_hour_resets TEXT,
  seven_day_pct REAL,
  seven_day_resets TEXT,
  scoped_json   TEXT,
  PRIMARY KEY (ts, machine, account)
);
CREATE INDEX IF NOT EXISTS idx_usage_machine_ts ON usage_snapshots(machine, ts);
CREATE INDEX IF NOT EXISTS idx_usage_account_ts ON usage_snapshots(account, ts);
"""


def _now():
    return time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())


def snapshot():
    """Run cswap and flatten its JSON into one row per account.

    Returns [] when no accounts are managed yet — the normal state on a freshly
    provisioned machine, before anyone has logged in. Not an error.
    """
    r = subprocess.run(["cswap", "list", "--json"], capture_output=True, text=True, timeout=60)
    if r.returncode != 0:
        raise RuntimeError(f"cswap list --json failed: {r.stderr.strip()[:300]}")
    data = json.loads(r.stdout)
    ts = _now()
    active = str(data.get("activeAccountNumber", ""))
    rows = []
    for a in data.get("accounts", []):
        u = a.get("usage") or {}
        five = u.get("fiveHour") or {}
        seven = u.get("sevenDay") or {}
        rows.append({
            "ts": ts,
            "machine": MACHINE,
            "account": str(a.get("number")),
            "email": a.get("email"),
            "org_name": a.get("organizationName"),
            "org_uuid": a.get("organizationUuid"),
            "is_active": 1 if str(a.get("number")) == active else 0,
            "usage_status": a.get("usageStatus"),
            "five_hour_pct": five.get("pct"),
            "five_hour_resets": five.get("resetsAt"),
            "seven_day_pct": seven.get("pct"),
            "seven_day_resets": seven.get("resetsAt"),
            "scoped_json": json.dumps(u.get("scoped") or []),
        })
    return rows


def write_jsonl(rows):
    OUT.mkdir(parents=True, exist_ok=True)
    with open(OUT / "usage.jsonl", "a") as f:
        for row in rows:
            f.write(json.dumps(row) + "\n")


def write_sqlite(rows):
    OUT.mkdir(parents=True, exist_ok=True)
    con = sqlite3.connect(OUT / "usage.db")
    try:
        con.executescript(DDL)
        con.executemany(
            "INSERT OR REPLACE INTO usage_snapshots "
            "(ts,machine,account,email,org_name,org_uuid,is_active,usage_status,"
            " five_hour_pct,five_hour_resets,seven_day_pct,seven_day_resets,scoped_json) "
            "VALUES (:ts,:machine,:account,:email,:org_name,:org_uuid,:is_active,:usage_status,"
            " :five_hour_pct,:five_hour_resets,:seven_day_pct,:seven_day_resets,:scoped_json)",
            rows,
        )
        con.commit()
    finally:
        con.close()


def push_to_remote(rows):
    """Ship rows to a central DB. No-op until you wire it — see REPORTING.md.

    Deliberately unconfigured: this installer is portable across machines and must not carry
    prod credentials. To enable, read a DSN from the environment and insert there, e.g.:

        dsn = os.environ.get("OMNI_REPORT_DSN")
        if not dsn: return
        import psycopg  # add to the service's env/venv
        with psycopg.connect(dsn) as con, con.cursor() as cur:
            cur.executemany("INSERT INTO usage_snapshots (...) VALUES (...) "
                            "ON CONFLICT (ts, machine, account) DO NOTHING", ...)

    Keep it best-effort: a reporting outage must never take down the loop.
    """
    return


def tick():
    rows = snapshot()
    if not rows:
        # Fresh machine, pre-login. Say so plainly and wait — the accounts will appear
        # once someone logs in, and the next tick picks them up with no restart needed.
        print(f"{_now()} {MACHINE}: no accounts managed yet — nothing to record", flush=True)
        return
    write_jsonl(rows)
    write_sqlite(rows)
    try:
        push_to_remote(rows)
    except Exception as e:
        sys.stderr.write(f"push_to_remote failed (local sinks still written): {e}\n")
    print(f"{rows[0]['ts']} {MACHINE}: recorded {len(rows)} accounts", flush=True)


def main():
    if os.environ.get("OMNI_REPORT_ONCE"):
        tick()
        return
    print(f"usage reporter start (machine={MACHINE} interval={INTERVAL}s out={OUT})", flush=True)
    while True:
        try:
            tick()
        except Exception as e:
            sys.stderr.write(f"tick error: {e}\n")
        time.sleep(INTERVAL)


if __name__ == "__main__":
    main()