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 many TFs; 1M is resampled locally from daily UTC months.
LOOKBACK = {
    "1h": 500,
    "2h": 400,
    "4h": 300,
    "6h": 280,
    "8h": 250,
    "12h": 220,
    "1d": 250,
    "1w": 104,
    "1M": 60,
}
# Default D/W/M stack (kept for compat); combos may request more TFs from provider.
TF_PROVIDER = ("1h", "4h", "8h", "1d", "1w")
TF_LIST = ("1d", "1w", "1M")
LOCAL_ONLY_TFS = frozenset({"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 is_intraday_tf(tf: str) -> bool:
    """True for minute/hour TFs that need clock time on charts."""
    t = (tf or "").strip()
    return t.endswith("m") or t.endswith("h")


def load_bars_with_ts(
    symbol: str, tf: str, lookback: int | None = None
) -> list[dict]:
    """Return OHLCV rows with UTC ms ts (for chart labels).

    ``datetime`` is wall-clock in Asia/Shanghai (UTC+8) for display.
    """
    from zoneinfo import ZoneInfo

    tz_cn = ZoneInfo("Asia/Shanghai")
    if lookback is None:
        try:
            from crypto_wyckoff.combos import lookback_for

            lookback = lookback_for(tf)
        except Exception:
            lookback = LOOKBACK.get(tf, 100)
    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()
    out = []
    for ts, o, h, l, c, v in rows:
        dt_utc = datetime.fromtimestamp(ts / 1000.0, tz=timezone.utc)
        dt_cn = dt_utc.astimezone(tz_cn)
        out.append(
            {
                "ts": int(ts),
                "datetime": dt_cn.strftime("%Y-%m-%dT%H:%M:%S+08:00"),
                "date": dt_cn.strftime("%Y-%m-%d"),
                "open": o,
                "high": h,
                "low": l,
                "close": c,
                "volume": v,
            }
        )
    return out


def load_frame(symbol: str, tf: str, lookback: int | None = None) -> OHLCVFrame | None:
    rows = load_bars_with_ts(symbol, tf, lookback)
    if not rows:
        return None
    return OHLCVFrame(
        ts_code=symbol,
        timeframe=tf,
        trade_dates=[
            datetime.fromtimestamp(r["ts"] / 1000.0, tz=timezone.utc).date() for r in rows
        ],
        open=[r["open"] for r in rows],
        high=[r["high"] for r in rows],
        low=[r["low"] for r in rows],
        close=[r["close"] for r in rows],
        volume=[r["volume"] 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 requested TFs; monthly derived from daily when needed."""
    wanted = list(dict.fromkeys(tfs))
    stats: dict = {}
    need_monthly = "1M" in wanted
    if need_monthly and "1d" not in wanted:
        wanted = ["1d", *wanted]

    for tf in wanted:
        if tf in LOCAL_ONLY_TFS:
            continue
        need = LOOKBACK.get(tf, 100)
        if tf == "1d" and need_monthly:
            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 need_monthly:
        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."""
    wanted = list(dict.fromkeys(tfs))
    changed = False
    for tf in wanted:
        if tf in LOCAL_ONLY_TFS:
            continue
        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)
    if "1M" in wanted:
        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 []