"""Read-only export of the minimax-rp data from production D1 (SELECT only), paged per run.
Run from a shell that has sourced cfenv.sh. Writes raw/export/<table>.jsonl and raw/export/queries.sql (every query issued)."""
import json, os, subprocess, sys
A = '${RPB_AUDIT_DIR}'
W = '${PROJECT_ROOT}/worker'
OUT = f'{A}/raw/export'; os.makedirs(OUT, exist_ok=True)
QLOG = open(f'{OUT}/queries.sql', 'a')

def q(sql):
    assert sql.lstrip().upper().startswith('SELECT'), sql
    QLOG.write(sql + ';\n'); QLOG.flush()
    for attempt in range(4):
        p = subprocess.run(['npx', 'wrangler', 'd1', 'execute', 'rp-bench-db', '--remote', '--json', '--command', sql], cwd=W, capture_output=True, text=True)
        if p.returncode == 0:
            return json.loads(p.stdout)[0]['results']
        err = p.stderr[-500:]
    raise RuntimeError(err)

def paged(name, sql_fn, page):
    rows, off = [], 0
    while True:
        r = q(sql_fn(page, off)); rows += r
        if len(r) < page: return rows
        off += page

runs = q("SELECT id FROM runs WHERE benchmark_id = 'minimax-rp' ORDER BY created_at")
ids = [r['id'] for r in runs]
which = sys.argv[1:] or ['sessions', 'judgments', 'simmodels', 'aimsg', 'logs', 'batches']
for table in which:
    f = open(f'{OUT}/{table}.jsonl', 'w')
    for rid in ids:
        imported = rid.startswith('mm-')
        if table == 'sessions':
            rows = paged(table, lambda n, o: f"SELECT id, run_id, seed_id, lang, sample_idx, status, attempts, substr(error,1,300) AS error, result_json, updated_at FROM sessions WHERE run_id = '{rid}' ORDER BY id LIMIT {n} OFFSET {o}", 150)
        elif table == 'judgments':
            if imported: continue
            rows = paged(table, lambda n, o: f"SELECT j.session_id, j.chunk_idx, j.judge_key, j.judge_model, j.status, json_extract(j.scores_json, '$.dims') AS dims, substr(j.error,1,300) AS error, j.completion_tokens, j.created_at FROM judgments j JOIN sessions s ON s.id = j.session_id WHERE s.run_id = '{rid}' ORDER BY j.session_id, j.chunk_idx, j.judge_key LIMIT {n} OFFSET {o}", 1500)
        elif table == 'simmodels':
            if imported: continue
            rows = q(f"SELECT m.session_id, m.source, m.model, COUNT(*) AS n, MIN(m.round) AS first_round FROM messages m JOIN sessions s ON s.id = m.session_id WHERE s.run_id = '{rid}' AND m.role = 'user' AND m.round >= 2 GROUP BY m.session_id, m.source, m.model")
        elif table == 'aimsg':
            if imported: continue
            rows = q(f"SELECT m.session_id, m.source, m.model, m.provider, m.finish_reason, COUNT(*) AS n FROM messages m JOIN sessions s ON s.id = m.session_id WHERE s.run_id = '{rid}' AND m.role = 'ai' AND m.round >= 2 GROUP BY m.session_id, m.source, m.model, m.provider, m.finish_reason")
        elif table == 'logs':
            rows = paged(table, lambda n, o: f"SELECT id, run_id, ts, level, substr(message,1,400) AS message FROM run_logs WHERE run_id = '{rid}' ORDER BY id LIMIT {n} OFFSET {o}", 1000)
        elif table == 'batches':
            rows = q(f"SELECT id, run_id, judge_key, model, status, json_array_length(items_json) AS items, provider_status, substr(error,1,200) AS error, cost, reserved, created_at, updated_at, polled_at FROM judge_batches WHERE run_id = '{rid}'")
        for r in rows:
            r.setdefault('run_id', rid); f.write(json.dumps(r, ensure_ascii=False) + '\n')
        print(table, rid, len(rows), flush=True)
    f.close()