#!/usr/bin/env python3
"""baostock 15min 下载（NAS 容器跑，本机零负载）。

为何 baostock：新浪 stock_zh_a_minute 15min 历史仅半年；baostock 5.5 年 + 复权齐全
（实测浦发 2021-01-04 起 21328 行；除权日 raw -5.9% vs qfq -0.7% 平滑）。

双源：qfq(adjustflag=2) + raw(adjustflag=3)。断点续传（skip 已存在，扩范围重下）。
数据落 NAS 本地盘（容器挂载 /volume1/stock → /stock，最快，不经网络写入）。

约束：单线程 + SLEEP 限速（baostock 服务器温和限频，避免封）。

用法（NAS 容器，本机编排）：
  /var/packages/Docker/target/usr/bin/docker run --rm -d --name dl15 \\
    --memory=512m -v /volume1/stock:/stock \\
    sanguo_vnpy_v2:with-sqlite \\
    sh -c "pip install baostock -q && python /stock/sanguo_vnpy/scripts/baostock_download.py --all --adjust both --start 2021-01-01"

env:
  DL_DIR       输出根（默认 /stock/A股数据/minute_kline）
  STOCK_LIST   全市场 csv（默认 NAS stock_info csv）
  SLEEP        每只每源间隔秒（默认 0.3）
"""
import argparse
import csv
import logging
import os
import sys
import time

import pandas as pd
import baostock as bs

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
log = logging.getLogger("dl_15min")

DL_DIR = os.environ.get("DL_DIR", "/stock/A股数据/minute_kline")
SLEEP = float(os.environ.get("SLEEP", "0.3"))
DEFAULT_STOCK_LIST = "/stock/A股数据/stock_info/stock_basic_info_raw_20260326_113530.csv"


def prefix(code: str) -> str:
    return "sh" if code.startswith(("60", "68", "51", "56", "58")) else "sz"


def bs_symbol(code: str) -> str:
    return f"{prefix(code)}.{code}"


def out_path(code: str, adjust: str) -> str:
    d = os.path.join(DL_DIR, f"15min_{adjust}")
    os.makedirs(d, exist_ok=True)
    return os.path.join(d, f"{prefix(code)}{code}_15min.parquet")


def exists(code: str, adjust: str, start_year: int) -> bool:
    """断点续传：文件存在且含 start_year 数据（扩范围 start 更早 → 重下覆盖）。"""
    f = out_path(code, adjust)
    if not os.path.exists(f):
        return False
    try:
        df = pd.read_parquet(f, columns=["date"])
        return str(start_year) in df["date"].astype(str).str[:4].unique()
    except Exception:
        return False


def reconnect() -> bool:
    """baostock 断连后重连（限流/broken pipe 恢复）。"""
    try:
        bs.logout()
    except Exception:  # noqa: BLE001
        pass
    return bs.login().error_code == "0"


def download_one(code: str, start: str, end: str, adjustflag: str, max_retries: int = 3):
    """baostock 15min，返回 (df, None) 或 (None, err)。断连/限流自动 re-login retry。"""
    last_err = ""
    for attempt in range(max_retries):
        try:
            rs = bs.query_history_k_data_plus(
                bs_symbol(code),
                "date,time,open,high,low,close,volume,amount",
                start_date=start, end_date=end, frequency="15", adjustflag=adjustflag,
            )
            rows, fields = [], rs.fields
            while (rs.error_code == "0") & rs.next():
                rows.append(rs.get_row_data())
            if not rows:
                last_err = rs.error_msg or "empty"
                # 空可能是限流（"接收数据异常"）→ re-login 重试
                if attempt < max_retries - 1:
                    reconnect()
                    time.sleep(2)
                    continue
                return None, last_err
            df = pd.DataFrame(rows, columns=fields)
            for c in ("open", "high", "low", "close", "volume", "amount"):
                df[c] = pd.to_numeric(df[c], errors="coerce")
            t = df["time"].astype(str)
            df["datetime"] = pd.to_datetime(
                df["date"] + " " + t.str[8:10] + ":" + t.str[10:12], errors="coerce"
            )
            df = df.dropna(subset=["close", "datetime"]).sort_values("datetime")
            return df, None
        except Exception as e:  # noqa: BLE001  broken pipe 等
            last_err = f"{type(e).__name__}: {str(e)[:80]}"
            if attempt < max_retries - 1:
                reconnect()
                time.sleep(2)
                continue
            return None, last_err
    return None, last_err


def save_one(code, df, adjust):
    df.to_parquet(out_path(code, adjust))
    return len(df)


def load_all_codes(stock_list):
    codes = []
    with open(stock_list, encoding="utf-8", errors="replace") as f:
        for row in csv.DictReader(f):
            for k in ("code", "symbol", "ts_code", "代码", "股票代码", "成分券代码"):
                if k in row and row[k]:
                    c = str(row[k]).strip().split(".")[0]
                    if c.isdigit() and len(c) == 6:
                        codes.append(c)
                        break
    return codes


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--symbols")
    ap.add_argument("--all", action="store_true")
    ap.add_argument("--start", default="2021-01-01")
    ap.add_argument("--end", default=None)
    ap.add_argument("--adjust", default="qfq", help="qfq / raw / both")
    ap.add_argument("--force", action="store_true")
    args = ap.parse_args()

    end = args.end or time.strftime("%Y-%m-%d")
    start_year = int(args.start[:4])
    adjusts = ["qfq", "raw"] if args.adjust == "both" else [args.adjust]
    flag = {"qfq": "2", "raw": "3"}

    if args.symbols:
        codes = [c.strip() for c in args.symbols.split(",") if c.strip()]
    elif args.all:
        sl = os.environ.get("STOCK_LIST", DEFAULT_STOCK_LIST)
        if not os.path.exists(sl):
            log.error("STOCK_LIST 不存在: %s", sl)
            sys.exit(1)
        codes = load_all_codes(sl)
        log.info("--all 读到 %d 只", len(codes))
    else:
        ap.error("需 --symbols 或 --all")

    lg = bs.login()
    log.info("baostock login: %s %s", lg.error_code, lg.error_msg)

    ok = fail = rows = skipped = 0
    for i, code in enumerate(codes, 1):
        for adj in adjusts:
            if not args.force and exists(code, adj, start_year):
                skipped += 1
                continue
            df, err = download_one(code, args.start, end, flag[adj])
            if df is None or df.empty:
                fail += 1
                if i <= 20 or i % 100 == 0:
                    log.warning("[%d/%d] %s %s empty %s", i, len(codes), code, adj, err)
            else:
                try:
                    n = save_one(code, df, adj)
                    ok += 1
                    rows += n
                    if i <= 20 or i % 100 == 0:
                        log.info("[%d/%d] %s %s ok %d", i, len(codes), code, adj, n)
                except Exception as e:  # noqa: BLE001
                    fail += 1
                    log.error("save %s %s: %s", code, adj, e)
            time.sleep(SLEEP)
    bs.logout()
    log.info("=== 完成 ok=%d skip=%d fail=%d rows=%d dir=%s ===",
             ok, skipped, fail, rows, DL_DIR)


if __name__ == "__main__":
    main()