backend/startup_utils.py

from __future__ import annotations

from datetime import datetime, timezone
from typing import Any, Awaitable, Callable, Dict, List, MutableMapping, Protocol, Sequence, Tuple


class LoggerLike(Protocol):
    def info(self, msg: str, *args: Any, **kwargs: Any) -> None: ...

    def warning(self, msg: str, *args: Any, **kwargs: Any) -> None: ...


class AsyncHttpClientLike(Protocol):
    async def get(self, url: str) -> Any: ...


class ForecasterLike(Protocol):
    is_ready: bool
    device: str

    def _lazy_load(self) -> Awaitable[None]: ...


def clear_stale_ip_limits(
    ip_limits: MutableMapping[str, List[float]],
    *,
    now: float,
    window_seconds: int = 60,
) -> int:
    stale_ips = [
        ip
        for ip, timestamps in ip_limits.items()
        if not any(now - timestamp < window_seconds for timestamp in timestamps)
    ]
    for ip in stale_ips:
        del ip_limits[ip]
    return len(stale_ips)


def build_source_selftest_urls(
    twelvedata_api_key: str | None,
    finnhub_api_key: str | None,
) -> List[Tuple[str, str]]:
    return [
        ("binance", "https://api.binance.com/api/v3/ping"),
        ("bybit", "https://api.bybit.com/v5/market/time"),
        ("coingecko", "https://api.coingecko.com/api/v3/ping"),
        ("twelvedata", f"https://api.twelvedata.com/api_usage?apikey={twelvedata_api_key or ''}"),
        ("finnhub", f"https://finnhub.io/api/v1/quote?symbol=AAPL&token={finnhub_api_key or ''}"),
    ]


async def run_source_selftest(
    *,
    client: AsyncHttpClientLike,
    tests: Sequence[Tuple[str, str]],
    startup_sources: MutableMapping[str, Dict[str, Any]],
    logger: LoggerLike,
    timestamp_provider: Callable[[], str] | None = None,
) -> None:
    get_timestamp = timestamp_provider or (lambda: datetime.now(timezone.utc).isoformat())
    for name, url in tests:
        try:
            response = await client.get(url)
            startup_sources[name] = {
                "reachable": True,
                "status_code": response.status_code,
                "checked_at": get_timestamp(),
            }
            logger.info("[selftest] %-15s HTTP %d", name, response.status_code)
        except Exception as ex:
            startup_sources[name] = {
                "reachable": False,
                "error": str(ex),
                "checked_at": get_timestamp(),
            }
            logger.warning("[selftest] %-15s FAILED: %s", name, ex)


async def warmup_timesfm(
    *,
    forecaster: ForecasterLike,
    startup_timesfm_state: MutableMapping[str, Any],
    logger: LoggerLike,
) -> None:
    await warmup_forecaster(
        forecaster=forecaster,
        startup_state=startup_timesfm_state,
        logger=logger,
        label="TimesFM",
    )


async def warmup_forecaster(
    *,
    forecaster: ForecasterLike,
    startup_state: MutableMapping[str, Any],
    logger: LoggerLike,
    label: str,
) -> None:
    startup_state["warming"] = True
    startup_state["last_error"] = None
    try:
        await forecaster._lazy_load()
        startup_state["loaded"] = forecaster.is_ready
        startup_state["device"] = forecaster.device
        logger.info("[startup] %s warmup finished on %s", label, forecaster.device)
    except Exception as ex:
        startup_state["loaded"] = False
        startup_state["device"] = forecaster.device
        startup_state["last_error"] = str(ex)
        logger.warning("[startup] %s warmup failed: %s", label, ex)
    finally:
        startup_state["warming"] = False