#!/usr/bin/env python3
"""AI 分帳：Claude Code 用量上傳程式（只用 Python 內建模組）

讀這台電腦 ~/.claude/projects 的紀錄，每 5 分鐘 × 模型加總成 token 與 API 等值美金，
上傳到 AI 分帳。只傳數字，不傳對話內容、專案名稱或檔案路徑。

  python ai_collector.py setup <邀請連結或群組碼> <你的名字> [這台電腦的名字]
  python ai_collector.py            # 跑一次（排程會每 15 分鐘呼叫）
  python ai_collector.py install    # 掛排程（Windows 工作排程器／macOS launchd）
  python ai_collector.py uninstall  # 拿掉排程
  python ai_collector.py status     # 看設定與上次結果
"""
import json
import os
import platform
import re
import subprocess
import sys
import urllib.error
import urllib.request
from datetime import datetime, timedelta, timezone
from pathlib import Path

VERSION = "1.0.0"
SB_URL = "https://wfvvqadfcgbganxgiskz.supabase.co"
SB_KEY = "sb_publishable__mR2qJp81SN3rK3EzyubMA_NJvXJipf"
HOME = Path.home()
CONF = HOME / ".ai-split.json"
ROOT = HOME / ".claude" / "projects"
TASK = "AI-Split-Collector"
LOOKBACK_DAYS = 3

# 每百萬 token 美金：(輸入, 輸出, 快取讀取)；快取寫入 5 分鐘＝輸入×1.25、1 小時＝輸入×2
PRICES = {
    "claude-fable-5-1": (10.0, 50.0, 0.25),
    "claude-fable-5": (10.0, 50.0, 1.00),
    "claude-opus-5-5": (4.0, 20.0, 0.20),
    "claude-opus-5": (5.0, 25.0, 0.50),
    "claude-opus-4-8": (5.0, 25.0, 0.50),
    "claude-opus-4-7": (5.0, 25.0, 0.50),
    "claude-opus-4-6": (5.0, 25.0, 0.50),
    "claude-sonnet-5-5": (2.0, 10.0, 0.20),
    "claude-sonnet-5": (2.0, 10.0, 0.20),
    "claude-sonnet-4-6": (3.0, 15.0, 0.30),
    "claude-sonnet-4-5": (3.0, 15.0, 0.30),
    "claude-haiku-5-5": (1.0, 5.0, 0.10),
    "claude-haiku-4-5": (1.0, 5.0, 0.10),
}
FAST_MULT = {"claude-opus-5-5": 2.0, "claude-opus-5": 2.0, "claude-opus-4-8": 2.0}
FALLBACK = (3.0, 15.0, 0.30)


def rpc(fn, args):
    req = urllib.request.Request(
        f"{SB_URL}/rest/v1/rpc/{fn}", data=json.dumps(args).encode(), method="POST",
        headers={"apikey": SB_KEY, "content-type": "application/json"})
    try:
        with urllib.request.urlopen(req, timeout=30) as r:
            body = r.read().decode()
            return json.loads(body) if body else None
    except urllib.error.HTTPError as e:
        msg = e.read().decode(errors="replace")
        try:
            msg = json.loads(msg).get("message", msg)
        except ValueError:
            pass
        raise RuntimeError(msg) from None


def load_conf():
    try:
        return json.loads(CONF.read_text(encoding="utf-8"))
    except (OSError, ValueError):
        return {}


def save_conf(c):
    CONF.write_text(json.dumps(c, ensure_ascii=False, indent=1), encoding="utf-8")
    try:
        os.chmod(CONF, 0o600)
    except OSError:
        pass


def norm_model(m):
    m = re.sub(r"\[.*\]$", "", m or "")
    return re.sub(r"-\d{8}$", "", m)


def cost(model, u):
    inp, out, cr = PRICES.get(model, FALLBACK)
    cc = u.get("cache_creation") or {}
    w5, w1 = cc.get("ephemeral_5m_input_tokens"), cc.get("ephemeral_1h_input_tokens")
    if w5 is None and w1 is None:
        w5, w1 = u.get("cache_creation_input_tokens", 0), 0
    usd = (u.get("input_tokens", 0) * inp + u.get("output_tokens", 0) * out
           + (w5 or 0) * inp * 1.25 + (w1 or 0) * inp * 2.0
           + u.get("cache_read_input_tokens", 0) * cr) / 1e6
    if u.get("speed") == "fast":
        usd *= FAST_MULT.get(model, 1.0)
    tokens = (u.get("input_tokens", 0) + u.get("output_tokens", 0) + (w5 or 0) + (w1 or 0)
              + u.get("cache_read_input_tokens", 0))
    return usd, tokens


def current_org():
    """Claude Code 目前登入的帳號（organization uuid）。讀不到回 None。"""
    for p in (HOME / ".claude.json", HOME / ".claude" / ".claude.json"):
        try:
            d = json.loads(p.read_text(encoding="utf-8"))
            return (d.get("oauthAccount") or {}).get("organizationUuid")
        except (OSError, ValueError):
            continue
    return None


def collect(since):
    """回傳 {(5 分鐘格, 模型): [n, usd, tokens]}"""
    seen = {}
    if not ROOT.exists():
        return {}
    cutoff_mtime = since.timestamp() - 3600
    for f in ROOT.rglob("*.jsonl"):
        try:
            if f.stat().st_mtime < cutoff_mtime:
                continue
            fh = f.open(encoding="utf-8", errors="replace")
        except OSError:
            continue
        with fh:
            for line in fh:
                if '"usage"' not in line:
                    continue
                try:
                    d = json.loads(line)
                except ValueError:
                    continue
                m = d.get("message") or {}
                u = m.get("usage")
                if d.get("type") != "assistant" or not isinstance(u, dict):
                    continue
                model = norm_model(m.get("model"))
                if not model or model == "<synthetic>":
                    continue
                ts = d.get("timestamp")
                if not ts:
                    continue
                try:
                    t = datetime.fromisoformat(ts.replace("Z", "+00:00"))
                except ValueError:
                    continue
                if t < since:
                    continue
                seen[(m.get("id"), d.get("requestId"))] = (t, model, u)
    buckets = {}
    for t, model, u in seen.values():
        b = t.replace(minute=t.minute - t.minute % 5, second=0, microsecond=0)
        usd, tok = cost(model, u)
        x = buckets.setdefault((b, model), [0, 0.0, 0])
        x[0] += 1
        x[1] += usd
        x[2] += tok
    return buckets


def run(quiet=False):
    c = load_conf()
    if not c.get("token"):
        sys.exit("還沒設定。先跑：python ai_collector.py setup <邀請連結> <你的名字>")
    org = current_org()
    status = {"v": VERSION, "kind": "cli", "os": platform.system(), "host": platform.node()[:40]}
    if c.get("claude_org") and org and org != c["claude_org"]:
        status["org_mismatch"] = True
        rpc("ai_push", {"p_token": c["token"], "p_events": [], "p_snaps": [], "p_status": status})
        c["last"] = {"at": datetime.now().isoformat(timespec="seconds"), "note": "Claude Code 登入的不是共用帳號，這次不上傳"}
        save_conf(c)
        if not quiet:
            print(c["last"]["note"])
        return
    since = datetime.now(timezone.utc) - timedelta(days=LOOKBACK_DAYS)
    # 第一次只從設定那天開始算，之前的用量不屬於這個群組
    if c.get("since"):
        since = max(since, datetime.fromisoformat(c["since"]))
    buckets = collect(since)
    events = [{"service": "claude_code", "model": model, "ts": b.isoformat(), "n": v[0],
               "usd": round(v[1], 6), "tokens": v[2], "key": f"cc:{b.strftime('%Y%m%dT%H%M')}:{model}"}
              for (b, model), v in sorted(buckets.items())]
    sent = 0
    for i in range(0, max(len(events), 1), 1000):
        chunk = events[i:i + 1000]
        rpc("ai_push", {"p_token": c["token"], "p_events": chunk, "p_snaps": [], "p_status": status})
        sent += len(chunk)
    usd = sum(v[1] for v in buckets.values())
    c["last"] = {"at": datetime.now().isoformat(timespec="seconds"), "rows": sent, "usd_3d": round(usd, 2)}
    save_conf(c)
    if not quiet:
        print(f"已上傳 {sent} 格（最近 {LOOKBACK_DAYS} 天 API 等值 US${usd:.2f}）")


def setup(args):
    if len(args) < 2:
        sys.exit("用法：python ai_collector.py setup <邀請連結或群組碼> <你的名字> [這台電腦的名字]")
    m = re.search(r"[0-9a-f]{32}", args[0])
    if not m:
        sys.exit("找不到群組碼，請貼總覽網頁「設定 → 裝置」裡的邀請連結")
    label = args[2] if len(args) > 2 else f"{platform.system()} {platform.node()[:20]} Claude Code"
    r = rpc("ai_register", {"p_code": m.group(0), "p_name": args[1], "p_kind": "cli", "p_label": label})
    c = {"token": r["token"], "device_id": r["device_id"], "member": args[1], "label": label,
         "claude_org": r.get("claude_org") or "", "since": datetime.now(timezone.utc).isoformat()}
    save_conf(c)
    org = current_org()
    print(f"已連結：{args[1]}（{label}）")
    if c["claude_org"] and org and org != c["claude_org"]:
        print("⚠ 這台 Claude Code 目前登入的不是共用帳號，換成共用帳號後才會開始上傳")
    run()
    print("下一步：python ai_collector.py install  （每 15 分鐘自動上傳）")


def install():
    py = Path(sys.executable)
    me = Path(__file__).resolve()
    if os.name == "nt":
        pyw = py.with_name("pythonw.exe")
        exe = pyw if pyw.exists() else py
        # 用 ScheduledTask cmdlet：schtasks 建的工作預設沒插電不跑、錯過也不補
        ps = (
            f"$a = New-ScheduledTaskAction -Execute '{exe}' -Argument '\"{me}\" --quiet';"
            "$t = New-ScheduledTaskTrigger -Once -At (Get-Date).AddMinutes(1) -RepetitionInterval (New-TimeSpan -Minutes 15);"
            "$s = New-ScheduledTaskSettingsSet -AllowStartIfOnBatteries -DontStopIfGoingOnBatteries -StartWhenAvailable "
            "-ExecutionTimeLimit (New-TimeSpan -Minutes 10) -MultipleInstances IgnoreNew;"
            f"Register-ScheduledTask -TaskName '{TASK}' -Action $a -Trigger $t -Settings $s -Force | Out-Null"
        )
        subprocess.run(["powershell", "-NoProfile", "-Command", ps], check=True)
        print(f"已掛排程「{TASK}」，每 15 分鐘上傳一次")
    elif sys.platform == "darwin":
        plist = HOME / "Library" / "LaunchAgents" / "dev.ai-split.collector.plist"
        plist.parent.mkdir(parents=True, exist_ok=True)
        plist.write_text(f"""<?xml version="1.0" encoding="UTF-8"?>
<!DOCTYPE plist PUBLIC "-//Apple//DTD PLIST 1.0//EN" "http://www.apple.com/DTDs/PropertyList-1.0.dtd">
<plist version="1.0"><dict>
<key>Label</key><string>dev.ai-split.collector</string>
<key>ProgramArguments</key><array><string>{py}</string><string>{me}</string><string>--quiet</string></array>
<key>StartInterval</key><integer>900</integer>
<key>RunAtLoad</key><true/>
<key>StandardErrorPath</key><string>{HOME}/.ai-split.log</string>
</dict></plist>
""", encoding="utf-8")
        subprocess.run(["launchctl", "unload", str(plist)], capture_output=True)
        subprocess.run(["launchctl", "load", str(plist)], check=True)
        print("已掛 launchd，每 15 分鐘上傳一次")
    else:
        print(f"Linux 請自己加 crontab：*/15 * * * * {py} {me} --quiet")


def uninstall():
    if os.name == "nt":
        subprocess.run(["powershell", "-NoProfile", "-Command",
                        f"Unregister-ScheduledTask -TaskName '{TASK}' -Confirm:$false"], check=False)
    elif sys.platform == "darwin":
        plist = HOME / "Library" / "LaunchAgents" / "dev.ai-split.collector.plist"
        subprocess.run(["launchctl", "unload", str(plist)], capture_output=True)
        plist.unlink(missing_ok=True)
    print("已拿掉排程")


def status():
    c = load_conf()
    if not c:
        print("還沒設定")
        return
    print(f"成員：{c.get('member')}　裝置：{c.get('label')}")
    print(f"上次：{c.get('last')}")
    org = current_org()
    print("Claude Code 帳號：" + ("共用帳號 ✓" if org == c.get("claude_org") else f"不是共用帳號（{org}）"))


def main():
    if hasattr(sys.stdout, "reconfigure"):
        sys.stdout.reconfigure(encoding="utf-8", errors="replace")
    a = [x for x in sys.argv[1:] if x != "--quiet"]
    quiet = "--quiet" in sys.argv
    cmd = a[0] if a else "run"
    try:
        if cmd == "setup":
            setup(a[1:])
        elif cmd == "install":
            install()
        elif cmd == "uninstall":
            uninstall()
        elif cmd == "status":
            status()
        elif cmd == "run":
            run(quiet)
        else:
            print(__doc__)
    except RuntimeError as e:
        msg = "這台裝置已被移除，請重新 setup" if "bad token" in str(e) else str(e)
        if quiet:
            (HOME / ".ai-split.log").open("a", encoding="utf-8").write(f"{datetime.now():%F %T} {msg}\n")
        sys.exit(msg)


if __name__ == "__main__":
    main()