#!/usr/bin/env python3
"""
usage_collector.py — 多 Provider 用量采集主控
=============================================
编排所有 provider 的采集，合并到统一的 usage_stats.json。

Provider:
  - opencode_go: 本地 HTTP/CDP 采集（usage_collector_ocg.py）
  - kimi: SSH 远程调用 Windows 采集器（usage_collector_kimi.py）

输出: gateway/temp/usage_stats.json（Dashboard /api/usage 读取）
"""
import json, os, sys, subprocess
from pathlib import Path
from datetime import datetime, timezone

SCRIPT_DIR = Path(__file__).resolve().parent
GATEWAY_DIR = SCRIPT_DIR.parent
TEMP_DIR = GATEWAY_DIR / "temp"
LOGS_DIR = GATEWAY_DIR / "logs"
os.makedirs(str(TEMP_DIR), exist_ok=True)
os.makedirs(str(LOGS_DIR), exist_ok=True)

OUTPUT_FILE = TEMP_DIR / "usage_stats.json"
WINDOWS_HOST = "192.168.1.16"
WINDOWS_PYTHON = r"C:\Users\hmo\AppData\Local\Programs\Python\Python310\python.exe"
WINDOWS_SCRIPT = r"D:\F\NewI\opencode\daily-workspace\projects\AgentsMeeting\gateway\scripts\usage_collector_kimi.py"


def log(msg):
    ts = datetime.now().strftime("%H:%M:%S")
    line = f"[{ts}] {msg}"
    print(line, flush=True)


def collect_opencode_go():
    """运行本地 OCG 采集器。"""
    log("OCG: starting...")
    ocg_script = SCRIPT_DIR / "usage_collector_ocg.py"
    if not ocg_script.exists():
        # Fallback: old name
        ocg_script = SCRIPT_DIR / "usage_collector.py"

    try:
        r = subprocess.run(
            [sys.executable, str(ocg_script)],
            capture_output=True, text=True, timeout=120,
            cwd=str(SCRIPT_DIR),
        )
        if r.returncode != 0 and r.stderr:
            log(f"OCG: FAILED (rc={r.returncode})")
            log(f"OCG stderr: {r.stderr[-500:]}")
        else:
            log(f"OCG: done (rc={r.returncode})")
    except Exception as e:
        log(f"OCG: EXCEPTION: {e}")


def collect_kimi():
    """SSH 到 Windows 运行 Kimi 采集器，scp 结果回来。"""
    log("Kimi: SSH to Windows...")
    try:
        r = subprocess.run(
            ["ssh", "-o", "StrictHostKeyChecking=no", "-o", "ConnectTimeout=10",
             "-o", "BatchMode=yes", f"hmo@{WINDOWS_HOST}",
             f'"{WINDOWS_PYTHON}" "{WINDOWS_SCRIPT}"'],
            capture_output=True, text=True, timeout=45,
        )
        for line in r.stdout.splitlines()[-5:]:
            log(f"Kimi: {line.strip()}")
        if r.returncode != 0:
            log(f"Kimi: SSH FAILED (rc={r.returncode}): {r.stderr[-200:]}")
            # Try scp alternative
            r2 = subprocess.run(
                ["scp", "-o", "StrictHostKeyChecking=no", "-o", "ConnectTimeout=10",
                 f"hmo@{WINDOWS_HOST}:{WINDOWS_SCRIPT}", "/dev/null"],
                capture_output=True, timeout=10,
            )
            if r2.returncode == 0:
                log("Kimi: SSH failed but scp works — check script path")
    except subprocess.TimeoutExpired:
        log("Kimi: SSH TIMEOUT")
    except Exception as e:
        log(f"Kimi: SSH EXCEPTION: {e}")

    # Try to scp the output file from Windows
    kimi_output = TEMP_DIR / "usage_stats_kimi.json"
    try:
        r = subprocess.run(
            ["scp", "-o", "StrictHostKeyChecking=no", "-o", "ConnectTimeout=10",
             f"hmo@{WINDOWS_HOST}:" + WINDOWS_SCRIPT.replace("usage_collector_kimi.py", "") + "../temp/usage_stats_kimi.json",
             str(kimi_output)],
            capture_output=True, text=True, timeout=15,
        )
        if kimi_output.exists():
            log(f"Kimi: output scp'd ({kimi_output.stat().st_size} bytes)")
        else:
            log("Kimi: scp failed or file empty")
    except Exception as e:
        log(f"Kimi: scp EXCEPTION: {e}")


def merge_all():
    """合并所有 provider 的数据到 usage_stats.json。"""
    accounts = []

    # 1. Read existing OCG cache (from usage_collector_ocg output)
    #    The OCG collector writes directly to usage_stats.json or a cache
    cache_file = TEMP_DIR / "usage_stats.json"
    if cache_file.exists():
        try:
            with open(str(cache_file), "r", encoding="utf-8") as f:
                cached = json.load(f)
            for a in cached.get("accounts", []):
                a["provider"] = a.get("provider", "opencode_go")
                accounts.append(a)
            log(f"Merge: {len(accounts)} OCG accounts from cache")
        except Exception as e:
            log(f"Merge: failed to read cache: {e}")

    # 2. Read Kimi output
    kimi_file = TEMP_DIR / "usage_stats_kimi.json"
    if kimi_file.exists():
        try:
            with open(str(kimi_file), "r", encoding="utf-8") as f:
                kimi_data = json.load(f)
            for a in kimi_data.get("accounts", []):
                a["provider"] = a.get("provider", "kimi")
                # Check if this key already exists (update instead of duplicate)
                existing = [i for i, x in enumerate(accounts) if x.get("key_id") == a.get("key_id")]
                if existing:
                    accounts[existing[0]] = a
                else:
                    accounts.append(a)
            log(f"Merge: +{len(kimi_data.get('accounts',[]))} Kimi accounts")
        except Exception as e:
            log(f"Merge: failed to read Kimi data: {e}")

    # 3. Count by provider
    providers = {}
    for a in accounts:
        p = a.get("provider", "unknown")
        providers[p] = providers.get(p, 0) + 1

    output = {
        "ok": True,
        "last_refresh_iso": datetime.now(timezone.utc).isoformat(timespec="seconds"),
        "account_count": len(accounts),
        "collected_count": sum(1 for a in accounts if a.get("error") is None),
        "providers": providers,
        "accounts": accounts,
    }

    with open(str(OUTPUT_FILE), "w", encoding="utf-8") as f:
        json.dump(output, f, ensure_ascii=False, indent=2)
    log(f"Merge: saved {len(accounts)} accounts ({providers}) to {OUTPUT_FILE}")


def main():
    log("=== Multi-provider collection started ===")

    # Run collectors in sequence (OCG takes longest, Kimi needs CDP sync)
    collect_opencode_go()
    collect_kimi()

    # Merge
    merge_all()

    log("=== Collection complete ===")


if __name__ == "__main__":
    main()