fithia2/apps/tools/build_dividend_pit_cache.py

"""Build a point-in-time dividend calendar cache from Oracle or Alpha Vantage."""
from __future__ import annotations

import argparse
import json
import os
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from itertools import islice
from pathlib import Path
from typing import Any

import pyarrow as pa
import pyarrow.parquet as pq
import requests

from apps.backtester.run import _build_merged_snapshot_store, load_manifest, resolve_config
from libs.common.logging import configure_logging


def _symbol_universe_from_config(config_path: str) -> list[str]:
    manifest = load_manifest(config_path)
    config = resolve_config(manifest)
    store = _build_merged_snapshot_store(
        manifest,
        config,
        snapshot_dir_override=None,
    )
    return sorted(str(symbol).upper() for symbol in store._bars.keys())


def _load_existing_records(path: Path) -> list[dict[str, Any]]:
    if not path.exists():
        return []
    table = pq.read_table(str(path))
    return table.to_pylist()


def _write_records(path: Path, rows: list[dict[str, Any]]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    table = pa.Table.from_pylist(rows)
    pq.write_table(table, str(path))


def _fetch_symbol_dividends(
    *,
    symbol: str,
    api_key: str,
    session: requests.Session,
) -> list[dict[str, Any]]:
    response = session.get(
        "https://www.alphavantage.co/query",
        params={
            "function": "DIVIDENDS",
            "symbol": symbol,
            "apikey": api_key,
        },
        timeout=30,
    )
    response.raise_for_status()
    payload = response.json()
    if isinstance(payload, dict) and payload.get("Note"):
        raise RuntimeError(str(payload["Note"]))
    if isinstance(payload, dict) and payload.get("Information"):
        raise RuntimeError(str(payload["Information"]))
    rows = payload.get("data") if isinstance(payload, dict) else None
    if not isinstance(rows, list):
        return []
    normalized: list[dict[str, Any]] = []
    for row in rows:
        declaration_date = str(row.get("declaration_date") or "").strip()
        ex_dividend_date = str(row.get("ex_dividend_date") or "").strip()
        amount_raw = row.get("amount")
        if not declaration_date or not ex_dividend_date or amount_raw is None:
            continue
        try:
            amount = float(amount_raw)
        except (TypeError, ValueError):
            continue
        normalized.append(
            {
                "symbol": symbol,
                "as_of_date": declaration_date[:10],
                "declaration_date": declaration_date[:10],
                "ex_dividend_date": ex_dividend_date[:10],
                "record_date": str(row.get("record_date") or "")[:10] or None,
                "payment_date": str(row.get("payment_date") or "")[:10] or None,
                "amount": amount,
                "source": "alphavantage_dividends",
            }
        )
    return normalized


def _fetch_symbol_dividends_from_oracle(
    *,
    symbol: str,
    base_url: str,
    session: requests.Session,
    force_refresh: bool,
) -> list[dict[str, Any]]:
    response = session.get(
        f"{base_url.rstrip('/')}/api/v1/dividends/history/{symbol}",
        params={
            "limit": 1000,
            "force_refresh": force_refresh,
        },
        timeout=60,
    )
    response.raise_for_status()
    payload = response.json()
    rows = payload.get("dividends") if isinstance(payload, dict) else None
    if not isinstance(rows, list):
        return []
    normalized: list[dict[str, Any]] = []
    for row in rows:
        ex_dividend_date = str(row.get("ex_dividend_date") or "").strip()
        as_of_date = str(row.get("as_of_date") or "").strip()
        amount_raw = row.get("amount")
        if not as_of_date or not ex_dividend_date or amount_raw is None:
            continue
        try:
            amount = float(amount_raw)
        except (TypeError, ValueError):
            continue
        normalized.append(
            {
                "symbol": str(row.get("ticker") or symbol).strip().upper(),
                "as_of_date": as_of_date[:10],
                "declaration_date": str(row.get("declaration_date") or "")[:10] or None,
                "ex_dividend_date": ex_dividend_date[:10],
                "record_date": str(row.get("record_date") or "")[:10] or None,
                "payment_date": str(row.get("payment_date") or "")[:10] or None,
                "amount": amount,
                "currency": str(row.get("currency") or "USD"),
                "dividend_type": str(row.get("dividend_type") or "regular"),
                "frequency": str(row.get("frequency") or "") or None,
                "source": str(row.get("source") or "oracle_dividends"),
            }
        )
    return normalized


def _chunked(items: list[str], size: int) -> list[list[str]]:
    iterator = iter(items)
    chunks: list[list[str]] = []
    while True:
        chunk = list(islice(iterator, size))
        if not chunk:
            break
        chunks.append(chunk)
    return chunks


def _ingest_oracle_dividends(
    *,
    symbols: list[str],
    base_url: str,
    session: requests.Session,
    force_refresh: bool,
    batch_size: int,
) -> None:
    for chunk in _chunked(symbols, max(1, int(batch_size))):
        response = session.post(
            f"{base_url.rstrip('/')}/api/v1/dividends/admin/ingest",
            json={
                "symbols": chunk,
                "force_refresh": force_refresh,
            },
            timeout=120,
        )
        response.raise_for_status()


def main() -> None:
    parser = argparse.ArgumentParser(description="Build PIT dividend calendar cache from Oracle or Alpha Vantage")
    parser.add_argument("--config", action="append", default=[])
    parser.add_argument("--symbol", action="append", default=[])
    parser.add_argument("--output", default="data/reference/dividend_calendar_pit.parquet")
    parser.add_argument("--raw-cache-dir", default=None)
    parser.add_argument("--limit", type=int, default=None)
    parser.add_argument("--sleep-seconds", type=float, default=12.5)
    parser.add_argument("--force-refresh", action="store_true")
    parser.add_argument("--provider", choices=["oracle", "alphavantage"], default="oracle")
    parser.add_argument("--oracle-base-url", default=os.environ.get("STOCK_ORACLE_URL") or "http://localhost:18001")
    parser.add_argument("--ingest-first", action="store_true")
    parser.add_argument("--ingest-batch-size", type=int, default=200)
    parser.add_argument("--max-workers", type=int, default=8)
    parser.add_argument("--api-key", default=None)
    args = parser.parse_args()

    configure_logging("INFO")
    api_key = ""
    if args.provider == "alphavantage":
        api_key = args.api_key or os.environ.get("ALPHAVANTAGE_API_KEY") or ""
        if not api_key:
            raise SystemExit("Missing Alpha Vantage API key. Pass --api-key or set ALPHAVANTAGE_API_KEY.")

    symbols: set[str] = {str(symbol).strip().upper() for symbol in args.symbol if str(symbol).strip()}
    for config_path in args.config:
        symbols.update(_symbol_universe_from_config(config_path))
    ordered_symbols = sorted(symbols)
    if args.limit is not None:
        ordered_symbols = ordered_symbols[: max(0, int(args.limit))]

    output_path = Path(args.output)
    raw_cache_dir = Path(
        args.raw_cache_dir
        or (
            "data/cache/oracle_dividends_raw"
            if args.provider == "oracle"
            else "data/cache/alpha_vantage_dividends_raw"
        )
    )
    raw_cache_dir.mkdir(parents=True, exist_ok=True)

    existing_rows = _load_existing_records(output_path)
    existing_by_symbol: dict[str, list[dict[str, Any]]] = {}
    for row in existing_rows:
        symbol = str(row.get("symbol") or "").strip().upper()
        if symbol:
            existing_by_symbol.setdefault(symbol, []).append(row)

    session = requests.Session()
    if args.provider == "oracle" and args.ingest_first and ordered_symbols:
        _ingest_oracle_dividends(
            symbols=ordered_symbols,
            base_url=str(args.oracle_base_url),
            session=session,
            force_refresh=bool(args.force_refresh),
            batch_size=int(args.ingest_batch_size),
        )
    updated_symbols = 0
    failed_symbols: list[str] = []
    cache_hits = [
        symbol
        for symbol in ordered_symbols
        if (raw_cache_dir / f"{symbol}.json").exists() and not args.force_refresh
    ]
    fetch_symbols = [symbol for symbol in ordered_symbols if symbol not in set(cache_hits)]

    for symbol in cache_hits:
        raw_path = raw_cache_dir / f"{symbol}.json"
        existing_by_symbol[symbol] = json.loads(raw_path.read_text())
        updated_symbols += 1

    if args.provider == "oracle" and fetch_symbols:
        max_workers = max(1, int(args.max_workers))

        def _oracle_fetch(symbol: str) -> tuple[str, list[dict[str, Any]]]:
            local_session = requests.Session()
            rows = _fetch_symbol_dividends_from_oracle(
                symbol=symbol,
                base_url=str(args.oracle_base_url),
                session=local_session,
                force_refresh=bool(args.force_refresh),
            )
            return symbol, rows

        with ThreadPoolExecutor(max_workers=max_workers) as executor:
            future_map = {executor.submit(_oracle_fetch, symbol): symbol for symbol in fetch_symbols}
            for future in as_completed(future_map):
                symbol = future_map[future]
                try:
                    _, payload_rows = future.result()
                except Exception:
                    failed_symbols.append(symbol)
                    continue
                raw_path = raw_cache_dir / f"{symbol}.json"
                raw_path.write_text(json.dumps(payload_rows, indent=2))
                existing_by_symbol[symbol] = payload_rows
                updated_symbols += 1
    else:
        for idx, symbol in enumerate(fetch_symbols):
            raw_path = raw_cache_dir / f"{symbol}.json"
            try:
                payload_rows = _fetch_symbol_dividends(symbol=symbol, api_key=api_key, session=session)
            except Exception:
                failed_symbols.append(symbol)
                continue
            raw_path.write_text(json.dumps(payload_rows, indent=2))
            should_sleep = args.provider == "alphavantage" and args.sleep_seconds > 0
            if idx < len(fetch_symbols) - 1 and should_sleep:
                time.sleep(args.sleep_seconds)
            existing_by_symbol[symbol] = payload_rows
            updated_symbols += 1

    merged_rows: list[dict[str, Any]] = []
    for symbol in sorted(existing_by_symbol):
        merged_rows.extend(existing_by_symbol[symbol])
    merged_rows.sort(key=lambda row: (row["symbol"], row["as_of_date"], row["ex_dividend_date"]))
    _write_records(output_path, merged_rows)
    print(
        json.dumps(
            {
                "output": str(output_path),
                "provider": args.provider,
                "symbols_requested": len(ordered_symbols),
                "symbols_written": updated_symbols,
                "symbols_failed": len(failed_symbols),
                "failed_symbols_preview": failed_symbols[:20],
                "rows_written": len(merged_rows),
            },
            indent=2,
        )
    )


if __name__ == "__main__":
    main()