TradeBot/tools/sync_orderbook_observations.py

from __future__ import annotations

import argparse
import base64
import json
import os
import sqlite3
from pathlib import Path
from typing import Any
from urllib.error import HTTPError, URLError
from urllib.parse import urlencode
from urllib.request import Request, urlopen


def sync_orderbook_observations(
    *,
    api_base_url: str,
    token: str,
    database_path: str | Path,
    timeout: int = 60,
    page_limit: int = 5000,
) -> dict[str, Any]:
    path = Path(database_path)
    path.parent.mkdir(parents=True, exist_ok=True)
    _init_schema(path)
    manifest = _get_json(
        api_base_url,
        "/api/training/market-observations/manifest",
        token=token,
        timeout=timeout,
    )
    rows = manifest.get("items") if isinstance(manifest.get("items"), list) else []
    downloaded = 0
    symbol_results: list[dict[str, Any]] = []
    for row in rows:
        if not isinstance(row, dict):
            continue
        symbol = str(row.get("symbol") or "").strip().upper()
        remote_max_id = int(row.get("max_id", 0) or 0)
        if not symbol or remote_max_id <= 0:
            continue
        after_id = _local_max_id(path, symbol)
        symbol_downloaded = 0
        while after_id < remote_max_id:
            query = urlencode(
                {
                    "symbol": symbol,
                    "after_id": after_id,
                    "limit": max(1, min(5000, int(page_limit))),
                }
            )
            payload = _get_json(
                api_base_url,
                f"/api/training/market-observations?{query}",
                token=token,
                timeout=timeout,
            )
            items = payload.get("items") if isinstance(payload.get("items"), list) else []
            if not items:
                break
            inserted = _insert_rows(path, items)
            symbol_downloaded += inserted
            downloaded += inserted
            next_after_id = int(payload.get("next_after_id", after_id) or after_id)
            if next_after_id <= after_id:
                break
            after_id = next_after_id
        symbol_results.append(
            {
                "symbol": symbol,
                "downloaded": symbol_downloaded,
                "local_max_id": _local_max_id(path, symbol),
                "remote_max_id": remote_max_id,
            }
        )
    return {
        "database_path": str(path.resolve()),
        "downloaded": downloaded,
        "symbols": symbol_results,
        "local_samples": _local_count(path),
    }


def _init_schema(path: Path) -> None:
    with sqlite3.connect(path) as connection:
        connection.executescript(
            """
            PRAGMA journal_mode=WAL;
            CREATE TABLE IF NOT EXISTS market_observations (
                id INTEGER PRIMARY KEY,
                symbol TEXT NOT NULL,
                bid_price REAL NOT NULL,
                bid_size REAL NOT NULL,
                ask_price REAL NOT NULL,
                ask_size REAL NOT NULL,
                mid_price REAL NOT NULL,
                microprice REAL NOT NULL,
                spread_bps REAL NOT NULL,
                imbalance REAL NOT NULL,
                last_price REAL NOT NULL,
                source_timestamp_ms INTEGER NOT NULL DEFAULT 0,
                created_at TEXT NOT NULL
            );
            CREATE INDEX IF NOT EXISTS idx_local_market_observations_symbol_id
                ON market_observations(symbol, id);
            """
        )


def _insert_rows(path: Path, rows: list[Any]) -> int:
    values = []
    for row in rows:
        if not isinstance(row, dict):
            continue
        values.append(
            (
                int(row.get("id", 0) or 0),
                str(row.get("symbol") or "").upper(),
                float(row.get("bid_price", 0.0) or 0.0),
                float(row.get("bid_size", 0.0) or 0.0),
                float(row.get("ask_price", 0.0) or 0.0),
                float(row.get("ask_size", 0.0) or 0.0),
                float(row.get("mid_price", 0.0) or 0.0),
                float(row.get("microprice", 0.0) or 0.0),
                float(row.get("spread_bps", 0.0) or 0.0),
                float(row.get("imbalance", 0.0) or 0.0),
                float(row.get("last_price", 0.0) or 0.0),
                int(row.get("source_timestamp_ms", 0) or 0),
                str(row.get("created_at") or ""),
            )
        )
    if not values:
        return 0
    with sqlite3.connect(path) as connection:
        before = connection.total_changes
        connection.executemany(
            """
            INSERT OR IGNORE INTO market_observations (
                id, symbol, bid_price, bid_size, ask_price, ask_size,
                mid_price, microprice, spread_bps, imbalance, last_price,
                source_timestamp_ms, created_at
            ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
            """,
            values,
        )
        return connection.total_changes - before


def _local_max_id(path: Path, symbol: str) -> int:
    with sqlite3.connect(path) as connection:
        row = connection.execute(
            "SELECT MAX(id) FROM market_observations WHERE symbol = ?",
            (symbol,),
        ).fetchone()
    return int(row[0] or 0) if row else 0


def _local_count(path: Path) -> int:
    with sqlite3.connect(path) as connection:
        row = connection.execute("SELECT COUNT(*) FROM market_observations").fetchone()
    return int(row[0] or 0) if row else 0


def _get_json(api_base_url: str, path: str, *, token: str, timeout: int) -> dict[str, Any]:
    headers = {"Accept": "application/json"}
    headers.update(_auth_headers(token))
    request = Request(api_base_url.rstrip("/") + path, headers=headers, method="GET")
    try:
        with urlopen(request, timeout=timeout) as response:
            text = response.read().decode("utf-8")
    except HTTPError as exc:
        detail = exc.read().decode("utf-8", errors="replace")
        raise RuntimeError(f"HTTP {exc.code} {path}: {detail[:300]}") from exc
    except URLError as exc:
        raise RuntimeError(f"network error {path}: {exc.reason}") from exc
    data = json.loads(text) if text.strip() else {}
    return data if isinstance(data, dict) else {}


def _auth_headers(token: str) -> dict[str, str]:
    value = token.strip()
    if not value:
        return {}
    headers = {"X-TradeBot-Token": value}
    if value.lower().startswith(("basic ", "bearer ")):
        headers["Authorization"] = value
    elif ":" in value:
        encoded = base64.b64encode(value.encode("utf-8")).decode("ascii")
        headers["Authorization"] = f"Basic {encoded}"
    else:
        headers["Authorization"] = f"Bearer {value}"
    return headers


def main() -> None:
    parser = argparse.ArgumentParser(description="Synchronize TradeBot L1 observations to a local SQLite cache.")
    parser.add_argument("--api-base-url", default=os.environ.get("TRADEBOT_API_BASE_URL", "https://tb.kusoft.xyz"))
    parser.add_argument("--api-auth", default=os.environ.get("TRADEBOT_API_AUTH", ""))
    parser.add_argument("--database", default="runtime/orderbook_observations.sqlite3")
    args = parser.parse_args()
    result = sync_orderbook_observations(
        api_base_url=args.api_base_url,
        token=args.api_auth,
        database_path=args.database,
    )
    print(json.dumps(result, ensure_ascii=False))


if __name__ == "__main__":
    main()