#!/usr/bin/env python3
"""Ship per-session token spend snapshots from agentsview to tokenspend.

Privacy model: only sessions whose project git remote matches an `allow` glob
in the config, or whose working directory sits under a `.tokenspend` marker
file (writing the marker is the consent), leave this machine. Everything else
(including sessions with no git remote) is counted locally and never sent.
Session ids are replaced by a salted HMAC; no content, titles, or paths are
ever read into the payload - a marker ships only its jobcode.

Each row carries a session's cumulative lifetime usage for one model. The
server timestamps every snapshot on receipt and attributes the difference from
the previous snapshot to that moment, so the shipper can run as often as you
like (daily is the recommended default; more often is safe) and nothing is
lost or double counted.

The shipper is installed once by the Setup page and then runs from cron, so it
can refresh itself: --check-for-updates compares this file with the copy the
server serves at /setup/ship.py, and --update replaces it. By default every
shipping run does that update itself first, so changes roll out to every
install; "auto_update": false in the config opts out, and then a run ends with
a one-line notice when an update is available (--no-nag silences it). The same
step keeps the private agentsview at AGENTSVIEW_VERSION, the release this
shipper was written against, so bumping that constant rolls out to every
install too (an agentsview from PATH, or one in a directory the shipper cannot
write to such as the nix flake's, is left alone).
"""
import argparse
import datetime as dt
import fnmatch
import hashlib
import hmac
import io
import json
import os
import platform
import secrets
import sqlite3
import subprocess
import sys
import tarfile
import urllib.parse
import urllib.request

CONFIG_DIR = os.path.expanduser("~/.config/tokenspend")
# JSON, not TOML: Apple's python3 is 3.9 and has no tomllib, and the shipper
# must run under cron/launchd on a stock Mac. Stdlib only, Python 3.9+.
CONFIG_PATH = os.environ.get("TOKENSPEND_CONFIG", os.path.join(CONFIG_DIR, "config.json"))
SALT_PATH = os.environ.get("TOKENSPEND_SALT", os.path.join(CONFIG_DIR, "salt"))
# The installer (the server's Setup page) puts a private, pinned agentsview
# here with its own data directory, so the shipper never depends on or touches
# an agentsview the developer may have installed for themselves. When the
# private copy is absent we fall back to `agentsview` on PATH and its default
# data directory (nix dev shell, older installs).
SHARE_DIR = os.path.expanduser("~/.local/share/tokenspend")
AGENTSVIEW_BIN = os.environ.get("TOKENSPEND_AGENTSVIEW", os.path.join(SHARE_DIR, "bin", "agentsview"))
AGENTSVIEW_DATA = os.environ.get("TOKENSPEND_AGENTSVIEW_DATA", os.path.join(SHARE_DIR, "agentsview-data"))
# The agentsview release this shipper is written against. The server's Setup
# page and flake.nix pin the same version; the server reads it from this line.
AGENTSVIEW_VERSION = "0.45.0"
AGENTSVIEW_RELEASES = "https://github.com/kenn-io/agentsview/releases/download/"
# Sessions active in the last N days are re-snapshotted each run. Anything
# untouched for longer has already been captured by an earlier run.
LOOKBACK_DAYS = 3

# export field name -> payload field name
TOKEN_FIELDS = {
    "input_tokens": "input_tokens",
    "output_tokens": "output_tokens",
    "cache_creation_input_tokens": "cache_creation_tokens",
    "cache_read_input_tokens": "cache_read_tokens",
}

EXAMPLE_CONFIG = '''{
  "developer": "yourname",
  "endpoint": "https://spend.example.com/api/ingest",
  "token": "<from the Setup page>",
  "allow": ["github.com/sixfeetup/*"],
  "auto_update": true
}
'''


def load_config(path=CONFIG_PATH, dry_run=False):
    try:
        with open(path) as f:
            cfg = json.load(f)
    except FileNotFoundError:
        if os.path.exists(os.path.splitext(path)[0] + ".toml"):
            sys.exit(f"config moved from TOML to JSON ({path}): rerun the one-line command "
                     "from the tokenspend site's Setup page, or convert it by hand")
        sys.exit(f"no config at {path}. Create it like:\n\n{EXAMPLE_CONFIG}")
    required = ("developer", "allow") if dry_run else ("developer", "allow", "endpoint", "token")
    for key in required:
        if key not in cfg:
            sys.exit(f"config {path}: missing '{key}'")
    return cfg


def load_salt(path=SALT_PATH):
    """Per-machine secret used to pseudonymize session ids. Created once, never shipped."""
    try:
        with open(path) as f:
            return f.read().strip()
    except FileNotFoundError:
        salt = secrets.token_hex(32)
        os.makedirs(os.path.dirname(path), exist_ok=True)
        with open(os.open(path, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600), "w") as f:
            f.write(salt)
        return salt


def private_agentsview():
    return os.path.exists(AGENTSVIEW_BIN)


def agentsview(*args):
    if private_agentsview():
        cmd = [AGENTSVIEW_BIN, *args]
        # NO_DAEMON keeps sync in-process: no background server, no port grab.
        env = {**os.environ, "AGENTSVIEW_DATA_DIR": AGENTSVIEW_DATA, "AGENTSVIEW_NO_DAEMON": "1"}
    else:
        cmd, env = ["agentsview", *args], None
    try:
        return subprocess.run(cmd, check=True, capture_output=True, text=True, env=env).stdout
    except subprocess.CalledProcessError as e:
        sys.exit(f"{cmd[0]} {args[0]} failed:\n{e.stderr[-2000:]}")
    except FileNotFoundError:
        sys.exit(f"{cmd[0]} not found; install it or rerun the Setup command from the tokenspend site")


def parse_ndjson(text):
    recs = [json.loads(line) for line in text.splitlines() if line.strip()]
    meta = [r for r in recs if r.get("type") == "meta"]
    sessions = [r for r in recs if "id" in r and "model_usage" in r]
    return meta, sessions


def remote_map(meta):
    """project_key -> normalized git remote, from export sessions meta records."""
    remotes = {}
    for rec in meta:
        for key, proj in rec.get("projects", {}).items():
            ident = proj.get("identity") or {}
            if ident.get("kind") == "git_remote" and ident.get("normalized_remote"):
                remotes[key] = ident["normalized_remote"]
    return remotes


def session_key(salt, session_id):
    return hmac.new(salt.encode(), session_id.encode(), hashlib.sha256).hexdigest()[:32]


def marker_code(cwd, cache):
    """Jobcode from the nearest .tokenspend file at or above cwd, else None.
    Walks up like git discovery, stopping at $HOME (or the filesystem root).
    Placing a marker is explicit consent to ship everything under it."""
    if not cwd:
        return None
    home = os.path.expanduser("~")
    code, walked, d = None, [], cwd
    while True:
        if d in cache:
            code = cache[d]
            break
        walked.append(d)
        try:
            with open(os.path.join(d, ".tokenspend")) as f:
                code = next((ln.strip() for ln in f if ln.strip()), None)
        except OSError:
            code = None
        parent = os.path.dirname(d)
        if code is not None or d == home or parent == d:
            break
        d = parent
    for w in walked:
        cache[w] = code
    return code


def session_cwds(ids):
    """id -> cwd for the given sessions, read straight from agentsview's SQLite
    (the only path-bearing source; `session list` would need a daemon). cwds
    never leave this machine; they are only used to find .tokenspend markers."""
    if not ids:
        return {}
    data_dir = AGENTSVIEW_DATA if private_agentsview() else os.environ.get(
        "AGENTSVIEW_DATA_DIR", os.path.expanduser("~/.agentsview"))
    db = os.path.join(data_dir, "sessions.db")
    if not os.path.exists(db):
        return {}
    # ponytail: the sessions(id, cwd) schema is pinned by the installer's agentsview version
    con = sqlite3.connect(f"file:{db}?mode=ro", uri=True)
    try:
        ids = list(ids)
        q = f"select id, cwd from sessions where cwd != '' and id in ({','.join('?' * len(ids))})"
        return dict(con.execute(q, ids))
    finally:
        con.close()


def snapshots(sessions, remotes, allow, salt, cwds=None):
    """Return (rows to ship, sessions kept home). Private by default."""
    rows, home, cache = [], [], {}
    cwds = cwds or {}
    for s in sessions:
        proj = s.get("project", {})
        code = marker_code(cwds.get(s["id"]), cache)
        if code == "-":
            # Reserved: a .tokenspend containing just "-" keeps everything
            # under it home, carving a private hole out of a marked tree
            # or an allow-listed repo.
            home.append(s)
            continue
        if code:
            # A .tokenspend marker overrides the git remote and needs no
            # allow match: writing the file is the consent.
            remote = f"tokenspend:{code}"
        else:
            remote = remotes.get(proj.get("project_key"))
            if not remote:
                # agentsview sometimes emits identity: None in the meta projects
                # map while the session record itself carries a resolved remote.
                ident = proj.get("identity") or {}
                if ident.get("kind") == "git_remote":
                    remote = ident.get("normalized_remote")
            if not (remote and any(fnmatch.fnmatch(remote, g) for g in allow)):
                home.append(s)
                continue
        key = session_key(salt, s["id"])
        for model, u in (s["model_usage"].get("by_model") or {}).items():
            rows.append({
                "session": key, "remote": remote, "model": model,
                **{out: u.get(src, 0) for src, out in TOKEN_FIELDS.items()},
                "microdollars": u.get("cost", {}).get("microdollars", 0),
            })
    rows.sort(key=lambda r: (r["remote"], r["session"], r["model"]))
    return rows, home


def build_payload(cfg, salt, today=None, lookback=LOOKBACK_DAYS):
    today = today or dt.date.today()  # the dev's local day, like the report's cutoffs
    since = (today - dt.timedelta(days=lookback)).isoformat()
    # The include flags are deliberate and not configurable: subagent, automated
    # (hooks, title generation) and one-shot sessions all cost money that belongs
    # to the client, and agentsview excludes them from exports by default.
    if private_agentsview():
        agentsview("sync")  # the private copy has no daemon watching; pull in new sessions first
    meta, sessions = parse_ndjson(agentsview(
        "export", "sessions", "--date-from", since, "--date-to", today.isoformat(),
        "--all", "--format", "ndjson",
        "--include-children", "--include-automated", "--include-one-shot"))
    rows, home = snapshots(sessions, remote_map(meta), cfg["allow"], salt,
                           cwds=session_cwds([s["id"] for s in sessions]))
    return {"developer": cfg["developer"], "snapshots": rows}, home, sessions


def _money(microdollars):
    return f"${microdollars / 1e6:,.4f}" if microdollars < 10_000 else f"${microdollars / 1e6:,.2f}"


def report(since, until, sessions, rows, home, salt):
    """Local gut-check in the style of `agentsview activity report`: what would
    ship (by remote/jobcode) versus what stays home (by project). Project labels
    are printed here only; they never enter the payload."""
    remote_of = {}
    for r in rows:
        remote_of[r["session"]] = r["remote"]
    lines = [f"Tokenspend {since} to {until} (sessions active in window; usage is each session's lifetime total)"]
    shipped = [s for s in sessions if session_key(salt, s["id"]) in remote_of]
    def cost(ss):  # same source the payload uses: per-model cost
        return sum(u.get("cost", {}).get("microdollars", 0)
                   for s in ss for u in (s["model_usage"].get("by_model") or {}).values())
    lines.append(f"Shipped    {len(shipped)} sessions  {_money(cost(shipped))}")
    lines.append(f"Kept home  {len(home)} sessions  {_money(cost(home))}")

    def group(title, ss, keyf):
        by = {}
        for s in ss:
            by.setdefault(keyf(s), []).append(s)
        lines.append(f"\n{title}:")
        for k, grp in sorted(by.items(), key=lambda kv: -cost(kv[1])):
            lines.append(f"  {k:<45} {len(grp):>3} sessions  {_money(cost(grp)):>10}")

    label = lambda s: s.get("project", {}).get("display_label") or "?"
    group("Shipped by remote", shipped, lambda s: remote_of[session_key(salt, s["id"])])
    group("Kept home by project", home, label)

    lines.append("\nSessions:")
    lines.append(f"  {'STATUS':<6} {'PROJECT':<24} {'AGENT':<8} {'MINUTES':>7} {'COST':>10}  REMOTE")
    for s in sorted(shipped + home, key=lambda s: -cost([s])):
        rem = remote_of.get(session_key(salt, s["id"]))
        lines.append(f"  {'ship' if rem else 'home':<6} {label(s)[:24]:<24} {s.get('agent', '?')[:8]:<8} "
                     f"{s.get('duration_seconds', 0) / 60:>7.1f} {_money(cost([s])):>10}  {rem or ''}")
    return "\n".join(lines) + "\n"


def post(cfg, payload):
    req = urllib.request.Request(
        cfg["endpoint"],
        data=json.dumps(payload).encode(),
        headers={"Content-Type": "application/json", "Authorization": f"Bearer {cfg['token']}"},
        method="POST",
    )
    with urllib.request.urlopen(req, timeout=30) as resp:
        return resp.status, json.loads(resp.read() or b"{}")


def ship_url(cfg):
    # The server serves its current ship.py (read from disk) next to the Setup
    # page; the config only carries the ingest endpoint, so derive it.
    return urllib.parse.urljoin(cfg["endpoint"], "/setup/ship.py")


def fetch_upstream(cfg):
    with urllib.request.urlopen(ship_url(cfg), timeout=30) as resp:
        return resp.read()


def check_for_updates(cfg, path):
    """(current, upstream_bytes). No version numbers: current means byte-identical
    to what the server serves right now, so it can never drift from reality."""
    upstream = fetch_upstream(cfg)
    with open(path, "rb") as f:
        return f.read() == upstream, upstream


def update(cfg, path):
    """Replace this file with the server's copy. Returns True if it changed."""
    current, upstream = check_for_updates(cfg, path)
    if current:
        return False
    # An error page, an empty body or a truncated download must never replace a
    # working shipper: it has to compile and it has to be this program.
    if b"\ndef main():" not in upstream:
        sys.exit("refusing to install the server's ship.py: not the shipper")
    try:
        compile(upstream, path, "exec")
    except (SyntaxError, ValueError) as e:
        sys.exit(f"refusing to install the server's ship.py: not valid Python ({e})")
    new = path + ".new"
    with open(new, "wb") as f:
        f.write(upstream)
    os.chmod(new, os.stat(path).st_mode & 0o777)
    os.replace(new, path)
    return True


def auto_update(cfg, path):
    """Replace this file with the server's copy before a shipping run, unless the
    config says "auto_update": false. Runs first so a shipper the server can no
    longer talk to still gets replaced. A new shipper re-execs itself so the
    run that fetched it already ships with it (TOKENSPEND_REEXEC stops a second
    hop). Then the private agentsview is brought up to the version the (new)
    shipper expects."""
    if cfg.get("auto_update", True) is False:
        return
    try:
        if not os.environ.get("TOKENSPEND_REEXEC") and update(cfg, path):
            print(f"updated {path}; restarting on the new shipper", file=sys.stderr)
            os.environ["TOKENSPEND_REEXEC"] = "1"
            os.execv(sys.executable, [sys.executable, path, *sys.argv[1:]])
    except Exception:  # ponytail: an update attempt must never stop the run that was about to ship
        pass
    try:
        if update_agentsview():
            print(f"updated {AGENTSVIEW_BIN} to v{AGENTSVIEW_VERSION}", file=sys.stderr)
    except Exception as e:  # ponytail: same rule; the old agentsview still ships
        print(f"agentsview update skipped: {e}", file=sys.stderr)


def _platform_asset():
    arch = {"arm64": "arm64", "aarch64": "arm64", "x86_64": "amd64", "amd64": "amd64"}[platform.machine().lower()]
    return f"{platform.system().lower()}_{arch}"


def agentsview_current(path=None):
    """True when the private agentsview at path reports AGENTSVIEW_VERSION.
    Installs without a private copy (agentsview from PATH) and copies in a
    directory the shipper cannot write to (the nix flake's, in the store) are
    never ours to update, so they count as current."""
    path = path or AGENTSVIEW_BIN
    if not os.path.exists(path) or not os.access(os.path.dirname(os.path.abspath(path)), os.W_OK):
        return True
    out = subprocess.run([path, "--version"], capture_output=True, text=True).stdout
    return f"v{AGENTSVIEW_VERSION}" in out.split()


def update_agentsview():
    """Replace the private agentsview when it is not the release this shipper
    was written against. Same download and checksum check as the installer;
    the new binary must report the expected version before it replaces the
    old one, so a bad download leaves a working install behind. Returns True
    if it changed."""
    if not private_agentsview() or agentsview_current():
        return False
    asset = f"agentsview_{AGENTSVIEW_VERSION}_{_platform_asset()}.tar.gz"
    base = f"{AGENTSVIEW_RELEASES}v{AGENTSVIEW_VERSION}/"
    with urllib.request.urlopen(base + asset, timeout=120) as resp:
        blob = resp.read()
    with urllib.request.urlopen(base + "SHA256SUMS", timeout=30) as resp:
        sums = resp.read().decode()
    if f"{hashlib.sha256(blob).hexdigest()}  {asset}" not in sums.splitlines():
        raise RuntimeError(f"checksum mismatch for {asset}")
    new = AGENTSVIEW_BIN + ".new"
    with tarfile.open(fileobj=io.BytesIO(blob)) as tar, open(new, "wb") as f:
        f.write(tar.extractfile("agentsview").read())
    os.chmod(new, 0o755)
    if not agentsview_current(new):
        os.remove(new)
        raise RuntimeError(f"downloaded {asset} does not report v{AGENTSVIEW_VERSION}")
    os.replace(new, AGENTSVIEW_BIN)
    return True


def nag(cfg, path):
    """One stderr line after a shipping run if the server has a newer shipper."""
    try:
        current, _ = check_for_updates(cfg, path)
    except Exception:  # ponytail: a nag must never fail the run that just shipped
        return
    if not current:
        print("update available: run tokenspend-ship --update", file=sys.stderr)


def backfill_days(today):
    """Lookback that reaches the 1st of today's month: what a first run wants.
    Older sessions belong to months that are already invoiced."""
    return today.day - 1


def main():
    ap = argparse.ArgumentParser(prog="tokenspend-ship", description=__doc__.splitlines()[0])
    ap.add_argument("--dry-run", action="store_true", help="print payload instead of posting")
    ap.add_argument("--backfill", action="store_true",
                    help="re-snapshot everything since the 1st of this month; run once after installing "
                         "or after placing a .tokenspend marker")
    ap.add_argument("--report", action="store_true",
                    help="print a local shipped-vs-home summary (like `agentsview activity report`) "
                         "instead of posting; defaults to today only, widen with --lookback")
    ap.add_argument("--lookback", type=int, default=LOOKBACK_DAYS, metavar="DAYS",
                    help=f"re-snapshot sessions active in the last DAYS days (default {LOOKBACK_DAYS}); "
                         "see --backfill for the usual first-run choice")
    ap.add_argument("--check-for-updates", action="store_true",
                    help="compare this file with the server's copy; exit 1 if an update is available")
    ap.add_argument("--update", action="store_true", help="replace this file with the server's copy")
    ap.add_argument("--no-nag", action="store_true", help="skip the update-available notice after shipping")
    a = ap.parse_args()
    cfg = load_config(dry_run=a.dry_run or a.report)
    if a.check_for_updates or a.update:
        # Before load_salt(): a check should never create a salt or sync agentsview.
        if a.update:
            print(f"updated {__file__} from {ship_url(cfg)}" if update(cfg, __file__) else "shipper already up to date")
            print(f"updated agentsview to v{AGENTSVIEW_VERSION}" if update_agentsview() else "agentsview already up to date")
            return
        if check_for_updates(cfg, __file__)[0] and agentsview_current():
            print("up to date")
            return
        print("update available: run tokenspend-ship --update")
        sys.exit(1)
    if not (a.dry_run or a.report):
        auto_update(cfg, __file__)
    salt = load_salt()
    if a.report and "--lookback" not in sys.argv:
        a.lookback = 0  # ponytail: a gut-check defaults to today; the shipper keeps its 3-day window
    if a.backfill:  # explicit, so it wins over the report default
        a.lookback = backfill_days(dt.date.today())
    payload, home, sessions = build_payload(cfg, salt, lookback=a.lookback)
    if a.report:
        today = dt.date.today()
        since = (today - dt.timedelta(days=a.lookback)).isoformat()
        print(report(since, today.isoformat(), sessions, payload["snapshots"], home, salt), end="")
        return
    n_sessions = len({r["session"] for r in payload["snapshots"]})
    print(f"shipping {n_sessions} sessions ({len(payload['snapshots'])} rows), "
          f"keeping {len(home)} sessions at home", file=sys.stderr)
    if a.dry_run:
        print(json.dumps(payload, indent=2))
        return
    status, body = post(cfg, payload)
    print(f"POST {cfg['endpoint']} -> {status} {body}", file=sys.stderr)
    if not a.no_nag:
        nag(cfg, __file__)


if __name__ == "__main__":
    main()