scripts/backfill_chart_svg.py

"""
Backfill chart_svg for cached forecasts that have null chart data.

Processes rows in batches as it scans — starts generating charts immediately
instead of waiting for a full scan. Handles Turso's HTTP timeout by keeping
each query small (200 rows).

Usage:
    cd backend
    python scripts/backfill_chart_svg.py [--dry-run] [--limit N] [--symbol AAPL]
"""

import asyncio
import json
import os
import sys
import time
import argparse
from pathlib import Path

import httpx

sys.path.insert(0, str(Path(__file__).parent.parent))

from dotenv import load_dotenv

load_dotenv()

from app.services.database_service import DatabaseService
from app.services.data_service import StockDataService
from app.services.chart_service import ChartService


async def backfill(
    dry_run: bool = False,
    limit: int = 0,
    symbol_filter: str = "",
):
    turso_url = os.getenv("TURSO_URL")
    turso_token = os.getenv("TURSO_TOKEN")
    if not turso_url or not turso_token:
        print("[FAIL] TURSO_URL and TURSO_TOKEN must be set")
        sys.exit(1)

    db = DatabaseService(url=turso_url, token=turso_token)
    data_service = StockDataService()
    chart_service = ChartService()

    await db.initialize()

    # Increase httpx timeout for bulk backfill operations
    db.client = httpx.AsyncClient(
        base_url=db.https_url,
        headers=dict(db.client.headers),
        timeout=60.0,
    )

    # Process in small batches as we scan — no pre-scan of all rows
    batch_size = 200
    offset = 0
    symbol_clause = "AND symbol = ?" if symbol_filter else ""
    symbol_arg = [symbol_filter.upper()] if symbol_filter else []

    updated = 0
    skipped = 0
    failed = 0
    scanned = 0
    # Cache historical data per symbol to avoid duplicate yfinance fetches
    price_cache: dict[str, object] = {}
    start_time = time.time()

    while True:
        rows = await db._execute(
            f"SELECT id, symbol, horizon_days, forecast_data FROM forecasts"
            f" WHERE 1=1 {symbol_clause}"
            f" ORDER BY id LIMIT ? OFFSET ?",
            symbol_arg + [batch_size, offset],
        )
        if not rows:
            break

        offset += batch_size
        scanned += len(rows)

        for row in rows:
            # Check limit
            if limit > 0 and updated >= limit:
                break

            fd = json.loads(row["forecast_data"])
            if fd.get("chart_svg"):
                continue  # Already has a chart, skip

            symbol = row["symbol"]
            horizon = row["horizon_days"]
            current_price = fd.get("current_price", 0)

            if dry_run:
                print(f"  [DRY] {symbol} h={horizon} — would generate chart")
                updated += 1
                continue

            print(f"  [{scanned}] {symbol} h={horizon}...", end=" ", flush=True)

            # Fetch historical prices (cached per symbol)
            if symbol not in price_cache:
                try:
                    stock_data = await data_service.get_stock_data(symbol, period="2y")
                    if stock_data is not None and len(stock_data) >= 100:
                        price_cache[symbol] = stock_data["Close"].tolist()
                    else:
                        price_cache[symbol] = None
                        print("[SKIP] insufficient data")
                        skipped += 1
                        continue
                    # Rate-limit yfinance fetches
                    await asyncio.sleep(0.5)
                except Exception as e:
                    price_cache[symbol] = None
                    print(f"[SKIP] fetch error: {e}")
                    skipped += 1
                    continue

            historical = price_cache[symbol]
            if historical is None:
                print("[SKIP] no data")
                skipped += 1
                continue

            # Build forecast dict for chart_service
            quantiles = fd.get("quantiles", {})
            p50 = quantiles.get("p50", [])
            forecast = {
                "point_forecast": p50 if p50 else [fd["point_forecast"]],
                "quantiles": quantiles,
            }

            # Generate chart
            try:
                chart_svg = chart_service.generate_forecast_chart(
                    symbol=symbol,
                    historical_prices=historical[-60:],
                    forecast=forecast,
                    current_price=current_price,
                )
            except Exception as e:
                print(f"[FAIL] chart error: {e}")
                failed += 1
                continue

            if not chart_svg:
                print("[FAIL] empty chart")
                failed += 1
                continue

            # Update forecast_data with chart_svg
            fd["chart_svg"] = chart_svg
            try:
                now = int(time.time())
                forecast_json = json.dumps(fd)
                await db._execute(
                    "UPDATE forecasts SET forecast_data = ?, created_at = ? WHERE id = ?",
                    [forecast_json, now, row["id"]],
                )
                print(f"[OK] chart={len(chart_svg)} chars")
                updated += 1
            except Exception as e:
                print(f"[FAIL] db update: {e}")
                failed += 1

        # Check limit after batch
        if limit > 0 and updated >= limit:
            break

        if len(rows) < batch_size:
            break

        # Progress report every 5 batches
        elapsed = time.time() - start_time
        rate = updated / elapsed if elapsed > 0 else 0
        print(f"  --- scanned {scanned}, updated {updated}, skipped {skipped}, failed {failed}, {rate:.1f}/s ---", flush=True)

    await db.close()

    elapsed = time.time() - start_time
    print(f"\n{'='*50}")
    print(f"  Backfill complete! ({elapsed:.0f}s)")
    print(f"  Scanned: {scanned}")
    print(f"  Updated: {updated}")
    print(f"  Skipped: {skipped}")
    print(f"  Failed:  {failed}")
    print(f"  Rate:    {updated/elapsed:.1f} charts/sec")
    print(f"{'='*50}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Backfill chart_svg for cached forecasts")
    parser.add_argument("--dry-run", action="store_true", help="Preview without writing")
    parser.add_argument("--limit", type=int, default=0, help="Max forecasts to process (0=all)")
    parser.add_argument("--symbol", type=str, default="", help="Process only this symbol")
    args = parser.parse_args()

    asyncio.run(backfill(dry_run=args.dry_run, limit=args.limit, symbol_filter=args.symbol))