#!/usr/bin/env python3
"""从币安公开数据仓库 data.binance.vision 下载 USDT-M 永续 K 线（CSV）。

（本服务器直连 fapi.binance.com 被区域限制，公开数据仓库可用。）

用法:
    python scripts/download_data.py --pairs BTC/USDT:USDT ETH/USDT:USDT \
        --timeframe 5m --start 2026-01-01
    python scripts/download_data.py --pairs BTC/USDT:USDT --timeframe 1d --start 2025-01-01
"""
import argparse
import io
import time
import zipfile
from pathlib import Path

import pandas as pd
import requests

DATA_DIR = Path(__file__).resolve().parents[1] / "data"
BASE = "https://data.binance.vision/data/futures/um"
COLS = ["open_time", "open", "high", "low", "close", "volume", "close_time",
        "quote_volume", "count", "taker_buy_volume", "taker_buy_quote_volume", "ignore"]


def filename(pair: str, timeframe: str) -> Path:
    safe = pair.replace("/", "").replace(":", "_")
    return DATA_DIR / f"{safe}_{timeframe}.csv"


def _symbol(pair: str) -> str:
    return pair.split(":")[0].replace("/", "")


def _fetch_zip(url: str) -> pd.DataFrame | None:
    r = requests.get(url, timeout=30)
    if r.status_code != 200:
        return None
    with zipfile.ZipFile(io.BytesIO(r.content)) as zf:
        with zf.open(zf.namelist()[0]) as f:
            df = pd.read_csv(f, header=None, names=COLS, dtype={"open_time": str})
    # 新版文件含表头行（open_time 为字符串 "open_time"）
    df = df[pd.to_numeric(df["open_time"], errors="coerce").notna()]
    df["open_time"] = df["open_time"].astype("int64")
    # open_time 微秒/毫秒兼容
    if df["open_time"].iloc[0] > 10**14:
        df["open_time"] //= 1000
    out = pd.DataFrame({
        "ts": (df["open_time"] // 1000).astype(int),
        "open": df["open"].astype(float), "high": df["high"].astype(float),
        "low": df["low"].astype(float), "close": df["close"].astype(float),
        "volume": df["volume"].astype(float),
    })
    return out


def download(pair: str, timeframe: str, start: pd.Timestamp,
             end: pd.Timestamp) -> pd.DataFrame:
    sym = _symbol(pair)
    frames = []
    # 1) 月频文件（不含当月）
    for m in pd.period_range(start, end, freq="M"):
        if m == pd.Period(end, freq="M"):
            break
        url = f"{BASE}/monthly/klines/{sym}/{timeframe}/{sym}-{timeframe}-{m}.zip"
        df = _fetch_zip(url)
        if df is not None:
            frames.append(df)
            print(f"  {sym} {timeframe} {m}: {len(df)} 根")
        time.sleep(0.2)
    # 2) 日频文件（当月 + 月频缺失的零散日）
    covered = pd.concat(frames)["ts"].max() if frames else 0
    for d in pd.date_range(start, end, freq="D"):
        if d.timestamp() <= covered:
            continue
        url = f"{BASE}/daily/klines/{sym}/{timeframe}/{sym}-{timeframe}-{d:%Y-%m-%d}.zip"
        df = _fetch_zip(url)
        if df is not None:
            frames.append(df)
        time.sleep(0.1)
    if not frames:
        raise SystemExit(f"{pair} {timeframe} 无数据")
    out = (pd.concat(frames).drop_duplicates("ts").sort_values("ts")
           .query("ts >= @start.timestamp() and ts <= @end.timestamp()"))
    return out.reset_index(drop=True)


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--pairs", nargs="+", default=["BTC/USDT:USDT"])
    ap.add_argument("--timeframe", default="5m")
    ap.add_argument("--start", default="2026-01-01")
    ap.add_argument("--end", default=None, help="默认昨天（保证 K 线已收盘）")
    args = ap.parse_args()

    start = pd.Timestamp(args.start, tz="UTC")
    end = (pd.Timestamp(args.end, tz="UTC") if args.end
           else pd.Timestamp.now(tz="UTC") - pd.Timedelta(days=1))
    DATA_DIR.mkdir(exist_ok=True)
    for pair in args.pairs:
        df = download(pair, args.timeframe, start, end)
        out = filename(pair, args.timeframe)
        df.to_csv(out, index=False)
        print(f"{pair} {args.timeframe}: {len(df)} 根 "
              f"[{pd.to_datetime(df.ts.min(), unit='s')} ~ "
              f"{pd.to_datetime(df.ts.max(), unit='s')}] → {out}")


if __name__ == "__main__":
    main()