fithia2/apps/tools/backfill_market_cap.py

"""Backfill market_cap_proxy + exchange_proxy into existing FeatureSnapshot rows.

Patches market_v1 FeatureSnapshots that are missing market_cap_proxy in
their feature_json.  This is needed for paper-trading backsim on events
that were feature-built before the feature-builder started persisting
company metadata.

Usage:
    python -m apps.tools.backfill_market_cap
    python -m apps.tools.backfill_market_cap --start 2026-01-01 --end 2026-03-31
    python -m apps.tools.backfill_market_cap --dry-run
"""
from __future__ import annotations

import argparse
import asyncio
import datetime as dt
import os
from pathlib import Path

_ENV_FILE = Path(__file__).parent.parent.parent / ".env"
if _ENV_FILE.exists():
    for _line in _ENV_FILE.read_text().splitlines():
        _line = _line.strip()
        if _line and not _line.startswith("#") and "=" in _line:
            _k, _, _v = _line.partition("=")
            os.environ.setdefault(_k.strip(), _v.strip())

from rich.console import Console
from rich.progress import track

console = Console(width=120)


async def run_backfill(
    start_date: dt.date | None,
    end_date: dt.date | None,
    dry_run: bool,
) -> None:
    from sqlalchemy import select, update
    from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine

    from libs.db.models import Event, FeatureSnapshot, SymbolMaster
    from libs.oracle_client import CompanyService, ScreenerService, make_oracle_client

    db_dsn = os.environ.get("POSTGRES_DSN", "")
    oracle_url = os.environ.get("STOCK_ORACLE_URL", "http://localhost:18001")

    if not db_dsn:
        console.print("[red]POSTGRES_DSN not set[/]")
        return

    engine = create_async_engine(db_dsn, echo=False)
    async_session = async_sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)

    # ── Step 1: find market_v1 snapshots without market_cap_proxy ──────────
    async with async_session() as session:
        stmt = (
            select(FeatureSnapshot, Event, SymbolMaster)
            .join(Event, FeatureSnapshot.event_id == Event.event_id)
            .outerjoin(SymbolMaster, Event.symbol_id == SymbolMaster.symbol_id)
            .where(FeatureSnapshot.snapshot_name == "market_v1")
        )
        if start_date:
            stmt = stmt.where(Event.event_date >= start_date)
        if end_date:
            stmt = stmt.where(Event.event_date <= end_date)

        rows = (await session.execute(stmt)).all()

    # Filter to those missing market_cap_proxy
    missing = [
        (snap, event, sym)
        for snap, event, sym in rows
        if not (snap.feature_json or {}).get("market_cap_proxy")
    ]
    console.print(f"Found [bold]{len(rows)}[/] market_v1 snapshots, "
                  f"[yellow]{len(missing)}[/] missing market_cap_proxy")

    if not missing or dry_run:
        if dry_run:
            console.print("[dim]Dry-run: no changes written.[/]")
        return

    # ── Step 2: collect unique tickers ─────────────────────────────────────
    ticker_to_snap: dict[str, list[FeatureSnapshot]] = {}
    for snap, event, sym in missing:
        if sym and sym.ticker:
            ticker_to_snap.setdefault(sym.ticker, []).append(snap)

    tickers = sorted(ticker_to_snap)
    console.print(f"Fetching company info for [bold]{len(tickers)}[/] unique tickers…")

    # ── Step 3: batch-fetch via screener first, then per-symbol CompanyService ─
    ticker_mcap: dict[str, float | None] = {t: None for t in tickers}
    ticker_exchange: dict[str, str | None] = {t: None for t in tickers}

    async with make_oracle_client() as client:
        # Try screener first (batch — faster, one call per page)
        try:
            svc = ScreenerService(client)
            stocks = await svc.search_all_stocks(
                market_cap_min=500_000_000,
                exchange="NYSE,NASDAQ,AMEX",
                exclude_types="ETF,FUND,ADR,SPAC",
            )
            screener_lookup = {(s.symbol or "").upper(): s for s in stocks}
            for t in tickers:
                s = screener_lookup.get(t.upper())
                if s:
                    ticker_mcap[t] = s.market_cap
                    ticker_exchange[t] = s.exchange
            resolved = sum(1 for t in tickers if ticker_mcap.get(t) is not None)
            console.print(f"  Screener resolved {resolved}/{len(tickers)}")
        except Exception as exc:
            console.print(f"  [yellow]Screener failed ({exc}), falling back to CompanyService[/]")

        # Per-symbol CompanyService for anything still missing
        unresolved = [t for t in tickers if ticker_mcap.get(t) is None]
        if unresolved:
            company_svc = CompanyService(client)
            semaphore = asyncio.Semaphore(16)

            async def _fetch_one(ticker: str) -> None:
                async with semaphore:
                    try:
                        info = await company_svc.get_company(ticker)
                        ticker_mcap[ticker] = info.market_cap
                        ticker_exchange[ticker] = info.exchange
                    except Exception:
                        pass

            await asyncio.gather(*(_fetch_one(t) for t in unresolved))
            resolved2 = sum(1 for t in unresolved if ticker_mcap.get(t) is not None)
            console.print(f"  CompanyService resolved {resolved2}/{len(unresolved)} remaining")

    # ── Step 4: update feature_json ─────────────────────────────────────────
    updated = 0
    skipped = 0
    async with async_session() as session:
        for ticker, snaps in track(ticker_to_snap.items(), description="Updating…"):
            mcap = ticker_mcap.get(ticker)
            exch = ticker_exchange.get(ticker)
            if mcap is None and exch is None:
                skipped += len(snaps)
                continue
            for snap in snaps:
                fj = dict(snap.feature_json or {})
                if mcap is not None:
                    fj["market_cap_proxy"] = mcap
                if exch is not None:
                    fj["exchange_proxy"] = exch
                await session.execute(
                    update(FeatureSnapshot)
                    .where(FeatureSnapshot.feature_snapshot_id == snap.feature_snapshot_id)
                    .values(feature_json=fj)
                )
                updated += 1

        await session.commit()

    await engine.dispose()
    console.print(f"\n[bold green]Done.[/] Updated {updated} snapshots, skipped {skipped} (no data).")


def main() -> None:
    parser = argparse.ArgumentParser(description="Backfill market_cap_proxy into FeatureSnapshots")
    parser.add_argument("--start", default=None, metavar="YYYY-MM-DD")
    parser.add_argument("--end", default=None, metavar="YYYY-MM-DD")
    parser.add_argument("--dry-run", action="store_true", help="Show what would be done without writing")
    args = parser.parse_args()

    start = dt.date.fromisoformat(args.start) if args.start else None
    end = dt.date.fromisoformat(args.end) if args.end else None
    asyncio.run(run_backfill(start, end, args.dry_run))


if __name__ == "__main__":
    main()