# sync_blackbox_flow.py — BlackBox API sync to PostgreSQL
# Fetches options flow data from BlackBox API and syncs to PostgreSQL database

import os, time, hashlib, re, argparse
from pathlib import Path
from typing import Dict, Optional, List
from datetime import datetime, date
import pandas as pd
import numpy as np
import requests
import json
import psycopg2
from psycopg2.extras import execute_values
from psycopg2 import sql
from dotenv import load_dotenv

load_dotenv()

# ─────────────────────────────────────────────
# Config
# ─────────────────────────────────────────────
WRITE_POSTGRES    = True

# BlackBox API Configuration
BLACKBOX_API_URL = "https://api.blackboxstocks.com/api/v2/options/getFlowMobile"
BLACKBOX_API_TOKEN = os.getenv("BLACKBOX_API_TOKEN", "eyJhbGciOiJodHRwOi8vd3d3LnczLm9yZy8yMDAxLzA0L3htbGRzaWctbW9yZSNobWFjLXNoYTI1NiIsInR5cCI6IkpXVCJ9.eyJodHRwOi8vc2NoZW1hcy54bWxzb2FwLm9yZy93cy8yMDA1LzA1L2lkZW50aXR5L2NsYWltcy9uYW1lIjoic3JlZGR5MDgxOUBnbWFpbC5jb20iLCJodHRwOi8vc2NoZW1hcy54bWxzb2FwLm9yZy93cy8yMDA1LzA1L2lkZW50aXR5L2NsYWltcy9uYW1laWRlbnRpZmllciI6Ijk3ZDA5NThjLTBhOTktNDA5OC05OWQwLWNkYTg3Njg5N2Q3ZiIsImh0dHA6Ly9zY2hlbWFzLnhtbHNvYXAub3JnL3dzLzIwMDUvMDUvaWRlbnRpdHkvY2xhaW1zL3NpZCI6ImRkODZTenZOMm9hajliT0dHNW1lajMwK2VaRDZsRkdvMU56bjc5MldBeVE9IiwiaHR0cDovL3NjaGVtYXMubWljcm9zb2Z0LmNvbS93cy8yMDA4LzA2L2lkZW50aXR5L2NsYWltcy9yb2xlIjpbIlN1YnNjcmliZXIiLCJPcHRpb25zIiwiRXF1aXRpZXMiXSwiZXhwIjoxNzY2MTU5Njc5fQ.8XPCxrjfRDLSbG_vcNNo59EO2OLsBZvNh-J8MmKHmgU")

# PostgreSQL connection parameters
POSTGRES_HOST     = os.getenv("POSTGRES_HOST", "192.168.8.151")
POSTGRES_PORT     = int(os.getenv("POSTGRES_PORT", "5432"))
POSTGRES_DB       = os.getenv("POSTGRES_DB", "institutional_trader")
POSTGRES_USER     = os.getenv("POSTGRES_USER", "postgres")
POSTGRES_PASSWORD = os.getenv("POSTGRES_PASSWORD", "postgres")
POSTGRES_SCHEMA   = os.getenv("POSTGRES_SCHEMA", "public")

# ⚠️ Per-table schema style: 'snake' or 'camel'
TABLE_STYLE = {
    "AlertStream":           "snake",
    "AlertStream_monthly":   "snake",
    "OptionsFlow":           "camel",   # ← your Supabase columns are CamelCase here
    "OptionsFlow_monthly":   "camel",   # ← same
    "OptionsVolume":         "camel",   # adjust if needed
    "Short_Long":            "camel",   # adjust if needed
}

_TARGET_STEMS = set(TABLE_STYLE.keys())

# Upsert behavior / logging
INCREMENTAL          = True
SINGLE_CHUNK         = False      # chunked to avoid 57014
UPSERT_CHUNK         = 3000       # drop to 1000 if still timeouts
MAX_RETRIES          = 3
PRINT_PROGRESS       = False
PRINT_PRUNE_LOGS     = False
USE_RETURNING_MINIMAL= True

print_flush = lambda *a, **k: print(*a, **k, flush=True)
def _ts(): return time.strftime("%Y-%m-%d %H:%M:%S")

# ─────────────────────────────────────────────
# PostgreSQL connection
# ─────────────────────────────────────────────
def _pg_conn():
    """Create a PostgreSQL connection."""
    return psycopg2.connect(
        host=POSTGRES_HOST,
        port=POSTGRES_PORT,
        database=POSTGRES_DB,
        user=POSTGRES_USER,
        password=POSTGRES_PASSWORD,
        options=f"-c search_path={POSTGRES_SCHEMA}"
    )

# ─────────────────────────────────────────────
# Expected headers (both styles) + mappings
# ─────────────────────────────────────────────
# OptionsFlow (CamelCase)
EXPECTED_OPTIONSFLOW_CAMEL = [
    "CreatedDate","CreatedTime","Symbol","Type","Volume","Price","Side",
    "CallPut","Strike","Spot","Premium","ExpirationDate","Color",
    "ImpliedVolatility","Dte","ER","StockEtf","Sector","Uoa",
    "Weekly","MktCap","OI"
]
# OptionsFlow (snake_case)
EXPECTED_OPTIONSFLOW_SNAKE = [
    "created_date","created_time","symbol","type","volume","price","side",
    "callput","strike","spot","premium","expiration_date","color",
    "implied_volatility","dte","er","stock_etf","sector","uoa",
    "weekly","mktcap","oi"
]
# AlertStream (CamelCase in files → snake in DB)
EXPECTED_ALERTSTREAM_SNAKE = [
    "date","timestamp","ticker","volume","price","pct_of_avg30",
    "notional","message","type","securitytype","industry","sector",
    "avg30day","float","earningsdate"
]

# Camel→snake renames (if file arrives CamelCase but DB expects snake)
MAP_OPTIONSFLOW_SNAKE = {
  "CreatedDate":"created_date","CreatedTime":"created_time","Symbol":"symbol","Type":"type",
  "Volume":"volume","Price":"price","Side":"side","CallPut":"callput","Strike":"strike",
  "Spot":"spot","Premium":"premium","ExpirationDate":"expiration_date","Color":"color",
  "ImpliedVolatility":"implied_volatility","Dte":"dte","ER":"er","StockEtf":"stock_etf",
  "Sector":"sector","Uoa":"uoa","Weekly":"weekly","MktCap":"mktcap","OI":"oi"
}
MAP_ALERTSTREAM_SNAKE = {
  "Date":"date","Timestamp":"timestamp","Ticker":"ticker","Volume":"volume","Price":"price",
  "Pct_of_Avg30Day":"pct_of_avg30","Notional":"notional","Message":"message","Type":"type",
  "SecurityType":"securitytype","Industry":"industry","Sector":"sector",
  "Avg30Day":"avg30day","Float":"float","EarningsDate":"earningsdate"
}

# snake→Camel renames (if file is snake but DB expects Camel)
MAP_OPTIONSFLOW_CAMEL = {v:k for k,v in MAP_OPTIONSFLOW_SNAKE.items()}

# ─────────────────────────────────────────────
# BlackBox API Integration
# ─────────────────────────────────────────────
def build_filter_bitmask(_filter_options: Optional[Dict] = None) -> int:
    """Build the filter bitmask based on filter options."""
    # Default filter value that enables common filters
    return 2198487171967

def build_filters(start_date: datetime, end_date: datetime, custom_filters: Optional[Dict] = None) -> Dict:
    """Build default filters object."""
    date_str = start_date.isoformat()
    
    filters = {
        "optionsDate": {
            "start": date_str,
            "end": end_date.isoformat()
        },
        "expireOptionsDate": {
            "start": date_str,
            "end": end_date.isoformat()
        },
        "optionsFlowPuts": True,
        "optionsFlowCalls": True,
        "optionsFlowYellow": True,
        "optionsFlowWhite": True,
        "optionsFlowMagenta": True,
        "optionsFlowAboveAskOnly": True,
        "optionsFlowBelowBidOnly": True,
        "optionsFlowAtOrAboveAsk": True,
        "optionsFlowAtOrBelowBid": True,
        "optionsFlowMultileg": False,
        "optionsFlowOnlyMultiLeg": False,
        "optionsFlowBelowPoint5": False,
        "optionsFlowBelow5": False,
        "optionsFlow100Contracts": False,
        "optionsFlow500Contracts": False,
        "optionsFlow5000Contracts": False,
        "optionsFlowStock": True,
        "optionsFlowEtf": True,
        "optionsFlowAbove50k": False,
        "optionsFlowAbove100k": False,
        "optionsFlowAbove200k": False,
        "optionsFlowAbove500k": False,
        "optionsFlowAbove1m": False,
        "marketCapAbove750B": False,
        "optionsFlowInTheMoney": False,
        "optionsFlowOutOfTheMoney": False,
        "optionsFlowSweepOnly": False,
        "optionsFlowWeeklyOnly": False,
        "optionsFlowEarningsReportOnly": False,
        "optionsFlowUnusualOnly": False,
        "optionsFlowExDiv": False,
        "optionsFlowConsumerDiscretionary": True,
        "optionsFlowIndustrials": True,
        "optionsFlowInformationTechnology": True,
        "optionsFlowRealEstate": True,
        "optionsFlowHealthCare": True,
        "optionsFlowEnergy": True,
        "optionsFlowFinancials": True,
        "optionsFlowMaterials": True,
        "optionsFlowConsumerStaples": True,
        "optionsFlowCommunicationServices": True,
        "optionsFlowUtilities": True,
        "optionsExpirationRange": False,
        "optionsFlowSectorNone": True,
    }
    
    if custom_filters:
        filters.update(custom_filters)
    
    return filters

# Constants
TIMEZONE_SUFFIX = "+00:00"

def fetch_blackbox_flow(options: Optional[Dict] = None) -> List[Dict]:
    """Fetch options flow data from BlackBox Stocks API."""
    if options is None:
        options = {}
    
    if not BLACKBOX_API_TOKEN:
        raise ValueError(
            "BLACKBOX_API_TOKEN not found in environment variables.\n"
            "Please add BLACKBOX_API_TOKEN to your .env file or set it as an environment variable."
        )
    
    # Parse dates - default to today if not provided
    if options.get("startDate"):
        if isinstance(options["startDate"], str):
            start_date = datetime.fromisoformat(options["startDate"].replace("Z", TIMEZONE_SUFFIX))
        elif isinstance(options["startDate"], date):
            start_date = datetime.combine(options["startDate"], datetime.min.time())
        else:
            start_date = options["startDate"]
    else:
        start_date = datetime.now()
    
    if options.get("endDate"):
        if isinstance(options["endDate"], str):
            end_date = datetime.fromisoformat(options["endDate"].replace("Z", TIMEZONE_SUFFIX))
        elif isinstance(options["endDate"], date):
            end_date = datetime.combine(options["endDate"], datetime.max.time())
        else:
            end_date = options["endDate"]
    else:
        end_date = start_date
    
    # Build request body
    body = {
        "historical": options.get("historical", False),
        "symbol": options.get("symbol", ""),
        "strike": options.get("strike", 0),
        "count": options.get("count") or options.get("limit", 300),
        "filter": build_filter_bitmask(options.get("filters")),
        "filters": build_filters(start_date, end_date, options.get("filters")),
        "fromDate": start_date.isoformat(),
        "toDate": end_date.isoformat()
    }
    
    try:
        response = requests.post(
            BLACKBOX_API_URL,
            headers={
                "Content-Type": "application/json",
                "Accept": "application/json",
                "Authorization": f"Bearer {BLACKBOX_API_TOKEN}"
            },
            json=body,
            timeout=30
        )
        
        if not response.ok:
            error_text = response.text
            raise RuntimeError(
                f"BlackBox API error: {response.status_code} {response.reason}\n"
                f"Response: {error_text}"
            )
        
        data = response.json()
        
        # Handle different response structures
        if isinstance(data, list):
            return data
        elif isinstance(data, dict):
            if "data" in data and isinstance(data["data"], list):
                return data["data"]
            elif "flows" in data and isinstance(data["flows"], list):
                return data["flows"]
            elif "results" in data and isinstance(data["results"], list):
                return data["results"]
            else:
                print_flush(f"{_ts()} | ⚠️ Unexpected API response structure: {list(data.keys())}")
                return []
        else:
            return []
    except Exception as e:
        print_flush(f"{_ts()} | ❌ Error fetching BlackBox flow data: {e}")
        raise

def map_blackbox_to_database(api_record: Dict) -> Dict:
    """Map BlackBox API response to database schema."""
    def get_value(obj: Dict, *keys: str) -> Optional[str]:
        """Safely extract values from object."""
        for key in keys:
            if key in obj and obj[key] is not None:
                return str(obj[key])
        return None
    
    def format_date(date_str: Optional[str]) -> Optional[str]:
        """Format date as YYYY-MM-DD."""
        if not date_str:
            return None
        try:
            dt = datetime.fromisoformat(str(date_str).replace("Z", TIMEZONE_SUFFIX))
            return dt.strftime("%Y-%m-%d")
        except (ValueError, AttributeError):
            try:
                dt = datetime.strptime(str(date_str), "%Y-%m-%d")
                return dt.strftime("%Y-%m-%d")
            except (ValueError, AttributeError):
                return str(date_str)
    
    def format_time(time_str: Optional[str]) -> Optional[str]:
        """Format time string."""
        if not time_str:
            return None
        return str(time_str)
    
    # Map fields - try multiple possible field names from API
    mapped = {
        "CreatedDate": format_date(
            get_value(api_record, "createdDate", "CreatedDate", "date", "Date", "timestamp", "Timestamp")
        ),
        "CreatedTime": format_time(
            get_value(api_record, "createdTime", "CreatedTime", "time", "Time", "timestamp", "Timestamp")
        ),
        "Symbol": get_value(api_record, "symbol", "Symbol", "ticker", "Ticker", "underlying", "Underlying"),
        "Type": get_value(api_record, "type", "Type", "tradeType", "TradeType"),
        "Volume": get_value(api_record, "volume", "Volume", "vol", "Vol", "contracts", "Contracts"),
        "Price": get_value(api_record, "price", "Price", "lastPrice", "LastPrice", "tradePrice", "TradePrice"),
        "Side": get_value(api_record, "side", "Side", "tradeSide", "TradeSide", "direction", "Direction"),
        "CallPut": get_value(api_record, "callPut", "CallPut", "optionType", "OptionType", "putCall", "PutCall", "type", "Type"),
        "Strike": get_value(api_record, "strike", "Strike", "strikePrice", "StrikePrice"),
        "Spot": get_value(api_record, "spot", "Spot", "underlyingPrice", "UnderlyingPrice", "stockPrice", "StockPrice"),
        "Premium": get_value(api_record, "premium", "Premium", "totalPremium", "TotalPremium", "notional", "Notional"),
        "ExpirationDate": format_date(
            get_value(api_record, "expirationDate", "ExpirationDate", "expiry", "Expiry", "expiration", "Expiration")
        ),
        "Color": get_value(api_record, "color", "Color", "tradeColor", "TradeColor"),
        "ImpliedVolatility": get_value(api_record, "impliedVolatility", "ImpliedVolatility", "iv", "IV", "volatility", "Volatility"),
        "Dte": get_value(api_record, "dte", "Dte", "DTE", "daysToExpiration", "DaysToExpiration", "daysToExpiry", "DaysToExpiry"),
        "ER": get_value(api_record, "er", "ER", "earnings", "Earnings", "earningsReport", "EarningsReport"),
        "StockEtf": get_value(api_record, "stockEtf", "StockEtf", "assetType", "AssetType", "securityType", "SecurityType"),
        "Sector": get_value(api_record, "sector", "Sector", "industry", "Industry"),
        "Uoa": get_value(api_record, "uoa", "Uoa", "UOA", "underlyingOfAsset", "UnderlyingOfAsset"),
        "Weekly": get_value(api_record, "weekly", "Weekly", "isWeekly", "IsWeekly", "weeklies", "Weeklies"),
        "MktCap": get_value(api_record, "mktCap", "MktCap", "marketCap", "MarketCap", "marketCapitalization", "MarketCapitalization"),
        "OI": get_value(api_record, "oi", "OI", "openInterest", "OpenInterest", "openInt", "OpenInt")
    }
    
    return mapped

# ─────────────────────────────────────────────
# Normalization + JSON safety
# ─────────────────────────────────────────────
WEIRD_STR = {"inf","+inf","-inf","infinity","+infinity","-infinity","∞","+∞","-∞",
             "nan","-nan","NaN","N/A","NA","NULL","null",""}

def _coerce_weird_numbers(df: pd.DataFrame) -> pd.DataFrame:
    df = df.copy()
    for c in df.columns:
        if df[c].dtype == object:
            df[c] = df[c].replace(list(WEIRD_STR), np.nan)
    return df

def _normalize_for_table(df: pd.DataFrame, table: str) -> pd.DataFrame:
    """Rename/select columns to match the DB style of this table."""
    style = TABLE_STYLE.get(table, "snake").lower()
    df = df.copy()
    df.columns = [c.strip() for c in df.columns]

    if table in ("OptionsFlow","OptionsFlow_monthly"):
        if style == "camel":
            # Ensure CamelCase headers, no renaming needed if file already CamelCase
            # If file is snake, map to Camel
            lower_cols = {c for c in df.columns if c.islower()}
            if lower_cols:
                df = df.rename(columns=MAP_OPTIONSFLOW_CAMEL)
            for c in EXPECTED_OPTIONSFLOW_CAMEL:
                if c not in df.columns: df[c] = None
            df = df[EXPECTED_OPTIONSFLOW_CAMEL]
        else:
            # snake_case target
            # If file CamelCase, map to snake
            has_upper = any(any(ch.isupper() for ch in c) for c in df.columns)
            if has_upper:
                df = df.rename(columns=MAP_OPTIONSFLOW_SNAKE)
            for c in EXPECTED_OPTIONSFLOW_SNAKE:
                if c not in df.columns: df[c] = None
            df = df[EXPECTED_OPTIONSFLOW_SNAKE]

    elif table in ("AlertStream","AlertStream_monthly"):
        # DB is snake_case per your screenshot
        # If file CamelCase, map to snake
        has_upper = any(any(ch.isupper() for ch in c) for c in df.columns)
        if has_upper:
            df = df.rename(columns=MAP_ALERTSTREAM_SNAKE)
        for c in EXPECTED_ALERTSTREAM_SNAKE:
            if c not in df.columns: df[c] = None
        df = df[EXPECTED_ALERTSTREAM_SNAKE]

    # Standardize any *date columns → YYYY-MM-DD* strings
    for col in [c for c in df.columns if c.lower().endswith("date")]:
        df[col] = pd.to_datetime(df[col], errors="coerce").dt.strftime("%Y-%m-%d")

    return df

def _json_safe(df: pd.DataFrame) -> pd.DataFrame:
    df = df.copy()
    # numeric: drop non-finite
    for c in df.columns:
        if pd.api.types.is_numeric_dtype(df[c]):
            s = pd.to_numeric(df[c], errors="coerce")
            s[~np.isfinite(s)] = np.nan
            df[c] = s
    # datetimes -> strings
    for c in df.columns:
        if pd.api.types.is_datetime64_any_dtype(df[c]):
            df[c] = df[c].dt.strftime("%Y-%m-%d %H:%M:%S")
    # NA -> None
    return df.astype(object).where(pd.notnull(df), None)

def _row_hash_from_series(s: pd.Series) -> str:
    vals=[]
    for _,v in s.items():
        if v is None or (isinstance(v,float) and pd.isna(v)): vals.append("NULL")
        elif isinstance(v,str): vals.append(v.strip())
        else: vals.append(str(v))
    return hashlib.sha1("\x1f".join(vals).encode("utf-8","ignore")).hexdigest()

def _df_prepare_for_postgres(df: pd.DataFrame) -> pd.DataFrame:
    if df.empty: return df
    df = _json_safe(df)
    df["row_hash"] = df.apply(_row_hash_from_series, axis=1)
    return df.drop_duplicates(subset=["row_hash"], keep="first").reset_index(drop=True)

def _extract_missing_col_from_error(msg: str) -> Optional[str]:
    """Extract missing column name from PostgreSQL error messages."""
    patterns = [
        r"column \"([^\"]+)\" does not exist",
        r"Could not find the '([^']+)' column",
    ]
    for pattern in patterns:
        m = re.search(pattern, msg, re.IGNORECASE)
        if m:
            return m.group(1)
    return None

# ─────────────────────────────────────────────
# PostgreSQL upload (quiet, chunked)
# ─────────────────────────────────────────────
def _ensure_table_exists(conn, table_name: str, columns: list):
    """Ensure table exists with row_hash column and unique constraint."""
    cur = conn.cursor()
    try:
        # Check if table exists
        cur.execute("""
            SELECT EXISTS (
                SELECT FROM information_schema.tables 
                WHERE table_schema = current_schema()
                AND table_name = %s
            )
        """, (table_name,))
        exists = cur.fetchone()[0]
        
        if not exists:
            # Create table with all columns as text initially (we'll let PostgreSQL infer types)
            # For now, we'll create a basic structure - actual schema should match your data
            col_defs = ", ".join([f'"{col}" TEXT' for col in columns if col != "row_hash"])
            col_defs += ', "row_hash" TEXT UNIQUE'
            
            cur.execute(f'CREATE TABLE IF NOT EXISTS "{table_name}" ({col_defs})')
            conn.commit()
        else:
            # Ensure row_hash column and unique constraint exist
            cur.execute("""
                SELECT column_name FROM information_schema.columns 
                WHERE table_schema = current_schema() 
                AND table_name = %s AND column_name = 'row_hash'
            """, (table_name,))
            if not cur.fetchone():
                cur.execute(f'ALTER TABLE "{table_name}" ADD COLUMN IF NOT EXISTS "row_hash" TEXT')
                conn.commit()
            
            # Check for unique constraint on row_hash
            cur.execute("""
                SELECT constraint_name FROM information_schema.table_constraints 
                WHERE table_schema = current_schema() 
                AND table_name = %s 
                AND constraint_type = 'UNIQUE'
                AND constraint_name LIKE %s
            """, (table_name, f'%{table_name}%row_hash%'))
            if not cur.fetchone():
                try:
                    cur.execute(f'CREATE UNIQUE INDEX IF NOT EXISTS "{table_name}_row_hash_idx" ON "{table_name}" ("row_hash")')
                    conn.commit()
                except Exception:
                    conn.rollback()
    finally:
        cur.close()

def _upsert_slice(conn, tname: str, rows: list, columns: list):
    """Upsert a slice of rows using PostgreSQL INSERT ... ON CONFLICT."""
    if not rows:
        return
    
    cur = conn.cursor()
    try:
        # Ensure table exists
        _ensure_table_exists(conn, tname, columns)
        
        # Build INSERT ... ON CONFLICT statement
        cols_quoted = ", ".join([f'"{col}"' for col in columns])
        updates = ", ".join([f'"{col}" = EXCLUDED."{col}"' for col in columns if col != "row_hash"])
        
        # Build the base query template for execute_values
        base_query = f'INSERT INTO "{tname}" ({cols_quoted}) VALUES %s'
        if updates:
            conflict_query = f'{base_query} ON CONFLICT ("row_hash") DO UPDATE SET {updates}'
        else:
            conflict_query = f'{base_query} ON CONFLICT ("row_hash") DO NOTHING'
        
        # Prepare values as list of tuples
        values = [tuple(row.get(col) for col in columns) for row in rows]
        
        execute_values(cur, conflict_query, values, template=None, page_size=len(rows))
        conn.commit()
    except Exception:
        conn.rollback()
        raise
    finally:
        cur.close()

def _postgres_upsert(table: str, df: pd.DataFrame):
    if df.empty:
        return
    
    tname = table
    total = len(df)
    columns = [col for col in df.columns if col != "row_hash"] + ["row_hash"]
    
    ranges = [(0, min(UPSERT_CHUNK,total))] if SINGLE_CHUNK else \
             [(i, min(i+UPSERT_CHUNK,total)) for i in range(0,total,UPSERT_CHUNK)]

    if not PRINT_PROGRESS:
        lbl = "(single-chunk)" if SINGLE_CHUNK else f"chunk_size={UPSERT_CHUNK}"
        print_flush(f"{_ts()} | ☁️ Starting upsert → [{tname}] rows={total:,} {lbl}")

    sent = 0
    conn = None
    
    try:
        for start, end in ranges:
            part = df.iloc[start:end].copy()
            if "row_hash" in part.columns:
                part = part.drop_duplicates(subset=["row_hash"], keep="first")
            
            # Ensure all columns are present
            for col in columns:
                if col not in part.columns:
                    part[col] = None
            
            rows = part[columns].to_dict(orient="records")
            
            tried_prune = False
            for attempt in range(1, MAX_RETRIES + 1):
                try:
                    if conn is None or conn.closed:
                        conn = _pg_conn()
                    _upsert_slice(conn, tname, rows, columns)
                    break
                except Exception as e:
                    missing = _extract_missing_col_from_error(str(e))
                    if (missing is not None) and (missing in part.columns) and (not tried_prune):
                        if PRINT_PRUNE_LOGS:
                            print_flush(f"{_ts()} | ⚠️ [{tname}] pruning missing column '{missing}'")
                        part = part.drop(columns=[missing])
                        columns = [c for c in columns if c != missing]
                        rows = part[columns].to_dict(orient="records")
                        tried_prune = True
                        continue
                    if attempt == MAX_RETRIES:
                        raise
                    time.sleep(1.0 * attempt)
                    if conn and not conn.closed:
                        conn.close()
                    conn = None

            sent += len(part)
            if PRINT_PROGRESS:
                print_flush(f"{_ts()} | ☁️ [{tname}] {sent:,}/{total:,}")

        if not PRINT_PROGRESS:
            print_flush(f"{_ts()} | ☁️ [{tname}] done {sent:,}/{total:,}")
    finally:
        if conn and not conn.closed:
            conn.close()

def _postgres_replace(table: str, df: pd.DataFrame):
    """Replace all data in table (delete then insert)."""
    conn = _pg_conn()
    try:
        cur = conn.cursor()
        cur.execute(f'DELETE FROM "{table}"')
        conn.commit()
        cur.close()
    except Exception:
        conn.rollback()
        # Table might not exist, that's okay
    finally:
        conn.close()
    
    _postgres_upsert(table, df)

def _load_to_postgres(df: pd.DataFrame, table_name: str, _source_path: str = ""):
    ndf = _normalize_for_table(df, table_name)
    if ndf.empty:
        print_flush(f"{_ts()} | ☁️ Empty after normalize. Skipping PostgreSQL [{table_name}]")
        return
    
    ndf = _df_prepare_for_postgres(ndf)
    
    if INCREMENTAL:
        _postgres_upsert(table_name, ndf)
        print_flush(f"{_ts()} | ☁️✅ Upserted {len(ndf):,} rows → PostgreSQL [{table_name}]")
    else:
        _postgres_replace(table_name, ndf)
        print_flush(f"{_ts()} | ☁️✅ Replaced table with {len(ndf):,} rows → PostgreSQL [{table_name}]")

# ─────────────────────────────────────────────
# Orchestrator
# ─────────────────────────────────────────────
def _load_records_to_databases(records: List[Dict], table_name: str):
    """Load records from API into both SQLite and PostgreSQL."""
    if not records:
        print_flush(f"{_ts()} | ⚠️ No records to process")
        return
    
    # Convert records to DataFrame
    df = pd.DataFrame(records)
    df = _coerce_weird_numbers(df)
    
    if WRITE_POSTGRES:
        try:
            _load_to_postgres(df, table_name)
        except Exception as e:
            print_flush(f"{_ts()} | ❌ (PostgreSQL) Error: {e}")

# ─────────────────────────────────────────────
# Main
# ─────────────────────────────────────────────
def sync_blackbox_flow():
    """Main sync function."""
    print_flush(f"{_ts()} | 🚀 Starting BlackBox Stocks flow data sync...\n")
    
    try:
        # Parse command line arguments
        parser = argparse.ArgumentParser(description="Sync BlackBox Stocks options flow data to databases")
        parser.add_argument("--start-date", type=str, help="Start date (YYYY-MM-DD)")
        parser.add_argument("--end-date", type=str, help="End date (YYYY-MM-DD)")
        parser.add_argument("--limit", type=int, help="Maximum number of records to fetch")
        parser.add_argument("--count", type=int, help="Maximum number of records to fetch (alias for --limit)")
        parser.add_argument("--symbol", type=str, help="Filter by specific symbol")
        parser.add_argument("--table", type=str, default="OptionsFlow_monthly", help="Target table name (default: OptionsFlow_monthly)")
        
        args = parser.parse_args()
        
        options = {}
        if args.start_date:
            options["startDate"] = args.start_date
        if args.end_date:
            options["endDate"] = args.end_date
        if args.limit:
            options["count"] = args.limit
        elif args.count:
            options["count"] = args.count
        if args.symbol:
            options["symbol"] = args.symbol
        
        table_name = args.table
        
        # Default to today if no dates provided
        if not options.get("startDate") and not options.get("endDate"):
            today = date.today().isoformat()
            options["startDate"] = today
            options["endDate"] = today
            print_flush(f"{_ts()} | 📅 No date range specified, using today: {today}")
        
        print_flush(f"{_ts()} | 📥 Fetching flow data from BlackBox API...")
        print_flush(f"{_ts()} |    Options: {options}")
        
        # Fetch data from API
        api_records = fetch_blackbox_flow(options)
        print_flush(f"{_ts()} |    ✅ Fetched {len(api_records)} records from API")
        
        if len(api_records) == 0:
            print_flush(f"{_ts()} | ⚠️  No records returned from API")
            return
        
        # Log sample record
        if len(api_records) > 0:
            print_flush(f"\n{_ts()} | 📋 Sample API record structure:")
            print_flush(json.dumps(api_records[0], indent=2, default=str))
        
        # Map API records to database schema
        print_flush(f"\n{_ts()} | 🔄 Mapping records to database schema...")
        mapped_records = [map_blackbox_to_database(record) for record in api_records]
        print_flush(f"{_ts()} |    ✅ Mapped {len(mapped_records)} records")
        
        # Log sample mapped record
        if len(mapped_records) > 0:
            print_flush(f"\n{_ts()} | 📋 Sample mapped record:")
            print_flush(json.dumps(mapped_records[0], indent=2, default=str))
        
        # Insert into databases
        print_flush(f"\n{_ts()} | 💾 Inserting records into databases...")
        _load_records_to_databases(mapped_records, table_name)
        
        print_flush(f"\n{_ts()} | ✅ Successfully synced {len(mapped_records)} records")
        
        # Summary
        print_flush(f"\n{_ts()} | 📊 Sync Summary:")
        print_flush(f"{_ts()} |    Fetched from API: {len(api_records)} records")
        print_flush(f"{_ts()} |    Inserted into DB: {len(mapped_records)} records")
        
        # Get total count in PostgreSQL database
        if WRITE_POSTGRES:
            try:
                conn = _pg_conn()
                cur = conn.cursor()
                cur.execute(f'SELECT COUNT(*) FROM "{table_name}"')
                total_count = cur.fetchone()[0]
                print_flush(f"{_ts()} |    Total records in PostgreSQL DB: {total_count}")
                conn.close()
            except Exception as e:
                print_flush(f"{_ts()} |    ⚠️ Could not get PostgreSQL count: {e}")
        
    except Exception as e:
        print_flush(f"\n{_ts()} | ❌ Sync failed: {e}")
        import traceback
        traceback.print_exc()
        raise

if __name__ == "__main__":
    if WRITE_POSTGRES:
        print_flush(f"{_ts()} | ☁️ PostgreSQL: {POSTGRES_HOST}:{POSTGRES_PORT}/{POSTGRES_DB} "
                    f"(schema={POSTGRES_SCHEMA}) | INCREMENTAL={INCREMENTAL} | chunk={UPSERT_CHUNK}")
    sync_blackbox_flow()
