src/data.py

"""OHLCV acquisition: provider chain, rate limiting, validation, refresh.

Two hard rules enforced here:

1. **Only the batch refresh path touches an external provider.** User-facing
   requests read exclusively from the cached store. `allow_network()` gates
   every outbound call and is off unless a refresh explicitly opens it.
2. **Never silently return partial data.** Every fetch reports what it actually
   got versus what was asked for; short coverage becomes an honest boundary in
   the manifest, never a silent truncation and never an error.

Providers are declared in `config.PROVIDER_CHAIN`; adding one means adding a
spec plus a fetch function to `_FETCHERS`, not restructuring the chain.
"""

from __future__ import annotations

import contextlib
import io
import logging
import os
import random
import threading
import time
from dataclasses import dataclass, field

import pandas as pd

from . import config
from .config import Asset, ProviderSpec
from .store import PRICE_COLUMNS, SignalStore, _iso, _utc, validate_price_frame

log = logging.getLogger("bit.data")

# --------------------------------------------------------------------------
# Network gate
# --------------------------------------------------------------------------

_network = threading.local()


def network_allowed() -> bool:
    return getattr(_network, "allowed", False)


@contextlib.contextmanager
def allow_network():
    """Open the gate for a batch refresh. Scoped to the calling thread."""
    prev = getattr(_network, "allowed", False)
    _network.allowed = True
    try:
        yield
    finally:
        _network.allowed = prev


class ProviderError(RuntimeError):
    """A provider failed in a way that should advance the chain."""


class NetworkNotAllowed(RuntimeError):
    """A user-facing code path tried to reach an external provider."""


# --------------------------------------------------------------------------
# Rate limiting + backoff
# --------------------------------------------------------------------------


class RateLimiter:
    """Per-provider minimum spacing between outbound calls, process-wide."""

    _locks: dict[str, threading.Lock] = {}
    _last: dict[str, float] = {}
    _guard = threading.Lock()

    @classmethod
    def wait(cls, name: str, min_interval_s: float) -> None:
        with cls._guard:
            lock = cls._locks.setdefault(name, threading.Lock())
        with lock:
            last = cls._last.get(name, 0.0)
            delta = time.monotonic() - last
            if delta < min_interval_s:
                time.sleep(min_interval_s - delta)
            cls._last[name] = time.monotonic()


def with_backoff(spec: ProviderSpec, fn, *args, **kwargs):
    """Run `fn` under the provider's rate limit with exponential backoff."""
    last_err: Exception | None = None
    for attempt in range(spec.max_retries):
        RateLimiter.wait(spec.name, spec.min_interval_s)
        try:
            return fn(*args, **kwargs)
        except Exception as e:  # provider libs raise a wide variety
            last_err = e
            if attempt == spec.max_retries - 1:
                break
            sleep_s = (spec.backoff_base_s ** attempt) + random.uniform(0, 0.4)
            log.warning(
                "provider %s attempt %d/%d failed (%s); backing off %.1fs",
                spec.name, attempt + 1, spec.max_retries, type(e).__name__, sleep_s,
            )
            time.sleep(sleep_s)
    raise ProviderError(f"{spec.name} exhausted retries: {last_err}") from last_err


# --------------------------------------------------------------------------
# Fetch result
# --------------------------------------------------------------------------


@dataclass
class FetchResult:
    frame: pd.DataFrame
    source: str
    requested_start: pd.Timestamp
    requested_end: pd.Timestamp
    # True when the provider's own history depth, not our request, set the floor.
    truncated_by_provider: bool = False
    provider_max_days: int | None = None
    notes: list[str] = field(default_factory=list)

    @property
    def rows(self) -> int:
        return len(self.frame)

    @property
    def actual_start(self) -> pd.Timestamp | None:
        return None if self.frame.empty else _utc(self.frame["ts"].iloc[0])

    @property
    def actual_end(self) -> pd.Timestamp | None:
        return None if self.frame.empty else _utc(self.frame["ts"].iloc[-1])


def _frame(rows: list[dict], source: str) -> pd.DataFrame:
    df = pd.DataFrame(rows, columns=["ts", "open", "high", "low", "close", "volume"])
    df["source"] = source
    if not df.empty:
        df["ts"] = df["ts"].map(_utc)
        df = df.drop_duplicates(subset="ts", keep="last").sort_values("ts")
    return df.loc[:, PRICE_COLUMNS].reset_index(drop=True)


# --------------------------------------------------------------------------
# Providers
# --------------------------------------------------------------------------


def _ccxt_exchange(name: str):
    import ccxt

    klass = getattr(ccxt, name)
    ex = klass({"enableRateLimit": True, "timeout": 20000})
    return ex


def _fetch_ccxt(exchange_name: str, spec: ProviderSpec, asset: Asset,
                tf: str, start, end) -> pd.DataFrame:
    if not asset.ccxt_symbol:
        raise ProviderError(f"{asset.slug} has no ccxt symbol")
    ex = _ccxt_exchange(exchange_name)
    ccxt_tf = config.TIMEFRAMES[tf].ccxt_tf
    step_ms = config.TIMEFRAMES[tf].minutes * 60_000
    since = int(_utc(start).timestamp() * 1000)
    end_ms = int(_utc(end).timestamp() * 1000)
    rows: list[dict] = []
    guard = 0

    while since < end_ms and guard < 4000:
        guard += 1
        batch = with_backoff(
            spec, ex.fetch_ohlcv, asset.ccxt_symbol, ccxt_tf, since, 1000
        )
        if not batch:
            break
        for ts, o, h, l, c, v in batch:
            if ts > end_ms:
                break
            rows.append({"ts": pd.Timestamp(ts, unit="ms", tz="UTC"),
                         "open": o, "high": h, "low": l, "close": c, "volume": v})
        last = batch[-1][0]
        if last <= since:
            break
        since = last + step_ms

    with contextlib.suppress(Exception):
        ex.close()
    return _frame(rows, exchange_name)


def _fetch_binance(spec, asset, tf, start, end):
    return _fetch_ccxt("binance", spec, asset, tf, start, end)


def _fetch_coinbase(spec, asset, tf, start, end):
    return _fetch_ccxt("coinbase", spec, asset, tf, start, end)


def _fetch_yfinance(spec: ProviderSpec, asset: Asset, tf: str, start, end) -> pd.DataFrame:
    if not asset.yahoo_symbol:
        raise ProviderError(f"{asset.slug} has no Yahoo symbol")
    import yfinance as yf

    interval = config.TIMEFRAMES[tf].yahoo_interval
    max_days = config.TIMEFRAMES[tf].yahoo_max_days
    s, e = _utc(start), _utc(end)
    if max_days is not None:
        floor = pd.Timestamp.now(tz="UTC") - pd.Timedelta(days=max_days - 1)
        s = max(s, floor)  # honest boundary; see refresh() notes

    def _call():
        return yf.download(
            asset.yahoo_symbol, start=s.date(), end=(e + pd.Timedelta(days=1)).date(),
            interval=interval, auto_adjust=False, progress=False, threads=False,
        )

    raw = with_backoff(spec, _call)
    if raw is None or raw.empty:
        raise ProviderError("yfinance returned no rows")
    if isinstance(raw.columns, pd.MultiIndex):
        raw.columns = raw.columns.get_level_values(0)
    # reset_index first: the timestamp arrives as the index ("Date"/"Datetime")
    # and only becomes a column here, so lowercasing must happen afterwards.
    raw = raw.reset_index()
    raw.columns = [str(c).lower() for c in raw.columns]
    tcol = next((c for c in ("datetime", "date", "index") if c in raw.columns), None)
    if tcol is None:
        raise ProviderError(f"yfinance frame has no timestamp column: {list(raw.columns)}")
    rows = [
        {"ts": r[tcol], "open": r["open"], "high": r["high"],
         "low": r["low"], "close": r["close"], "volume": r.get("volume", 0.0)}
        for _, r in raw.iterrows()
    ]
    return _frame(rows, "yfinance")


def _fetch_stooq(spec: ProviderSpec, asset: Asset, tf: str, start, end) -> pd.DataFrame:
    """Stooq CSV endpoint -- no key, daily only."""
    if tf != "1d":
        raise ProviderError("stooq serves daily bars only")
    if not asset.stooq_symbol:
        raise ProviderError(f"{asset.slug} has no Stooq symbol")
    import requests

    url = f"https://stooq.com/q/d/l/?s={asset.stooq_symbol}&i=d"

    def _call():
        r = requests.get(url, timeout=20)
        r.raise_for_status()
        if "Date" not in r.text[:64]:
            raise ProviderError("stooq returned no CSV header (rate limited?)")
        return r.text

    text = with_backoff(spec, _call)
    raw = pd.read_csv(io.StringIO(text))
    raw.columns = [c.lower() for c in raw.columns]
    raw = raw.dropna(subset=["open", "high", "low", "close"])
    s, e = _utc(start), _utc(end)
    rows = []
    for _, r in raw.iterrows():
        ts = _utc(r["date"])
        if ts < s or ts > e:
            continue
        rows.append({"ts": ts, "open": r["open"], "high": r["high"],
                     "low": r["low"], "close": r["close"], "volume": r.get("volume", 0.0)})
    return _frame(rows, "stooq")


def _fetch_tiingo(spec: ProviderSpec, asset: Asset, tf: str, start, end) -> pd.DataFrame:
    if tf != "1d":
        raise ProviderError("tiingo adapter covers daily bars only")
    key = os.environ.get("TIINGO_KEY")
    if not key:
        raise ProviderError("TIINGO_KEY not set")
    import requests

    sym = asset.tiingo_symbol or asset.slug
    url = f"https://api.tiingo.com/tiingo/daily/{sym}/prices"
    params = {"startDate": _utc(start).date().isoformat(),
              "endDate": _utc(end).date().isoformat(), "token": key}

    def _call():
        r = requests.get(url, params=params, timeout=20)
        r.raise_for_status()
        return r.json()

    payload = with_backoff(spec, _call)
    rows = [
        {"ts": _utc(d["date"]), "open": d["open"], "high": d["high"],
         "low": d["low"], "close": d["close"], "volume": d.get("volume", 0.0)}
        for d in payload
    ]
    return _frame(rows, "tiingo")


_FETCHERS = {
    "binance": _fetch_binance,
    "coinbase": _fetch_coinbase,
    "yfinance": _fetch_yfinance,
    "stooq": _fetch_stooq,
    "tiingo": _fetch_tiingo,
}


# --------------------------------------------------------------------------
# Chain walk
# --------------------------------------------------------------------------


def fetch_ohlcv(asset_slug: str, timeframe: str, start, end) -> FetchResult:
    """Walk the provider chain for `asset_slug` until one returns rows."""
    if not network_allowed():
        raise NetworkNotAllowed(
            "external providers are reachable only from the batch refresh path; "
            "user-facing requests must read from the cached store"
        )
    asset = config.ASSETS.get(asset_slug)
    if asset is None:
        raise ProviderError(f"unknown asset {asset_slug!r}")
    if timeframe not in config.TIMEFRAMES:
        raise ProviderError(f"unknown timeframe {timeframe!r}")

    s, e = _utc(start), _utc(end)
    chain = config.providers_for(asset.kind)
    if not chain:
        raise ProviderError(f"no usable provider for {asset.kind}")

    notes: list[str] = []
    for spec in chain:
        fetcher = _FETCHERS.get(spec.name)
        if fetcher is None:
            continue
        try:
            frame = fetcher(spec, asset, timeframe, s, e)
        except Exception as exc:
            notes.append(f"{spec.name}: {type(exc).__name__}: {exc}")
            log.warning("provider %s failed for %s %s: %s", spec.name, asset_slug, timeframe, exc)
            continue
        if frame.empty:
            notes.append(f"{spec.name}: returned 0 rows")
            continue

        max_days = config.TIMEFRAMES[timeframe].yahoo_max_days if spec.name == "yfinance" else None
        actual_start = _utc(frame["ts"].iloc[0])
        truncated = max_days is not None and actual_start > s + pd.Timedelta(days=1)
        if truncated:
            notes.append(
                f"{spec.name} serves at most ~{max_days}d of {timeframe} bars; "
                f"coverage starts {_iso(actual_start)}"
            )
        return FetchResult(
            frame=frame, source=spec.name, requested_start=s, requested_end=e,
            truncated_by_provider=truncated, provider_max_days=max_days, notes=notes,
        )

    raise ProviderError(
        f"all providers failed for {asset_slug} {timeframe} "
        f"[{_iso(s)} .. {_iso(e)}]: " + " | ".join(notes)
    )


# --------------------------------------------------------------------------
# Refresh
# --------------------------------------------------------------------------


@dataclass
class RefreshReport:
    asset: str
    timeframe: str
    fetched_ranges: list[tuple[str, str]] = field(default_factory=list)
    rows_added: int = 0
    sources: list[str] = field(default_factory=list)
    skipped_cached: bool = False
    gaps: int = 0
    boundary_notes: list[str] = field(default_factory=list)
    errors: list[str] = field(default_factory=list)

    @property
    def ok(self) -> bool:
        return not self.errors

    def summary(self) -> str:
        if self.skipped_cached:
            return f"{self.asset} {self.timeframe}: already cached, nothing fetched"
        if self.errors:
            return f"{self.asset} {self.timeframe}: FAILED -- {'; '.join(self.errors)}"
        return (
            f"{self.asset} {self.timeframe}: +{self.rows_added} rows "
            f"from {','.join(self.sources) or 'n/a'} ({self.gaps} gaps)"
        )


def missing_price_ranges(
    store: SignalStore, asset: str, timeframe: str, start, end
) -> list[tuple[pd.Timestamp, pd.Timestamp]]:
    """Sub-ranges of [start, end] absent from the price cache."""
    s, e = _utc(start), _utc(end)
    if s > e:
        return []
    cov = store.load_manifest().prices.get(f"{asset}|{timeframe}")
    if cov is None or cov.rows == 0:
        return [(s, e)]
    cs, ce = _utc(cov.start_ts), _utc(cov.end_ts)
    step = pd.Timedelta(minutes=config.TIMEFRAMES[timeframe].minutes)
    out = []
    if s < cs:
        out.append((s, min(e, cs - step)))
    if e > ce:
        out.append((max(s, ce + step), e))
    return [(a, b) for a, b in out if a <= b]


def refresh(
    store: SignalStore, asset: str, timeframe: str, start, end, *, strict: bool = False
) -> RefreshReport:
    """Fetch only what the cache is missing, validate it, and write it.

    Never fetches a range the manifest already covers.
    """
    rep = RefreshReport(asset=asset, timeframe=timeframe)
    try:
        gaps = missing_price_ranges(store, asset, timeframe, start, end)
    except Exception as e:
        rep.errors.append(f"coverage lookup failed: {e}")
        return rep

    if not gaps:
        rep.skipped_cached = True
        return rep

    with allow_network():
        for gs, ge in gaps:
            try:
                res = fetch_ohlcv(asset, timeframe, gs, ge)
            except Exception as e:
                rep.errors.append(f"[{_iso(gs)}..{_iso(ge)}] {e}")
                continue

            _, report = validate_price_frame(res.frame, timeframe, strict=strict)
            if report.problems:
                # Partial or dirty data is surfaced, never silently accepted.
                rep.errors.append(
                    f"[{_iso(gs)}..{_iso(ge)}] validation: {'; '.join(report.problems)}"
                )
                continue

            cov = store.write_prices(asset, timeframe, res.frame, strict=strict)
            rep.fetched_ranges.append((_iso(gs), _iso(ge)))
            rep.rows_added += res.rows
            rep.gaps = max(rep.gaps, cov.gaps)
            if res.source not in rep.sources:
                rep.sources.append(res.source)
            rep.boundary_notes.extend(res.notes)
    return rep