#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""BaoStock 静态数据下载 (basic / adjust_factor / dividend 3 类, 只产 parquet)

目标:
  在 Windows Server 2022 VPS (49.232.102.198) 上用 baostock 全量拉取 A 股
  (含退市) 3 类静态数据:
    1. basic         - 基础信息 (全市场单文件, 不按只)
    2. adjust_factor - 复权因子 (per-stock)
    3. dividend      - 分红送转 (per-stock, 循环年份 2020..2026)

  输出 import-ready parquet, 后续数据平台直接读. 与 15min 下载脚本共用硬化机制.

与 baostock_15min_fullmarket_download.py 共用的硬约束 (踩过的坑):
  1. 单进程单登录, 严禁并发 (baostock 并发会拉黑封 IP 6-24h)
  2. 直连不走代理 (脚本开头 unset proxy)
  3. 每次 baostock query 后 sleep BS_INTERVAL=0.4s 限速
  4. baostock login 受 socket.setdefaulttimeout(30) 保护
     (Windows 无 SIGALRM, 这是跨平台等价方案)
  5. empty vs failed 区分:
     - 空数据 (退市/无除权事件) → fetch 返空 df → status='empty' 中性, 不重试
     - 真错误 → fetch raise → retry → status='failed' 计断路器

与 15min 脚本差异:
  - 3 类 fetch 各自实现 (basic 单调用 / adjust_factor per-stock / dividend per-year)
  - per-stock empty **也写 marker** (静态空 = 确定性"无除权事件",
    重跑只是浪费 baostock 配额, 与 K 线 empty 不写 marker 的语义不同)
  - 输出 3 个子目录: basic / adjust_factor / dividend
  - 单次进程可顺序跑多类 (--types basic,adjust_factor,dividend)
  - basic 单文件全量刷新, 不受 --limit / --codes 影响

⚠️ 构建期禁止任何 live baostock 调用 (VPS 15min 下载并行会封 IP).
   本脚本只 py_compile + 代码 review, smoke 测试延后.

用法:
  # 全跑 (basic + adjust_factor + dividend)
  python baostock_static_download.py
  # 只跑指定类型 (逗号分隔)
  python baostock_static_download.py --types basic,adjust_factor
  python baostock_static_download.py --types dividend
  # 测试 (前 10 只, basic 不受影响)
  python baostock_static_download.py --types adjust_factor --limit 10
  # 指定股票 (6 位 code, 不带前缀)
  python baostock_static_download.py --types dividend --codes 600519,000001
  # 强制重下, 忽略 marker
  python baostock_static_download.py --types basic --force

输出目录结构:
  {OUT_DIR}/
  ├── basic/
  │   ├── stock_basic.parquet
  │   └── .stock_basic.baostock        (marker)
  ├── adjust_factor/
  │   ├── 600519.SH_factor.parquet
  │   ├── .600519.SH_factor.baostock    (marker)
  │   └── ...
  ├── dividend/
  │   ├── 600519.SH_dividend.parquet
  │   ├── .600519.SH_dividend.baostock  (marker)
  │   └── ...
  └── logs/baostock_static_YYYYMMDD_HHMMSS.log

退出码: 0=完成, 1=致命错误 (登录/列表拉取失败), 2=断路器触发 (可重试)
"""

import argparse
import datetime
import json
import logging
import os
import socket
import sys
import time
from functools import partial
from pathlib import Path
from typing import Callable, List, Optional, Tuple

# ======================== 硬约束: unset proxy + socket timeout ========================
# 必须在 import baostock 之前清理 (baostock 底层 urllib 会读 proxy 环境变量)
for _k in ("http_proxy", "https_proxy", "HTTP_PROXY", "HTTPS_PROXY"):
    os.environ.pop(_k, None)

# SIGALRM 跨平台替代: 全局 socket 超时 30s (Windows 无 SIGALRM)。
# baostock 的 socket 连接会继承这个超时, connect 卡死时抛 socket.timeout。
socket.setdefaulttimeout(30)

# stdout 行缓冲 (SSH-detached / pipe 重定向时也能看到实时进度)
try:
    sys.stdout.reconfigure(line_buffering=True)
except (AttributeError, ValueError):
    pass

import baostock as bs  # noqa: E402
import pandas as pd  # noqa: E402


# ======================== 配置 (环境变量可覆盖) ========================

DEFAULT_OUT_DIR = r"C:\sanguo_vnpy_v2\data\static"
OUT_DIR = Path(os.environ.get("BS_STATIC_OUT_DIR", DEFAULT_OUT_DIR))
DEFAULT_LOG_DIR = OUT_DIR / "logs"
_log_dir_env = os.environ.get("BS_STATIC_LOG_DIR")
LOG_DIR = Path(_log_dir_env) if _log_dir_env else DEFAULT_LOG_DIR

START_DATE_DEFAULT = "2020-01-01"
# 分红查询年份范围 (含两端). 2020..2026 = 7 年, 与 START_DATE_DEFAULT 对齐。
DIVIDEND_YEAR_RANGE = (2020, 2026)

BS_INTERVAL = 0.4             # 每次 baostock query 后间隔秒 (防封 IP)
BS_MAX_RETRIES = 3            # 单次下载重试次数 (含强制重登)
PROGRESS_LOG_EVERY = 500      # 每 N 只打一次进度日志
RELOGIN_EVERY = 400           # 每 N 只定期重登保持连接
CIRCUIT_BREAKER = 30          # 连续失败 N 只 → 断路退出

# 子目录名 (同时也是 --types 合法值)
DIR_BASIC = "basic"
DIR_ADJUST = "adjust_factor"
DIR_DIVIDEND = "dividend"

# 文件名后缀 (parquet / marker stem 解析用)
SUFFIX_BASIC = "stock_basic"             # basic 单文件 stem (无 per-stock)
SUFFIX_ADJUST = "_factor"                # 600519.SH_factor
SUFFIX_DIVIDEND = "_dividend"            # 600519.SH_dividend

VALID_TYPES = (DIR_BASIC, DIR_ADJUST, DIR_DIVIDEND)


# ======================== 日志 ========================

def setup_logging() -> Tuple[logging.Logger, Path]:
    """配置 root logger: 同时写 stdout + 文件。返回 (logger, log_file)。"""
    LOG_DIR.mkdir(parents=True, exist_ok=True)
    ts = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    log_file = LOG_DIR / f"baostock_static_{ts}.log"

    fmt = logging.Formatter("%(asctime)s %(levelname)s %(message)s")
    root = logging.getLogger()
    root.setLevel(logging.INFO)
    for h in list(root.handlers):
        root.removeHandler(h)

    sh = logging.StreamHandler(sys.stdout)
    sh.setFormatter(fmt)
    root.addHandler(sh)

    fh = logging.FileHandler(log_file, encoding="utf-8")
    fh.setFormatter(fmt)
    root.addHandler(fh)

    return logging.getLogger(__name__), log_file


logger, LOG_FILE = setup_logging()


# ======================== 工具函数 ========================

def normalize_date(s: str) -> str:
    """接受 YYYYMMDD 或 YYYY-MM-DD, 返回 YYYY-MM-DD; 空串 → 今天。"""
    s = (s or "").strip()
    if not s:
        return datetime.date.today().strftime("%Y-%m-%d")
    if "-" in s:
        return s
    if len(s) == 8 and s.isdigit():
        return f"{s[:4]}-{s[4:6]}-{s[6:8]}"
    raise ValueError(f"无效日期格式: {s} (期望 YYYYMMDD 或 YYYY-MM-DD)")


def baostock_code_to_parts(bs_code: str) -> Tuple[str, str]:
    """sh.600519 → ('600519', 'SH'); sz.000001 → ('000001', 'SZ')。
    前缀直接从 baostock code 拿, 不猜。
    """
    if "." not in bs_code:
        raise ValueError(f"无效 baostock code: {bs_code}")
    prefix, num = bs_code.split(".", 1)
    if prefix == "sh":
        return num, "SH"
    if prefix == "sz":
        return num, "SZ"
    raise ValueError(f"未知 baostock 前缀: {bs_code}")


def parts_to_baostock(code: str, exchange: str) -> str:
    """('600519', 'SH') → 'sh.600519'。"""
    pfx = "sh" if exchange == "SH" else "sz"
    return f"{pfx}.{code}"


def guess_exchange_by_code(code: str) -> str:
    """6/68/51 开头 → SH, 其他 → SZ (无前缀代码用, 与 15min 模板一致)。"""
    if code.startswith(("6", "68", "51")):
        return "SH"
    return "SZ"


def _rs_to_rows(rs) -> Tuple[List[str], List[List[str]]]:
    """读取 ResultData, 返回 (fields, rows)。
    用 rs.fields 动态拿字段名, 不硬编码位置 (踩过坑: baostock 字段顺序与文档不一致)。
    """
    fields = list(rs.fields)
    rows: List[List[str]] = []
    while rs.next():
        rows.append(rs.get_row_data())
    return fields, rows


# ======================== 股票列表 (baostock query_stock_basic) ========================

def fetch_all_stocks() -> List[Tuple[str, str]]:
    """baostock query_stock_basic 拉全市场 A 股列表 (含退市, type=='1')。

    返回: [(code, exchange), ...], 例 [('600519', 'SH'), ('000001', 'SZ')]
    过滤: type=='1' (股票); **不过滤 status** (保留退市, 避免生存偏差)。

    baostock 0.9.3 实测字段顺序与文档不一致, 用 rs.fields 动态建索引按名取,
    避免字段位置差异导致过滤错位 (照抄 15min 模板已验证逻辑)。
    """
    rs = bs.query_stock_basic()
    if rs.error_code != "0":
        raise RuntimeError(
            f"query_stock_basic 失败: code={rs.error_code} msg={rs.error_msg}"
        )

    fields = list(rs.fields)
    idx = {name: i for i, name in enumerate(fields)}
    logger.info("query_stock_basic fields=%s", fields)

    out: List[Tuple[str, str]] = []
    n_stock = 0
    n_skip_type = 0
    n_skip_code = 0
    while rs.next():
        r = rs.get_row_data()
        type_ = r[idx["type"]] if "type" in idx and len(r) > idx["type"] else ""
        if type_ != "1":  # 只要股票 (排除指数/债券/其他)
            n_skip_type += 1
            continue
        bs_code = r[idx["code"]] if "code" in idx and len(r) > idx["code"] else ""
        try:
            code, exchange = baostock_code_to_parts(bs_code)
        except ValueError:
            n_skip_code += 1
            continue
        if len(code) != 6 or not code.isdigit():
            n_skip_code += 1
            continue
        out.append((code, exchange))
        n_stock += 1
    logger.info(
        "股票列表: stocks(type=1)=%d skip_type=%d skip_code=%d (含退市, 不过滤 status)",
        n_stock, n_skip_type, n_skip_code)
    return out


# ======================== baostock 登录 ========================

def _login_once() -> bool:
    """bs.login() — 已被 socket.setdefaulttimeout(30) 保护。
    connect 卡死会抛 socket.timeout, 这里捕获。
    """
    try:
        lg = bs.login()
        if lg.error_code == "0":
            return True
        logger.error("bs.login() 失败: code=%s msg=%s", lg.error_code, lg.error_msg)
        return False
    except (socket.timeout, TimeoutError) as e:
        logger.error("bs.login() 超时: %s (baostock 疑似冷却)", e)
        return False
    except Exception as e:
        logger.error("bs.login() 异常: %s", e)
        return False


def _relogin() -> bool:
    """强制重登: logout + login。失败等 2s 再试 1 次。"""
    try:
        bs.logout()
    except Exception:
        pass
    if _login_once():
        return True
    time.sleep(2)
    try:
        bs.logout()
    except Exception:
        pass
    return _login_once()


# ======================== 数据 fetch (3 类各自实现) ========================

def fetch_basic() -> pd.DataFrame:
    """拉全市场基础信息, 单次调用返 df。

    返回列: baostock 原始字段 (code, code_name, ipoDate, outDate, type, status) +
            派生列 code (6位纯数字) / exchange (SH/SZ) / baostock_code (sh.XXXXXX)
    原始 code 列被覆盖为 6 位纯数字, baostock_code 保留 sh.XXXXXX 形式。

    baostock 字段顺序与文档不一致时, 用 rs.fields 动态拿列名, 不硬编码位置。
    """
    rs = bs.query_stock_basic()
    if rs.error_code != "0":
        raise RuntimeError(
            f"query_stock_basic 失败: code={rs.error_code} msg={rs.error_msg}"
        )

    fields, rows = _rs_to_rows(rs)
    extra_cols = ["code_6digit", "exchange", "baostock_code"]
    if not rows:
        return pd.DataFrame(columns=fields + extra_cols)

    df = pd.DataFrame(rows, columns=fields)
    # 派生列: 从原始 baostock code 拆出 6 位 + 交易所
    df["baostock_code"] = df["code"]
    parts = df["code"].apply(baostock_code_to_parts)
    df["code_6digit"] = parts.str[0]
    df["exchange"] = parts.str[1]
    # 列重排: 把派生 6 位 code 放前面 (与原始 baostock code 区分, 避免歧义)
    return df.reset_index(drop=True)


def fetch_adjust_factor(bs_code: str, start_date: str, end_date: str) -> pd.DataFrame:
    """拉单只复权因子, 返回 df (可能空)。

    返回字段 (baostock 文档): code, dividOperateDate, foreAdjustFactor,
                              backAdjustFactor, adjustFactor
    数值列 to_numeric; 空数据 (无除权事件的票) → empty df, 上层计 empty 中性。

    字段顺序按 rs.fields 动态拿, 不硬编码位置 (容错 baostock 字段顺序差异)。
    """
    rs = bs.query_adjust_factor(
        code=bs_code, start_date=start_date, end_date=end_date,
    )
    if rs.error_code != "0":
        raise RuntimeError(
            f"query_adjust_factor {bs_code} 错误: {rs.error_code} {rs.error_msg}"
        )

    fields, rows = _rs_to_rows(rs)
    if not rows:
        return pd.DataFrame(columns=fields)

    df = pd.DataFrame(rows, columns=fields)
    # 数值列 to_numeric (baostock 返字符串)
    for col in ("foreAdjustFactor", "backAdjustFactor", "adjustFactor"):
        if col in df.columns:
            df[col] = pd.to_numeric(df[col], errors="coerce")
    return df


def fetch_dividend(bs_code: str, year_range: Tuple[int, int]) -> pd.DataFrame:
    """拉单只分红送转, 循环年份 concat。返回 df (可能空)。

    每年调用 bs.query_dividend_data(code, year, yearType="report"),
    返回该票当年分红记录 (多行)。concat 所有年份。
    **保留 baostock 返回的全部列** (用 rs.fields 动态拿, 不硬编码字段名)。

    限速: 每次年份 query 后 sleep BS_INTERVAL (除最后一次, 由主循环负责)。
    yearType="report" = 预案公告年份 (与用户指定一致)。
    """
    years = list(range(year_range[0], year_range[1] + 1))
    frames: List[pd.DataFrame] = []
    fields_ref: Optional[List[str]] = None

    for i, year in enumerate(years):
        rs = bs.query_dividend_data(
            code=bs_code, year=year, yearType="report",
        )
        if rs.error_code != "0":
            raise RuntimeError(
                f"query_dividend_data {bs_code} year={year} 错误: "
                f"{rs.error_code} {rs.error_msg}"
            )
        fields, rows = _rs_to_rows(rs)
        if fields_ref is None:
            fields_ref = fields  # 锁定首次响应的字段列表
        if rows:
            # 各年字段应一致; 以当前响应字段为准构造 df
            frames.append(pd.DataFrame(rows, columns=fields))
        # 限速: 非最后一年时 sleep (最后一次由主循环负责 stock 间隔)
        if i < len(years) - 1:
            time.sleep(BS_INTERVAL)

    if not frames:
        cols = fields_ref if fields_ref is not None else []
        return pd.DataFrame(columns=cols)
    return pd.concat(frames, ignore_index=True)


# ======================== 路径 / marker ========================

def subdir_for(data_type: str) -> Path:
    """data_type → OUT_DIR / <subdir>。"""
    return OUT_DIR / data_type


def parquet_path_basic() -> Path:
    """basic 单文件 parquet 路径。"""
    return subdir_for(DIR_BASIC) / f"{SUFFIX_BASIC}.parquet"


def marker_path_basic() -> Path:
    """basic 单文件 marker。"""
    return subdir_for(DIR_BASIC) / f".{SUFFIX_BASIC}.baostock"


def parquet_path_per_stock(
    code: str, exchange: str, data_type: str, suffix: str,
) -> Path:
    """('600519', 'SH', 'adjust_factor', '_factor')
    → OUT_DIR/adjust_factor/600519.SH_factor.parquet
    """
    return subdir_for(data_type) / f"{code}.{exchange}{suffix}.parquet"


def marker_path_for(parquet_path: Path) -> Path:
    """parquet → 同目录 .{stem}.baostock marker。"""
    return parquet_path.parent / f".{parquet_path.stem}.baostock"


def load_done_set_per_stock(data_type: str, suffix: str) -> set:
    """扫子目录 marker 构造已完成 (code, exchange) 集合 (真相源)。

    marker 文件名格式: .600519.SH_factor.baostock
    解析 stem = '600519.SH_factor', 按 suffix='_factor' 剥出 head='600519.SH',
    再拆 code='600519' / exc='SH', 校验 6 位数字 + 交易所。
    """
    done: set = set()
    d = subdir_for(data_type)
    if not d.exists():
        return done
    ext = ".baostock"
    for marker in d.glob(f".*{ext}"):
        name = marker.name
        if not name.startswith(".") or not name.endswith(ext):
            continue
        stem = name[1:-len(ext)]  # e.g. '600519.SH_factor'
        if not suffix or not stem.endswith(suffix):
            continue
        head = stem[:-len(suffix)] if suffix else stem
        # head: '600519.SH'
        try:
            code, exc = head.split(".", 1)
            if len(code) == 6 and code.isdigit() and exc in ("SH", "SZ"):
                done.add((code, exc))
        except ValueError:
            continue
    return done


# ======================== 单只下载 ========================

def download_basic(force: bool) -> Tuple[str, int]:
    """basic 单文件全量刷新。返 (status, rows), status ∈ {'ok','skipped','failed'}。

    basic 全市场不应为空 (5537 票级别), 空数据 = 失败 (与 per-stock 语义不同)。
    """
    parquet_path = parquet_path_basic()
    marker_path = marker_path_basic()

    if not force and marker_path.exists():
        return "skipped", 0

    df: Optional[pd.DataFrame] = None
    for attempt in range(BS_MAX_RETRIES):
        try:
            df = fetch_basic()
            break  # 成功 (df 空仍 break, 下方判 failed)
        except (socket.timeout, TimeoutError, OSError) as e:
            logger.warning("basic socket 异常重试 %d/%d: %s — 强制重登",
                           attempt + 1, BS_MAX_RETRIES, e)
            if not _relogin():
                logger.error("重登失败, 放弃 basic")
                df = None
                break
        except Exception as e:
            logger.warning("basic 异常重试 %d/%d: %s — 强制重登",
                           attempt + 1, BS_MAX_RETRIES, e)
            if not _relogin():
                logger.error("重登失败, 放弃 basic")
                df = None
                break

    if df is None:
        return "failed", 0
    if df.empty:
        # 全市场不应为空 — 视为失败 (重试或人工排查)
        logger.error("basic 返回空 (异常, 全市场 type=1 不应为空)")
        return "failed", 0

    try:
        d = subdir_for(DIR_BASIC)
        d.mkdir(parents=True, exist_ok=True)
        df.to_parquet(parquet_path, index=False)
        marker_path.write_text(datetime.datetime.now().isoformat())
        return "ok", len(df)
    except Exception as e:
        logger.error("写入 %s 失败: %s", parquet_path, e)
        return "failed", 0


def download_one_per_stock(
    code: str,
    exchange: str,
    data_type: str,
    suffix: str,
    fetch_fn: Callable[[str], pd.DataFrame],
    force: bool,
) -> Tuple[str, int]:
    """通用 per-stock 下载 → 写 parquet + marker。

    fetch_fn(bs_code) → df (可能空) 或 raise。返 (status, rows), status ∈
    {'ok', 'skipped', 'empty', 'failed'}。

    与 15min 模板差异: per-stock empty **也写 marker** (静态空 = "查过了确实无除权事件",
    重跑只是浪费 baostock 配额; 与 K 线 empty "可能是延迟" 语义不同)。
    """
    parquet_path = parquet_path_per_stock(code, exchange, data_type, suffix)
    marker_path = marker_path_for(parquet_path)

    if not force and marker_path.exists():
        return "skipped", 0

    bs_code = parts_to_baostock(code, exchange)

    # retry 循环: fetch_fn 仅在真错误时抛 (走 retry/relogin);
    # 空数据返空 df (合法, 不重试), 成功 break。
    df: Optional[pd.DataFrame] = None
    for attempt in range(BS_MAX_RETRIES):
        try:
            df = fetch_fn(bs_code)
            break  # fetch 成功 (df 可能空 = 合法无数据)
        except (socket.timeout, TimeoutError, OSError) as e:
            logger.warning(
                "%s %s socket 异常重试 %d/%d: %s — 强制重登",
                code, data_type, attempt + 1, BS_MAX_RETRIES, e,
            )
            if not _relogin():
                logger.error("重登失败, 放弃 %s %s", code, data_type)
                df = None
                break
        except Exception as e:
            logger.warning(
                "%s %s 下载异常重试 %d/%d: %s — 强制重登",
                code, data_type, attempt + 1, BS_MAX_RETRIES, e,
            )
            if not _relogin():
                logger.error("重登失败, 放弃 %s %s", code, data_type)
                df = None
                break

    if df is None:
        return "failed", 0   # 多次重试仍报错 (真失败, 计断路器)

    # df 可能空 (合法) 或非空 — 两种都写 marker (静态数据语义, 见 docstring)
    try:
        d = subdir_for(data_type)
        d.mkdir(parents=True, exist_ok=True)
        df.to_parquet(parquet_path, index=False)
        marker_path.write_text(datetime.datetime.now().isoformat())
        return ("ok" if not df.empty else "empty"), len(df)
    except Exception as e:
        logger.error("写入 %s 失败: %s", parquet_path, e)
        return "failed", 0


# ======================== per-stock 主循环 (通用) ========================

def run_per_stock_type(
    data_type: str,
    suffix: str,
    fetch_fn: Callable[[str], pd.DataFrame],
    all_codes: List[Tuple[str, str]],
    args: argparse.Namespace,
    skip_loop_sleep: bool,
) -> Tuple[dict, bool]:
    """运行 per-stock 类型 (adjust_factor / dividend)。

    fetch_fn: 已绑定参数的 callable, 接受 bs_code, 返回 df (可能空) 或 raise。
    skip_loop_sleep: True 时主循环不再 sleep (fetch_fn 内部已自限速, 如 dividend)。
    返回 (stats, circuit_triggered)。
    """
    # --codes 过滤 (无前缀, 按代码开头猜交易所, 与 15min 模板一致)
    if args.codes:
        todo_codes = [
            (c.strip(), guess_exchange_by_code(c.strip()))
            for c in args.codes.split(",") if c.strip()
        ]
    else:
        todo_codes = list(all_codes)

    # --limit 截断
    if args.limit > 0:
        todo_codes = todo_codes[:args.limit]
        logger.info("[%s] limit=%d 截断", data_type, args.limit)

    # marker 断点续传
    done_set = load_done_set_per_stock(data_type, suffix)
    if args.force:
        todo = todo_codes
    else:
        todo = [(c, e) for c, e in todo_codes if (c, e) not in done_set]
    logger.info("[%s] 待处理: %d (已完成 %d)", data_type, len(todo), len(done_set))

    stats = {"ok": 0, "skipped": 0, "empty": 0, "failed": 0, "rows": 0}
    consec_fail = 0
    circuit_triggered = False
    t_start = time.time()

    for i, (code, exc) in enumerate(todo):
        # 定期重登保持连接
        if i > 0 and i % RELOGIN_EVERY == 0:
            logger.info("[%s] 定期重登 @ %d/%d", data_type, i, len(todo))
            if not _relogin():
                logger.warning("[%s] 定期重登失败, 继续 (单次失败不致命)", data_type)

        try:
            status, rows = download_one_per_stock(
                code, exc, data_type, suffix, fetch_fn, args.force,
            )
        except Exception as e:
            status, rows = "failed", 0
            logger.debug("[%s] %s.%s 异常: %s", data_type, code, exc, e)

        stats[status] = stats.get(status, 0) + 1
        if status == "ok":
            stats["rows"] += rows
            consec_fail = 0
        elif status == "failed":
            consec_fail += 1
        # skipped / empty 中性: 不重置也不递增

        # 断路器: 连续 N 只全 failed → baostock 疑似不可达, 保存进度主动退出
        if consec_fail >= CIRCUIT_BREAKER:
            logger.error(
                "[%s] [FATAL] 断路器触发: 连续 %d 只失败, baostock 疑似不可达, "
                "退出 (done_set 不含 failed 票, 复跑会重试)",
                data_type, consec_fail,
            )
            circuit_triggered = True
            break

        # 进度日志
        if (i + 1) % PROGRESS_LOG_EVERY == 0:
            elapsed = time.time() - t_start
            logger.info(
                "[%s] 进度 %d/%d ok=%d skipped=%d empty=%d failed=%d rows=%d (%.0f秒)",
                data_type, i + 1, len(todo), stats["ok"], stats["skipped"],
                stats["empty"], stats["failed"], stats["rows"], elapsed,
            )

        # 限速 (dividend 内部已 sleep, 跳过)
        if not skip_loop_sleep and i < len(todo) - 1:
            time.sleep(BS_INTERVAL)

    elapsed = time.time() - t_start
    logger.info(
        "[%s] 完成, 耗时 %.1f 秒, 统计: %s",
        data_type, elapsed, json.dumps(stats, ensure_ascii=False),
    )
    return stats, circuit_triggered


# ======================== CLI / main ========================

def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(
        description="BaoStock 静态数据下载 (basic / adjust_factor / dividend)",
    )
    p.add_argument(
        "--types", default=",".join(VALID_TYPES),
        help=f"逗号分隔类型, 默认全部 ({','.join(VALID_TYPES)})",
    )
    p.add_argument(
        "--start", default=START_DATE_DEFAULT,
        help=f"起始日期 YYYYMMDD 或 YYYY-MM-DD, 默认 {START_DATE_DEFAULT}",
    )
    p.add_argument("--end", default="", help="结束日期, 默认今天")
    p.add_argument(
        "--codes",
        help="指定代码逗号分隔 (6 位无前缀), 如 600519,000001 (basic 不受影响)",
    )
    p.add_argument(
        "--limit", type=int, default=0,
        help="限制处理股票数, 测试用 (basic 不受影响)",
    )
    p.add_argument("--force", action="store_true", help="强制重下, 忽略 marker")
    return p.parse_args()


def resolve_types(s: str) -> List[str]:
    """解析 --types 字符串 → 去重保序的类型列表。无效类型 SystemExit。"""
    parts = [t.strip() for t in s.split(",") if t.strip()]
    bad = [t for t in parts if t not in VALID_TYPES]
    if bad:
        raise SystemExit(f"未知 --types: {bad}, 可选 {list(VALID_TYPES)}")
    if not parts:
        return list(VALID_TYPES)
    seen: set = set()
    out: List[str] = []
    for t in parts:
        if t not in seen:
            out.append(t)
            seen.add(t)
    return out


def main() -> None:
    args = parse_args()
    start_date = normalize_date(args.start)
    end_date = normalize_date(args.end)
    types = resolve_types(args.types)

    logger.info("=" * 60)
    logger.info("BaoStock 静态数据下载 (basic / adjust_factor / dividend)")
    logger.info("  输出目录: %s", OUT_DIR)
    logger.info("  日志文件: %s", LOG_FILE)
    logger.info("  日期范围: %s ~ %s", start_date, end_date)
    logger.info("  分红年份: %d..%d", *DIVIDEND_YEAR_RANGE)
    logger.info("  类型: %s", types)
    logger.info("  socket.setdefaulttimeout(30) (SIGALRM 替代)")
    logger.info("  当前时间: %s", datetime.datetime.now().isoformat())
    logger.info("=" * 60)

    # 登录
    if not _login_once():
        logger.error("[FATAL] baostock 登录失败 (疑似冷却), 退出")
        sys.exit(1)
    logger.info("baostock 登录成功")

    all_codes: Optional[List[Tuple[str, str]]] = None
    any_circuit = False

    try:
        for t in types:
            logger.info("-" * 50)
            logger.info(">>> 类型: %s", t)

            if t == DIR_BASIC:
                status, rows = download_basic(args.force)
                logger.info("[basic] status=%s rows=%d", status, rows)
                if status == "failed":
                    any_circuit = True
                    logger.error("[basic] 失败, 跳过后续类型")
                    break
                continue

            # per-stock 类型: 需股票列表 (同进程内缓存, 不重复拉)
            if all_codes is None:
                try:
                    all_codes = fetch_all_stocks()
                except Exception as e:
                    logger.error("[FATAL] 获取股票列表失败: %s", e)
                    any_circuit = True
                    break
            logger.info("全市场 A 股 (含退市): %d 只", len(all_codes))

            if t == DIR_ADJUST:
                fetch_fn = partial(fetch_adjust_factor,
                                   start_date=start_date, end_date=end_date)
                _, circuit = run_per_stock_type(
                    DIR_ADJUST, SUFFIX_ADJUST, fetch_fn,
                    all_codes, args, skip_loop_sleep=False,
                )
            elif t == DIR_DIVIDEND:
                # fetch_dividend 内部已对每年 query sleep, 主循环跳过 sleep
                fetch_fn = partial(fetch_dividend,
                                   year_range=DIVIDEND_YEAR_RANGE)
                _, circuit = run_per_stock_type(
                    DIR_DIVIDEND, SUFFIX_DIVIDEND, fetch_fn,
                    all_codes, args, skip_loop_sleep=True,
                )
            else:
                logger.error("未知类型 (跳过): %s", t)
                continue

            if circuit:
                any_circuit = True
                logger.error("[%s] 断路器触发, 跳过后续类型", t)
                break
    finally:
        try:
            bs.logout()
        except Exception:
            pass

    if any_circuit:
        sys.exit(2)
    sys.exit(0)


if __name__ == "__main__":
    main()