src/prices.py

"""Price refresh. Batch path only -- never called while rendering a page.

The Arena reads prices from the shared store and writes new bars back into the
same tree the rest of the estate uses, so a refresh here also benefits the
Backtest Lab. Provider order is crypto via ccxt, equities via yfinance falling
back to Stooq, which is the convention the store was built on.

Two properties matter more than throughput:

**Only closed bars are stored.** A bar for the period currently in progress
will change before it closes, and a forecast issued against it would be
conditioning on a number that later moves -- a subtle lookahead that no
timestamp check would catch.

**A refresh can never shrink the store.** Bars are appended per year file and
deduplicated on timestamp, and a year file is only rewritten when it would gain
rows. That guard matters more than it looks: if the read of the existing
history fails -- an unreachable Hub, an empty local mirror -- the merge sees no
history and the write would replace years of bars with whatever one provider
call returned. That was reproducible, and it is what this refuses.

**Providers fall back.** Binance answers `451` from large parts of the world,
so crypto tries Coinbase after it; equities try yfinance then Stooq. A chain
that stops at the first provider is a chain that works until someone runs it
somewhere else.
"""

from __future__ import annotations

import logging

import pandas as pd

from . import config
from .store import ArenaStore, now_utc

log = logging.getLogger("arena.prices")

PRICE_COLUMNS = ["ts", "open", "high", "low", "close", "volume"]

# ccxt symbols for the crypto assets, and Stooq/yfinance symbols for equities.
# Per-exchange symbols: Coinbase quotes in USD, Binance in USDT.
CCXT_SYMBOLS = {
    "binance": {"BTC-USD": "BTC/USDT", "ETH-USD": "ETH/USDT", "SOL-USD": "SOL/USDT"},
    "coinbase": {"BTC-USD": "BTC/USD", "ETH-USD": "ETH/USD", "SOL-USD": "SOL/USD"},
}
CRYPTO_EXCHANGES = ("binance", "coinbase")
CCXT_TIMEFRAMES = {"1h": "1h", "1d": "1d"}
STOOQ_SYMBOLS = {"SPY": "spy.us", "NVDA": "nvda.us", "QQQ": "qqq.us"}


class ProviderError(RuntimeError):
    pass


def refresh(store: ArenaStore, asset: str, timeframe: str,
            lookback_days: int = 30) -> int:
    """Fetch recent bars and merge them into the store. Returns bars added."""
    spec = config.ASSETS.get(asset)
    if spec is None:
        raise ProviderError(f"unknown asset {asset!r}")

    existing = store.get_prices(asset, timeframe)
    since = None
    if len(existing):
        since = pd.to_datetime(existing["ts"], utc=True).max()

    if spec.asset_class == "crypto":
        fresh, errors = None, []
        for exchange in CRYPTO_EXCHANGES:
            try:
                fresh = _fetch_ccxt(exchange, asset, timeframe, since, lookback_days)
                break
            except Exception as e:
                # Binance answers 451 from restricted regions. That is not a
                # reason to give up on the asset, only on the exchange.
                errors.append(f"{exchange}: {type(e).__name__}")
                log.info("%s failed for %s (%s); trying the next exchange",
                         exchange, asset, type(e).__name__)
        if fresh is None:
            raise ProviderError(f"every crypto provider failed for {asset}: "
                                f"{', '.join(errors)}")
    else:
        try:
            fresh = _fetch_yfinance(asset, timeframe, since, lookback_days)
        except Exception as e:
            log.info("yfinance failed for %s (%s); falling back to Stooq", asset, e)
            fresh = _fetch_stooq(asset, timeframe)

    if fresh is None or not len(fresh):
        return 0

    fresh = _drop_open_bar(fresh, timeframe)
    if since is not None:
        fresh = fresh[pd.to_datetime(fresh["ts"], utc=True) > since]
    if not len(fresh):
        return 0

    added = 0
    combined = pd.concat([existing, fresh], ignore_index=True) if len(existing) else fresh
    combined["ts"] = pd.to_datetime(combined["ts"], utc=True)
    combined = (combined.drop_duplicates(subset=["ts"], keep="first")
                .sort_values("ts").reset_index(drop=True))

    for year, chunk in combined.groupby(combined["ts"].dt.year):
        path = config.prices_path(asset, timeframe, int(year))
        before = store.read_parquet(path)
        have = len(before) if before is not None else 0
        if len(chunk) <= have:
            # Nothing gained. Writing anyway would, in the case where `before`
            # could not be read at all, replace a full year with a fragment.
            if len(chunk) < have:
                log.warning("refusing to shrink %s from %d to %d rows",
                            path, have, len(chunk))
            continue
        store.write_parquet(path, chunk[PRICE_COLUMNS])
        added += len(chunk) - have
    return max(0, added)


def _drop_open_bar(frame: pd.DataFrame, timeframe: str) -> pd.DataFrame:
    """Discard the bar for the period still in progress.

    An in-progress bar's close is not its close. Storing it would let a
    forecast condition on a value that changes afterwards.
    """
    if not len(frame):
        return frame
    minutes = config.TIMEFRAMES[timeframe].minutes
    cutoff = now_utc().floor(f"{minutes}min")
    ts = pd.to_datetime(frame["ts"], utc=True)
    return frame[ts < cutoff].reset_index(drop=True)


def _fetch_ccxt(exchange_name: str, asset: str, timeframe: str, since,
                lookback_days: int) -> pd.DataFrame:
    import ccxt

    symbol = CCXT_SYMBOLS.get(exchange_name, {}).get(asset)
    if not symbol:
        raise ProviderError(f"no {exchange_name} symbol for {asset}")
    tf = CCXT_TIMEFRAMES.get(timeframe)
    if not tf:
        raise ProviderError(f"ccxt cannot serve {timeframe}")

    exchange = getattr(ccxt, exchange_name)({"enableRateLimit": True})
    start = since if since is not None else now_utc() - pd.Timedelta(days=lookback_days)
    ms = int(pd.Timestamp(start).timestamp() * 1000)

    rows = []
    for _ in range(20):          # bounded: 20 pages of 1000 bars is plenty
        batch = exchange.fetch_ohlcv(symbol, tf, since=ms, limit=1000)
        if not batch:
            break
        rows.extend(batch)
        ms = batch[-1][0] + 1
        if len(batch) < 1000:
            break

    if not rows:
        return pd.DataFrame(columns=PRICE_COLUMNS)
    frame = pd.DataFrame(rows, columns=["ts", "open", "high", "low", "close", "volume"])
    frame["ts"] = pd.to_datetime(frame["ts"], unit="ms", utc=True)
    return frame[PRICE_COLUMNS]


def _fetch_yfinance(asset: str, timeframe: str, since, lookback_days: int) -> pd.DataFrame:
    import yfinance as yf

    interval = {"1h": "1h", "1d": "1d"}[timeframe]
    period = f"{max(lookback_days, 7)}d" if timeframe == "1h" else "2y"
    raw = yf.Ticker(asset).history(period=period, interval=interval, auto_adjust=False)
    if raw is None or not len(raw):
        raise ProviderError("yfinance returned no rows")

    frame = raw.reset_index()
    stamp = next((c for c in ("Datetime", "Date", "index") if c in frame.columns), None)
    if stamp is None:
        raise ProviderError(f"no timestamp column in {list(frame.columns)}")
    frame = frame.rename(columns={stamp: "ts", "Open": "open", "High": "high",
                                  "Low": "low", "Close": "close", "Volume": "volume"})
    frame["ts"] = pd.to_datetime(frame["ts"], utc=True)
    return frame[PRICE_COLUMNS]


def _fetch_stooq(asset: str, timeframe: str) -> pd.DataFrame:
    import io

    import requests

    if timeframe != "1d":
        raise ProviderError("Stooq serves daily bars only")
    symbol = STOOQ_SYMBOLS.get(asset)
    if not symbol:
        raise ProviderError(f"no Stooq symbol for {asset}")

    response = requests.get(f"https://stooq.com/q/d/l/?s={symbol}&i=d", timeout=30)
    response.raise_for_status()
    frame = pd.read_csv(io.StringIO(response.text))
    frame = frame.rename(columns={c: c.lower() for c in frame.columns})
    if "date" not in frame.columns:
        raise ProviderError("Stooq returned no date column")
    frame["ts"] = pd.to_datetime(frame["date"], utc=True)
    return frame[PRICE_COLUMNS]