"""Replicate prod data into local Postgres Docker container.

Prerequisites:
    docker compose -f docs/run1-14022026/docker-compose.yml up -d

Run:
    .venv/bin/python docs/run1-14022026/setup_local_db.py
"""
import sys, os, json, math
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..", "notebooks"))

import psycopg
from psycopg.types.json import Jsonb
from utils.connection import init as prod_init

LOCAL_DSN = "postgresql://localdev:localdev@localhost:15432/qc_trace"
SCHEMA_PATH = os.path.join(os.path.dirname(__file__), "..", "..", "qc_trace", "db", "schema.sql")

print("Connecting to prod...", flush=True)
prod = prod_init()

print("Connecting to local...", flush=True)
local = psycopg.connect(LOCAL_DSN, autocommit=True)

print("Creating schema...", flush=True)
with open(SCHEMA_PATH) as f:
    schema_sql = f.read()
with local.cursor() as cur:
    cur.execute(schema_sql)

# Disable FK checks during load
with local.cursor() as cur:
    cur.execute("SET session_replication_role = 'replica'")

TABLES = ["sessions", "messages", "token_usage", "tool_calls", "tool_results"]

# JSONB columns per table
JSONB_COLS = {"raw_data", "tool_input", "source_stats", "recent_errors"}
# Integer columns that might come as float from pandas
INT_COLS = {"raw_line_number", "input_tokens", "output_tokens", "cached_tokens",
            "thinking_tokens", "source_schema_version", "queue_size", "messages_this_session"}

BATCH = 1000

for table in TABLES:
    print(f"\n{table}:", end=" ", flush=True)

    # Clear
    with local.cursor() as cur:
        cur.execute(f"DELETE FROM {table}")

    # Fetch from prod
    with prod.cursor() as cur:
        cur.execute(f"SELECT * FROM {table}")
        cols = [d.name for d in cur.description]
        total = 0
        while True:
            rows = cur.fetchmany(BATCH)
            if not rows:
                break

            # Build insert
            placeholders = ", ".join(["%s"] * len(cols))
            col_list = ", ".join(cols)
            sql = f"INSERT INTO {table} ({col_list}) VALUES ({placeholders}) ON CONFLICT DO NOTHING"

            clean_rows = []
            for row in rows:
                clean = []
                for col_name, val in zip(cols, row):
                    if col_name in INT_COLS and val is not None:
                        try:
                            val = int(val) if not (isinstance(val, float) and math.isnan(val)) else None
                        except (ValueError, TypeError):
                            val = None
                    if col_name in JSONB_COLS and isinstance(val, (dict, list)):
                        val = Jsonb(val)
                    clean.append(val)
                clean_rows.append(tuple(clean))

            with local.cursor() as lcur:
                lcur.executemany(sql, clean_rows)

            total += len(rows)
            print(f"{total}", end=" ", flush=True)

    print("done")

# Re-enable FK checks
with local.cursor() as cur:
    cur.execute("SET session_replication_role = 'origin'")

prod.close()
local.close()
print(f"\nAll data replicated to local Postgres at {LOCAL_DSN}")
