Chan/crypto_wyckoff/io.py

"""Paths + OHLCV cache + DATA_SERVICE fetch (crypto continuous calendar)."""

from __future__ import annotations

import json
import logging
import os
import sqlite3
import time
from datetime import date, datetime, timezone
from pathlib import Path
from typing import Iterable

import requests

from crypto_wyckoff.domain_models import OHLCVFrame

logger = logging.getLogger(__name__)

_REPO_ROOT = Path(__file__).resolve().parents[1]
DATA_DIR = Path(os.environ.get("CRYPTO_WYCKOFF_DATA", str(_REPO_ROOT / "data" / "crypto_wyckoff")))
BARS_DB = DATA_DIR / "bars.sqlite"
SCAN_DB = DATA_DIR / "scan.sqlite"

DATA_SERVICE_URL = os.environ.get(
    "DATA_SERVICE_URL",
    os.environ.get("DATASVC_URL", "https://provider.jackyu66.com"),
).rstrip("/")

# Continuous crypto: bar counts (not A-share weekend-padded calendar multipliers)
# Provider has 1d/1w but no 1M — monthly is resampled locally from daily UTC months.
LOOKBACK = {"1d": 250, "1w": 104, "1M": 60}
TF_PROVIDER = ("1d", "1w")
TF_LIST = ("1d", "1w", "1M")


def ensure_dirs() -> None:
    DATA_DIR.mkdir(parents=True, exist_ok=True)


def _symbol_key(symbol: str) -> str:
    return symbol.replace("/", "_").replace(":", "_")


def _bars_conn() -> sqlite3.Connection:
    ensure_dirs()
    conn = sqlite3.connect(str(BARS_DB), timeout=60)
    conn.execute(
        """
        CREATE TABLE IF NOT EXISTS bars (
            symbol TEXT NOT NULL,
            tf TEXT NOT NULL,
            ts INTEGER NOT NULL,
            open REAL, high REAL, low REAL, close REAL, volume REAL,
            PRIMARY KEY (symbol, tf, ts)
        )
        """
    )
    conn.execute("CREATE INDEX IF NOT EXISTS idx_bars_sym_tf ON bars(symbol, tf)")
    return conn


def fetch_candles(
    symbol: str,
    tf: str,
    *,
    limit: int | None = None,
    start_ms: int | None = None,
    end_ms: int | None = None,
    timeout: float = 15.0,
) -> list[dict]:
    params: dict = {"symbol": symbol, "tf": tf}
    if limit is not None:
        params["limit"] = int(limit)
    if start_ms is not None:
        params["start"] = int(start_ms)
    if end_ms is not None:
        params["end"] = int(end_ms)
    resp = requests.get(f"{DATA_SERVICE_URL}/api/candles", params=params, timeout=timeout)
    resp.raise_for_status()
    data = resp.json()
    if not isinstance(data, list):
        return []
    out = []
    for row in data:
        try:
            ts = int(float(row["timestamp"]))
            out.append(
                {
                    "ts": ts,
                    "open": float(row["open"]),
                    "high": float(row["high"]),
                    "low": float(row["low"]),
                    "close": float(row["close"]),
                    "volume": float(row.get("volume") or 0),
                }
            )
        except (KeyError, TypeError, ValueError):
            continue
    out.sort(key=lambda r: r["ts"])
    return out


def upsert_bars(symbol: str, tf: str, rows: list[dict]) -> int:
    if not rows:
        return 0
    conn = _bars_conn()
    try:
        conn.executemany(
            """
            INSERT INTO bars(symbol, tf, ts, open, high, low, close, volume)
            VALUES (?, ?, ?, ?, ?, ?, ?, ?)
            ON CONFLICT(symbol, tf, ts) DO UPDATE SET
              open=excluded.open, high=excluded.high, low=excluded.low,
              close=excluded.close, volume=excluded.volume
            """,
            [
                (symbol, tf, r["ts"], r["open"], r["high"], r["low"], r["close"], r["volume"])
                for r in rows
            ],
        )
        conn.commit()
        return len(rows)
    finally:
        conn.close()


def load_frame(symbol: str, tf: str, lookback: int | None = None) -> OHLCVFrame | None:
    lookback = lookback or LOOKBACK.get(tf, 100)
    conn = _bars_conn()
    try:
        cur = conn.execute(
            """
            SELECT ts, open, high, low, close, volume FROM bars
            WHERE symbol=? AND tf=?
            ORDER BY ts DESC LIMIT ?
            """,
            (symbol, tf, lookback),
        )
        rows = list(reversed(cur.fetchall()))
    finally:
        conn.close()
    if not rows:
        return None
    trade_dates: list[date] = []
    for ts, *_ in rows:
        trade_dates.append(datetime.fromtimestamp(ts / 1000.0, tz=timezone.utc).date())
    return OHLCVFrame(
        ts_code=symbol,
        timeframe=tf,
        trade_dates=trade_dates,
        open=[r[1] for r in rows],
        high=[r[2] for r in rows],
        low=[r[3] for r in rows],
        close=[r[4] for r in rows],
        volume=[r[5] for r in rows],
    )


def bar_count(symbol: str, tf: str) -> int:
    conn = _bars_conn()
    try:
        cur = conn.execute(
            "SELECT COUNT(*) FROM bars WHERE symbol=? AND tf=?", (symbol, tf)
        )
        return int(cur.fetchone()[0])
    finally:
        conn.close()


def rebuild_monthly_from_daily(symbol: str) -> int:
    """Aggregate UTC calendar-month OHLCV from local daily bars (provider has no 1M)."""
    conn = _bars_conn()
    try:
        cur = conn.execute(
            """
            SELECT ts, open, high, low, close, volume FROM bars
            WHERE symbol=? AND tf='1d' ORDER BY ts ASC
            """,
            (symbol,),
        )
        daily = cur.fetchall()
    finally:
        conn.close()
    if not daily:
        return 0

    months: dict[tuple[int, int], dict] = {}
    for ts, o, h, l, c, v in daily:
        dt = datetime.fromtimestamp(ts / 1000.0, tz=timezone.utc)
        key = (dt.year, dt.month)
        # month bar open timestamp = first day 00:00 UTC
        month_ts = int(datetime(dt.year, dt.month, 1, tzinfo=timezone.utc).timestamp() * 1000)
        if key not in months:
            months[key] = {
                "ts": month_ts,
                "open": o,
                "high": h,
                "low": l,
                "close": c,
                "volume": v or 0.0,
            }
        else:
            m = months[key]
            m["high"] = max(m["high"], h)
            m["low"] = min(m["low"], l)
            m["close"] = c
            m["volume"] = (m["volume"] or 0) + (v or 0)

    rows = sorted(months.values(), key=lambda r: r["ts"])
    # drop stale months then upsert
    conn = _bars_conn()
    try:
        conn.execute("DELETE FROM bars WHERE symbol=? AND tf='1M'", (symbol,))
        conn.commit()
    finally:
        conn.close()
    return upsert_bars(symbol, "1M", rows)


def backfill_symbol(symbol: str, tfs: Iterable[str] = TF_LIST) -> dict:
    """Pull history for continuous crypto TFs; monthly derived from daily."""
    stats = {}
    for tf in TF_PROVIDER:
        if tf not in tfs and "1M" not in tfs:
            continue
        need = LOOKBACK.get(tf, 100)
        # need extra daily for monthly history
        if tf == "1d":
            need = max(need, LOOKBACK["1M"] * 31)
        try:
            rows = fetch_candles(symbol, tf, limit=need)
            n = upsert_bars(symbol, tf, rows)
            stats[tf] = n
        except Exception as e:
            logger.warning("backfill %s %s failed: %s", symbol, tf, e)
            stats[tf] = 0
        time.sleep(0.05)
    if "1M" in tfs or True:
        try:
            stats["1M"] = rebuild_monthly_from_daily(symbol)
        except Exception as e:
            logger.warning("monthly rebuild %s failed: %s", symbol, e)
            stats["1M"] = 0
    return stats


def tip_update_symbol(symbol: str, tfs: Iterable[str] = TF_LIST) -> bool:
    """Update forming tip bars (limit=3). Returns True if any bar changed."""
    changed = False
    for tf in TF_PROVIDER:
        try:
            rows = fetch_candles(symbol, tf, limit=3)
            if not rows:
                continue
            before = _tip_fingerprint(symbol, tf)
            upsert_bars(symbol, tf, rows)
            after = _tip_fingerprint(symbol, tf)
            if before != after:
                changed = True
        except Exception as e:
            logger.debug("tip %s %s: %s", symbol, tf, e)
        time.sleep(0.02)
    # Always rebuild current month tip from daily
    before_m = _tip_fingerprint(symbol, "1M")
    try:
        rebuild_monthly_from_daily(symbol)
    except Exception as e:
        logger.debug("monthly tip %s: %s", symbol, e)
    after_m = _tip_fingerprint(symbol, "1M")
    if before_m != after_m:
        changed = True
    return changed


def _tip_fingerprint(symbol: str, tf: str) -> tuple | None:
    conn = _bars_conn()
    try:
        cur = conn.execute(
            """
            SELECT ts, open, high, low, close, volume FROM bars
            WHERE symbol=? AND tf=? ORDER BY ts DESC LIMIT 1
            """,
            (symbol, tf),
        )
        row = cur.fetchone()
        return tuple(row) if row else None
    finally:
        conn.close()


def fetch_symbols_from_provider() -> list[str]:
    try:
        resp = requests.get(f"{DATA_SERVICE_URL}/health", timeout=8)
        resp.raise_for_status()
        payload = resp.json()
        symbols = payload.get("symbols") or payload.get("symbol_list") or []
        return [s for s in symbols if isinstance(s, str)]
    except Exception as e:
        logger.warning("health symbols failed: %s", e)
        return []