"""Dump all prod data to local Parquet files for offline analysis.

Run once:
    python docs/run1-14022026/dump_local_data.py

Creates Parquet files in docs/run1-14022026/data/ so agents never touch prod.
"""
import sys, os
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..", "notebooks"))

from utils.connection import init, query_df

conn = init()

DATA_DIR = os.path.join(os.path.dirname(__file__), "data")
os.makedirs(DATA_DIR, exist_ok=True)

TABLES = {
    "sessions": "SELECT * FROM sessions ORDER BY first_seen",
    "messages": "SELECT * FROM messages ORDER BY timestamp",
    "token_usage": "SELECT * FROM token_usage",
    "tool_calls": """
        SELECT tc.*, m.session_id, m.timestamp
        FROM tool_calls tc
        JOIN messages m ON m.id = tc.message_id
        ORDER BY m.timestamp
    """,
    "tool_results": """
        SELECT tr.*, m.session_id
        FROM tool_results tr
        JOIN messages m ON m.id = tr.message_id
    """,
}

import json as _json

# Columns that are JSONB in Postgres — serialize to JSON strings for Parquet
JSONB_COLS = {"tool_input", "raw_data", "source_stats", "recent_errors"}

for name, sql in TABLES.items():
    print(f"Dumping {name}...", end=" ", flush=True)
    df = query_df(conn, sql)
    for col in JSONB_COLS & set(df.columns):
        df[col] = df[col].apply(lambda v: _json.dumps(v) if v is not None else None)
    path = os.path.join(DATA_DIR, f"{name}.parquet")
    df.to_parquet(path, index=False)
    print(f"{len(df)} rows -> {path}")

conn.close()
print("\nDone. All data saved to", DATA_DIR)
