import argparse
import sys, time, re, requests, random, os
from typing import List, Tuple, Set, Optional, Iterable
import pandas as pd
import yfinance as yf
from bs4 import BeautifulSoup
from datetime import datetime, timedelta, timezone, date
from io import StringIO
import psycopg2
from psycopg2.extras import execute_values
from dotenv import load_dotenv
try:
    from zoneinfo import ZoneInfo
    CT = ZoneInfo("America/Chicago")
except Exception:
    CT = None

load_dotenv()



# ========= USER CONFIG =======================================================
# PostgreSQL connection - uses environment variables or defaults to local
# Set DATABASE_URL or USE_LOCAL_DB=true with LOCAL_DB_* variables
SEGMENTS               = ["large", "mid", "nasdaq"]      # auto universes to include
TAG_AUTO_SEGMENTS      = True                            # tag auto symbols as AUTO:SP500/SP400/NDX100
EXTRA_SYMBOLS          = []                              # optional flat list, e.g. ["TSLA","COIN","SPY"]

EXTRA_GROUPS = {
    "OIL": ["HUSA","INDO","USEG","MARPS","BRN","PED","GBR","KLXE","CEIN","IMPP","CVX","OXY","GUSH","RCON","APA","VIVK"],
    "FASTFOOD_REST": ["SG","CMG","WEN","QSR","WING","BROS","SHAK","CAVA","SBUX","MCD","YUM"],
    "ROBOTICS": ["SERV","RR","KITT","MBOT","IRBT","ARBE","SYM","TSLA","NVDA","ROBO","BOTZ","TER","ROK"],
    "SEMIS": ["NVDA","MCHP","NVTS","ADI","QCOM","AVGO","ON","TXN","QQQ","OLED","STM","ASML","RMBS","ARM","SMH","INTC","LRCX","DELL","MRVL","TSM","SMCI","MU","SOXL","AMD"],
    "QUANTUM": ["IONQ","RGTI","QUBT","QBTS","ARQQ","QMCO","COHR","MRVL","HON","TSEM","FORM"],
    "CYBER": ["CYBR","FTNT","PANW","ZS","CRWD","OKTA","S"],
    "FINTECH": ["SOFI","AFRM","NU","V","MA","PYPL","UPST","FOUR","AXP","PSFE","PAYO","ADYEY","QFIN"],
    "CATHY": ["TXG","TWST","CRSP","COIN","PACB","DNA","SOFI","PD","HOOD","ROKU","ZM","EXAS","U","NTLA","SPY","TER","PATH","PINS","META","ARKK","CERS","BEAM","VCYT","RBLX","TTD","TSLA","TDOC","ACHR"],
    "CHINA": ["FUTU","KC","LI","EDU","BILI","KWEB","PDD","JD","BABA","BIDU","NIO","XPEV","GDS"],
    "SOLAR": ["DQ","SEDG","JKS","ENPH","TAN","CSIQ","ARRY","FSLR","RUN"],
    "BANKS": ["MS","C","GS","KRE","WFC","JPM","BAC","SCHW","XLF"],
    "TRAVEL": ["EXPE","BKNG","WYNN","LVS","TRIP","ABNB","JETS","RCL","LUV","DAL","UAL","NCLH","CCL","MAR","HLT"],
    "WEED": ["ACB","MSOS","CGC","CURLF","GTBIF","TLRY","SNDL","MSOX","MJ"],
    "DOW": ["WMT","CRM","PG","HD","MSFT","V","MCD","IBM","AMGN","VZ","CSCO","AAPL","MRK","DIS","TRV","HON","JNJ","JPM","NKE","AXP","MMM","GS","CVX","CAT","UNH","BA","BE"],
    "NUCLEAR": ["SMR","OKLO","NNE","ASPI","LTBR","LEU","TLN","CEG","PEG"],
    "URANIUM": ["URA","NLR","URNM","NXE","URAN","URNJ","CCJ","UEC"],
    "SAAS-CLOUD": ["NOW","CRM","WDAY","OKTA","MDB","DDOG","TWLO","DOCU","WIX","TEAM","SNOW","NET","U","PLTR","BILL","DOCN","VEEV","CRWV"],
    "SPACE-ROCKET": ["RKLB","SPCE","SIDU","MNTS","RDW","LUNR","ASTS","SATS","ARKX","VSAT","IRDM","SES"],
    "MEME": ["NEGG","AMC","BBAI","BYND","OPEN","UPST","NVAX","GME","DJT","BB","TRUP","KOSS","KODK","SPWR","RDDT","HOOD"],
    "DRONES": ["NOC","LMT","KTOS","BA","AIRO","UMAC","DPRO","ONDS","RTX","RCAT","EVTL","AMBA","ACHR","JOBY","GPRO","HON","GOGO","EH"]
}
SYMBOLS_ONLY           = False            # True = ignore SEGMENTS and use only EXTRA_SYMBOLS + EXTRA_GROUPS
LOOKBACK_DAYS_INTRADAY = 30               # 30 days (1 month) - uses 5m interval for >7 days
DAILY_PERIOD           = "2y"
BATCH_SIZE             = 110               # 50–120 good
SLEEP_BETWEEN_BATCHES  = 0.25             # seconds
LOOP                   = False  # Set to False for one-shot update, True for continuous loop
INTRADAY_INTERVAL_MIN  = 1
DAILY_INTERVAL_MIN     = 1
# ============================================================================

USER_AGENT = "Mozilla/5.0 (Windows NT 10.0; Win64; x64)"
HEADERS    = {"User-Agent": USER_AGENT}
URL_SP500  = "https://en.wikipedia.org/wiki/List_of_S%26P_500_companies"
URL_SP400  = "https://en.wikipedia.org/wiki/List_of_S%26P_400_companies"
URL_NASDAQ = "https://en.wikipedia.org/wiki/Nasdaq-100"
VALID_TICKER_RE = re.compile(r"^[A-Z][A-Z0-9.\-]{0,6}$")
BAD_WORDS = {"CLOSING","INTRADAY","TICKER","SYMBOL","INDEX","COMPANY","WEIGHT","WEIGHTS"}


# ---------- time helpers ----------
def now_utc() -> str:
    return datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M:%S UTC")
def now_ct() -> str:
    if CT:
        return datetime.now(CT).strftime("%Y-%m-%d %I:%M:%S %p CT")
    return time.strftime("%Y-%m-%d %I:%M:%S %p (local)")
def stamp(label: str) -> str:
    return f"[{label}] {now_ct()}  |  {now_utc()}"


# ---------- db ----------
def connect_rw() -> psycopg2.extensions.connection:
    """Create a PostgreSQL connection using environment variables."""
    use_local_db = os.getenv('USE_LOCAL_DB', 'false').lower() == 'true'
    
    if use_local_db:
        config = {
            'host': os.getenv('LOCAL_DB_HOST', '192.168.8.151'),
            'port': int(os.getenv('LOCAL_DB_PORT', '5432')),
            'user': os.getenv('LOCAL_DB_USER', 'postgres'),
            'password': os.getenv('LOCAL_DB_PASSWORD', 'postgres'),
            'database': os.getenv('LOCAL_DB_NAME', 'institutional_trader')
        }
    elif os.getenv('DATABASE_URL'):
        # Parse connection string
        import urllib.parse
        url = urllib.parse.urlparse(os.getenv('DATABASE_URL'))
        config = {
            'host': url.hostname,
            'port': url.port or 5432,
            'user': url.username,
            'password': url.password,
            'database': url.path.lstrip('/')
        }
    else:
        raise ValueError(
            'No database configuration found. Set either:\n'
            '  - USE_LOCAL_DB=true with LOCAL_DB_* variables\n'
            '  - OR DATABASE_URL'
        )
    
    return psycopg2.connect(**config)


def ensure_metadata_tables(con: psycopg2.extensions.connection):
    with con.cursor() as cur:
        cur.execute("""
        CREATE TABLE IF NOT EXISTS universe_symbols (
          symbol    VARCHAR(20) PRIMARY KEY,
          source    VARCHAR(50),
          added_at  TIMESTAMP WITH TIME ZONE
        )""")
        cur.execute("""
        CREATE TABLE IF NOT EXISTS symbol_sectors (
          symbol    VARCHAR(20) NOT NULL,
          sector    VARCHAR(100) NOT NULL,
          PRIMARY KEY (symbol, sector)
        )""")
        con.commit()

def ensure_prices_daily_schema(con: psycopg2.extensions.connection):
    with con.cursor() as cur:
        # Note: PostgreSQL schema uses "Date" (capital D) to match existing schema
        cur.execute("""
        CREATE TABLE IF NOT EXISTS prices_daily (
          symbol     VARCHAR(10) NOT NULL,
          "Date"     DATE NOT NULL,
          open       NUMERIC(12, 4),
          high       NUMERIC(12, 4),
          low        NUMERIC(12, 4),
          close      NUMERIC(12, 4),
          volume     NUMERIC(15, 2),
          PRIMARY KEY (symbol, "Date")
        )""")
        cur.execute("""
        CREATE INDEX IF NOT EXISTS idx_prices_daily_symbol_date
          ON prices_daily(symbol, "Date" DESC)
        """)
        con.commit()
def ensure_intraday_schema(con: psycopg2.extensions.connection):
    with con.cursor() as cur:
        # PostgreSQL schema uses TIMESTAMP WITH TIME ZONE instead of INTEGER
        cur.execute("""
        CREATE TABLE IF NOT EXISTS prices_intraday_1m (
          symbol   VARCHAR(10) NOT NULL,
          ts       TIMESTAMP WITH TIME ZONE NOT NULL,
          open     NUMERIC(12, 4),
          high     NUMERIC(12, 4),
          low      NUMERIC(12, 4),
          close    NUMERIC(12, 4),
          volume   NUMERIC(15, 2),
          PRIMARY KEY (symbol, ts)
        )""")
        cur.execute("""
        CREATE INDEX IF NOT EXISTS idx_prices_intraday_symbol_ts
          ON prices_intraday_1m(symbol, ts DESC)
        """)
        cur.execute("""
        CREATE INDEX IF NOT EXISTS idx_prices_intraday_ts
          ON prices_intraday_1m(ts DESC)
        """)
        con.commit()


def get_last_daily_date(con: psycopg2.extensions.connection, sym: str) -> Optional[str]:
    """Get the last date for a symbol from prices_daily table."""
    with con.cursor() as cur:
        cur.execute('SELECT MAX("Date") FROM prices_daily WHERE symbol=%s', (sym,))
        row = cur.fetchone()
        if row and row[0]:
            # Convert date to string format YYYY-MM-DD
            if isinstance(row[0], date):
                return row[0].isoformat()
            return str(row[0])
        return None


# ---------- symbol universe ----------
def normalize_yahoo_symbol(s: str) -> str:
    return s.strip().upper().replace(".", "-")

def _unique(seq: Iterable[str]) -> List[str]:
    seen, out = set(), []
    for x in seq:
        if x not in seen:
            out.append(x); seen.add(x)
    return out

def _parse_first_wikitable_with_columns(html: str, preferred_cols: Tuple[str, ...]) -> List[str]:
    soup = BeautifulSoup(html, "html.parser")
    for t in soup.select("table.wikitable"):
        try:
            df = pd.read_html(StringIO(str(t)))[0]
        except Exception:
            continue
        for pc in preferred_cols:
            if pc in df.columns:
                ser = (df[pc].astype(str)
                         .str.replace(r"\[.*?\]", "", regex=True)
                         .str.strip())
                out = []
                for raw in ser.tolist():
                    up = raw.upper().replace(" ", "").replace("\u200b","")
                    if up in BAD_WORDS:  # safety against header junk
                        continue
                    if VALID_TICKER_RE.match(up):
                        out.append(normalize_yahoo_symbol(up))
                if out:
                    return out
    return []

def build_ticker_universe(segments: List[str]) -> Tuple[List[str], dict]:
    # manual group tags
    sector_map = {}
    manual_group_syms = []
    for sector, lst in (EXTRA_GROUPS or {}).items():
        for s in lst:
            sym = normalize_yahoo_symbol(s)
            manual_group_syms.append(sym)
            sector_map.setdefault(sym, set()).add(sector)
    manual_group_syms = _unique(manual_group_syms)

    manual_flat = [normalize_yahoo_symbol(s) for s in EXTRA_SYMBOLS]
    for s in manual_flat:
        sector_map.setdefault(s, set()).add("MANUAL_LIST")

    if SYMBOLS_ONLY:
        return _unique(manual_group_syms + manual_flat), sector_map

    auto_syms: List[str] = []
    if "large" in segments:
        print("Fetching S&P 500 tickers…")
        r = requests.get(URL_SP500, headers=HEADERS, timeout=20); r.raise_for_status()
        spx = _parse_first_wikitable_with_columns(r.text, ("Symbol","Ticker","Ticker symbol"))
        auto_syms += spx
        if TAG_AUTO_SEGMENTS:
            for s in spx: sector_map.setdefault(s, set()).add("AUTO:SP500")

    if "mid" in segments:
        print("Fetching S&P 400 tickers…")
        r = requests.get(URL_SP400, headers=HEADERS, timeout=20); r.raise_for_status()
        mid = _parse_first_wikitable_with_columns(r.text, ("Symbol","Ticker","Ticker symbol"))
        auto_syms += mid
        if TAG_AUTO_SEGMENTS:
            for s in mid: sector_map.setdefault(s, set()).add("AUTO:SP400")

    if "nasdaq" in segments:
        print("Fetching Nasdaq-100 tickers…")
        r = requests.get(URL_NASDAQ, headers=HEADERS, timeout=20); r.raise_for_status()
        ndx = _parse_first_wikitable_with_columns(r.text, ("Ticker","Ticker symbol","Symbol"))
        auto_syms += ndx
        if TAG_AUTO_SEGMENTS:
            for s in ndx: sector_map.setdefault(s, set()).add("AUTO:NDX100")

    return _unique(auto_syms + manual_group_syms + manual_flat), sector_map


# ---------- yfinance helpers ----------
def yf_download_multi(tickers: List[str], period: str, interval: str) -> pd.DataFrame:
    return yf.download(tickers=tickers, period=period, interval=interval,
                       auto_adjust=False, progress=False, threads=True, group_by="ticker")

def chunks(seq, n):
    for i in range(0, len(seq), n):
        yield seq[i:i+n]


# ---------- metadata upserts ----------
def ensure_meta(con: psycopg2.extensions.connection):
    ensure_metadata_tables(con)

def upsert_universe(con: psycopg2.extensions.connection, symbols: List[str], sector_map: dict):
    now_iso = datetime.now(timezone.utc)
    rows = []
    for s in symbols:
        tags = sector_map.get(s, set())
        is_manual = ("MANUAL_LIST" in tags) or any(not t.startswith("AUTO:") for t in tags)
        rows.append((s, "manual" if is_manual else "auto", now_iso))
    
    with con.cursor() as cur:
        execute_values(cur, """
          INSERT INTO universe_symbols(symbol, source, added_at)
          VALUES %s
          ON CONFLICT (symbol) DO NOTHING
        """, rows)
        
        pairs = []
        for s, tags in sector_map.items():
            for t in tags:
                pairs.append((s, t))
        if pairs:
            execute_values(cur, """
              INSERT INTO symbol_sectors(symbol, sector)
              VALUES %s
              ON CONFLICT (symbol, sector) DO NOTHING
            """, pairs)
        con.commit()


# ---------- data upserts ----------
def upsert_daily(con: psycopg2.extensions.connection, sym: str, df: pd.DataFrame, source="yfinance", diag: bool=False) -> int:
    if df is None or df.empty:
        print(f"[daily upsert] {sym}: EMPTY dataframe from yfinance")
        return 0

    # Handle single or multi-ticker DataFrame
    if isinstance(df.columns, pd.MultiIndex):
        g = None
        for level in (0, 1):
            try:
                g = df.xs(sym, axis=1, level=level)
                break
            except Exception:
                pass
        if g is None:
            print(f"[daily upsert] {sym}: could not slice MultiIndex for ticker")
            return 0
    else:
        g = df

    g = g.copy()
    idx = pd.to_datetime(g.index, errors="coerce")
    g = g[~idx.isna()]
    try:
        idx = idx.tz_localize(None)
    except Exception:
        pass
    g.index = idx

    # Normalize column names to our schema
    g = g.rename(columns={
        "Open":"open", "High":"high", "Low":"low", "Close":"close",
        "Adj Close":"adj_close", "Volume":"volume"
    })
    for c in ["open","high","low","close","volume"]:
        if c not in g.columns:
            g[c] = None

    # Use "Date" (capital D) to match PostgreSQL schema, convert to date object
    g["Date"] = g.index.date
    g["symbol"] = sym

    if diag:
        dmin = g["Date"].min() if len(g) else "NA"
        dmax = g["Date"].max() if len(g) else "NA"
        print(f"[daily upsert] {sym}: normalized rows={len(g)}  range=({dmin} -> {dmax})")
        try:
            g.reset_index(drop=True).to_csv(f"daily_{sym}.csv", index=False)
            print(f"[diag] wrote preview CSV -> daily_{sym}.csv")
        except Exception as e:
            print(f"[diag] CSV error: {e}")

    rows = g[["symbol","Date","open","high","low","close","volume"]].values.tolist()
    try:
        with con.cursor() as cur:
            execute_values(cur, """
              INSERT INTO prices_daily(symbol,"Date",open,high,low,close,volume)
              VALUES %s
              ON CONFLICT(symbol,"Date") DO UPDATE SET
                open=EXCLUDED.open, high=EXCLUDED.high, low=EXCLUDED.low, 
                close=EXCLUDED.close, volume=EXCLUDED.volume
            """, rows)
            con.commit()
    except Exception as e:
        print(f"[daily upsert] {sym}: postgres error -> {e}")
        con.rollback()
        return 0
    return len(rows)

# ---------- single-symbol seeder (CLI) ----------
def seed_single_daily(con: psycopg2.extensions.connection, sym: str, period: str = "2y", diag: bool=False) -> None:
    print(stamp(f"SEED DAILY {sym} (period={period})"))
    df = yf.download(sym, period=period, interval="1d", auto_adjust=False, progress=False)
    if df is None or df.empty:
        print(f"[SEED] {sym}: yfinance returned EMPTY — cannot seed")
        return

    idx = pd.to_datetime(df.index, errors="coerce")
    try: idx = idx.tz_localize(None)
    except Exception: pass
    dmin = str(pd.to_datetime(idx.min()).date()) if len(idx) else "NA"
    dmax = str(pd.to_datetime(idx.max()).date()) if len(idx) else "NA"
    print(f"[SEED] {sym}: fetched {len(df)} rows ({dmin} -> {dmax})")

    wrote = upsert_daily(con, sym, df, diag=diag)
    print(f"[SEED] {sym}: inserted/updated {wrote} rows")

    with con.cursor() as cur:
        cur.execute("""
            SELECT COUNT(*), MIN("Date"), MAX("Date")
            FROM prices_daily
            WHERE symbol=%s
        """, (sym,))
        row = cur.fetchone()
        print(f"[SEED CHECK] {sym}: count={row[0]}  range=({row[1]} -> {row[2]})")


def upsert_intraday(con: psycopg2.extensions.connection, sym: str, df: pd.DataFrame, source="yfinance") -> int:
    if df is None or df.empty:
        return 0

    # Slice the multi-index (group_by="ticker") or pass-through single frame
    g = df.xs(sym, axis=1, level=0) if isinstance(df.columns, pd.MultiIndex) else df
    if g is None or g.empty:
        return 0

    g = g.rename(columns={"Open": "open", "High": "high", "Low": "low", "Close": "close", "Volume": "volume"}).copy()

    # Ensure UTC time index
    idx = pd.to_datetime(g.index, errors="coerce")
    if idx.tz is None:
        idx = idx.tz_localize("UTC", nonexistent="shift_forward", ambiguous="NaT")
    else:
        idx = idx.tz_convert("UTC")

    # Convert to pandas Timestamp for PostgreSQL TIMESTAMP WITH TIME ZONE
    # PostgreSQL will handle the timezone conversion
    g["symbol"] = sym.upper()
    g["ts"] = idx

    # Order matches table columns (no date_cst, min_cst, source in PostgreSQL schema)
    rows = g[["symbol","ts","open","high","low","close","volume"]].values.tolist()

    # Use PostgreSQL ON CONFLICT syntax
    try:
        with con.cursor() as cur:
            execute_values(cur, """
                INSERT INTO prices_intraday_1m
                  (symbol, ts, open, high, low, close, volume)
                VALUES %s
                ON CONFLICT (symbol, ts) DO UPDATE SET
                  open=EXCLUDED.open, high=EXCLUDED.high, low=EXCLUDED.low,
                  close=EXCLUDED.close, volume=EXCLUDED.volume
            """, rows)
            con.commit()
    except Exception as e:
        print(f"[intraday upsert] {sym}: postgres error -> {e}")
        con.rollback()
        return 0

    return len(g)



# ---------- runners ----------
def run_intraday_batched(con, symbols, lookback_days, batch_size, sleep_between):
    # Determine interval based on lookback days (Yahoo Finance limitations)
    # 1m: Last 7 days maximum
    # 5m: Last 60 days maximum
    # 15m: Last 60 days maximum
    if lookback_days <= 7:
        interval = "1m"
        period = f"{lookback_days}d"
    elif lookback_days <= 60:
        interval = "5m"  # 5-minute bars for longer lookback (30 days)
        period = f"{lookback_days}d"
    else:
        interval = "15m"  # 15-minute for very long lookback
        period = "60d"  # Max 60 days for 15m
    
    total_rows, errors = 0, 0
    print(stamp(f"INTRADAY START last {lookback_days}d using {interval} interval"))
    print(f"Universe size: {len(symbols)}; batch={batch_size}", flush=True)

    for bi, batch in enumerate(chunks(symbols, batch_size), 1):
        tag = f"[batch {bi}/{(len(symbols)+batch_size-1)//batch_size}]"
        print(stamp(f"{tag} download {len(batch)} tickers"), flush=True)
        try:
            df = yf_download_multi(batch, period=period, interval=interval)
        except KeyboardInterrupt:
            print("\n[interrupt] stopping intraday loop gracefully."); break
        except Exception as e:
            errors += len(batch); print(f"{tag} download error: {e}")
            time.sleep(max(0.2, sleep_between)); continue

        try:
            # Process batch - each upsert commits individually
            for sym in batch:
                try:
                    total_rows += upsert_intraday(con, sym, df)
                except Exception as e:
                    errors += 1; print(f"{tag} upsert error {sym}: {e}")
        except KeyboardInterrupt:
            print("\n[interrupt] stopping mid-batch."); break

        print(stamp(f"{tag} done; cumulative rows={total_rows}, errors={errors}"), flush=True)
        time.sleep(max(0.2, sleep_between))

    print(stamp("INTRADAY END"))
    print(f"[INTRADAY] rows upserted: {total_rows}, errors: {errors}")
    return total_rows, errors


def _daily_count(con: psycopg2.extensions.connection, sym: str) -> int:
    with con.cursor() as cur:
        cur.execute("SELECT COUNT(*) FROM prices_daily WHERE symbol=%s", (sym,))
        row = cur.fetchone()
        return int(row[0] or 0)

def run_daily_incremental(con: psycopg2.extensions.connection, symbols: List[str], period: str) -> Tuple[int,int]:
    print(stamp(f"DAILY START (period={period})"))
    total, errors = 0, 0
    for i, sym in enumerate(symbols, 1):
        try:
            have = _daily_count(con, sym)
            need_seed = have < 10
            last = None if need_seed else get_last_daily_date(con, sym)

            if last is None:
                mode = "FULL"
                df = yf.download(sym, period=period, interval="1d",
                                 auto_adjust=False, progress=False)
            else:
                mode = f"INC from {last} -5d"
                start_dt = (datetime.strptime(last, "%Y-%m-%d") - timedelta(days=5)).date().isoformat()
                df = yf.download(sym, start=start_dt, interval="1d",
                                 auto_adjust=False, progress=False)

            if df is None or df.empty:
                print(f"[DAILY {mode}] {sym}: yfinance returned EMPTY (have={have})")
            else:
                i0 = pd.to_datetime(df.index, errors="coerce")
                try: i0 = i0.tz_localize(None)
                except Exception: pass
                dmin = str(pd.to_datetime(i0.min()).date()) if len(i0) else "NA"
                dmax = str(pd.to_datetime(i0.max()).date()) if len(i0) else "NA"
                print(f"[DAILY {mode}] {sym}: fetched {len(df)} rows ({dmin} -> {dmax}), have={have}")

            wrote = upsert_daily(con, sym, df)
            total += wrote

            if i % 50 == 0:
                print(stamp(f"DAILY progress {i}/{len(symbols)} (rows={total})"), flush=True)

        except Exception as e:
            errors += 1
            print(f"[daily error] {sym}: {e}")
        time.sleep(0.02)
    # Note: upsert_daily commits individually, no need for final commit
    print(stamp("DAILY END"))
    print(f"[DAILY] rows upserted: {total}, errors: {errors}")
    return total, errors


# ---------- single-symbol seeder (CLI) ----------
# ---------- main loop ----------
def main():
    db_info = os.getenv('DATABASE_URL', 'local PostgreSQL')
    print(stamp(f"BOOT — PostgreSQL: {db_info}"))
    con = connect_rw()
    try:
        ensure_prices_daily_schema(con)
        ensure_intraday_schema(con)

        symbols, sector_map = build_ticker_universe(SEGMENTS)
        ensure_meta(con)
        upsert_universe(con, symbols, sector_map)
        print(f"Total tickers in universe: {len(symbols)}")

        if not LOOP:
            run_intraday_batched(con, symbols, LOOKBACK_DAYS_INTRADAY, BATCH_SIZE, SLEEP_BETWEEN_BATCHES)
            run_daily_incremental(con, symbols, DAILY_PERIOD)
            print(stamp("ONE-SHOT DONE"))
            return 0

        last_daily = 0.0
        cycle = 0
        while True:
            cycle += 1
            print("\n" + "="*7 + f" LOOP cycle {cycle} " + "="*7)
            print(stamp("LOOP START"))

            run_intraday_batched(con, symbols, LOOKBACK_DAYS_INTRADAY, BATCH_SIZE, SLEEP_BETWEEN_BATCHES)

            now = time.time()
            if DAILY_INTERVAL_MIN == 0 or now - last_daily >= DAILY_INTERVAL_MIN * 60:
                run_daily_incremental(con, symbols, DAILY_PERIOD)
                last_daily = time.time()

            if INTRADAY_INTERVAL_MIN <= 0:
                time.sleep(1.0)
            else:
                interval = INTRADAY_INTERVAL_MIN * 60
                now = time.time()
                next_ts = (int(now // interval) + 1) * interval
                jitter = random.uniform(0, 4)
                delay = max(0.0, (next_ts - now) + jitter)
                print(stamp(f"SLEEP {delay:.1f}s until next intraday run"))
                time.sleep(delay)

    finally:
        con.close()
        print(stamp("CLOSED DB"))


# ---------- CLI entry ----------
def cli_entry():
    import argparse
    parser = argparse.ArgumentParser(description="Load prices to PostgreSQL (daily + intraday)")
    parser.add_argument("--only", help="Run only for one symbol (e.g. --only ASTS)")
    parser.add_argument("--daily", action="store_true", help="Run only daily for the symbol (with --only)")
    parser.add_argument("--intraday", action="store_true", help="Run only intraday for the symbol (with --only)")
    parser.add_argument("--diag", action="store_true", help="Diagnostic: print ranges and write CSV previews")
    args = parser.parse_args()

    db_info = os.getenv('DATABASE_URL', 'local PostgreSQL')
    print(stamp(f"BOOT — PostgreSQL: {db_info}"))
    con = connect_rw()
    ensure_prices_daily_schema(con)
    ensure_intraday_schema(con)

    if args.only:
        sym = args.only.strip().upper()
        if args.daily:
            seed_single_daily(con, sym, DAILY_PERIOD, diag=args.diag)
        elif args.intraday:
            # diag not needed for intraday path
            run_intraday_batched(con, [sym], LOOKBACK_DAYS_INTRADAY, 1, 0.1)
        else:
            seed_single_daily(con, sym, DAILY_PERIOD, diag=args.diag)
            run_intraday_batched(con, [sym], LOOKBACK_DAYS_INTRADAY, 1, 0.1)
        con.close()
        print(stamp("DONE (CLI one-shot)"))
        sys.exit(0)

    # Default = full loop
    rc = main()
    sys.exit(rc if isinstance(rc, int) else 0)


if __name__ == "__main__":
    cli_entry()