#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""BaoStock 成分股历史下载脚本 (只产 parquet)

目标:
  在 Mac 本机用 baostock 拉取 HS300/ZZ500/SZ50 历史成分股快照,
  按历史日期循环 (每周一一个快照, 从 2006-01 至今)。
  输出 parquet 到 staging 目录。

硬约束 (踩过的坑):
  1. 单进程单登录, 严禁并发 (baostock 并发会拉黑封 IP 6-24h)
  2. 直连不走代理 (脚本开头 unset proxy)
  3. 每次 query 后 sleep 0.4s 限速
  4. staging parquet, 绝不直接写主库

用法:
  python baostock_constituent_download.py                    # 全指数 2006-01~至今
  python baostock_constituent_download.py --limit 5          # 测试: 前 5 个快照
  python baostock_constituent_download.py --indices hs300    # 只跑指定指数
  python baostock_constituent_download.py --start 2020-01-01  # 指定起始日期

退出码: 0=完成, 1=致命错误, 2=断路器触发 (可重试)
"""

import argparse
import datetime
import json
import logging
import os
import socket
import sys
import time
from datetime import date, timedelta
from pathlib import Path
from typing import Dict, List, Optional, Tuple

# ======================== 硬约束: unset proxy + socket timeout ========================
for _k in ("http_proxy", "https_proxy", "HTTP_PROXY", "HTTPS_PROXY"):
    os.environ.pop(_k, None)

socket.setdefaulttimeout(30)

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 = "/Users/chufeng/.openclaw/sanguo_projects/sanguo_vnpy_v2/data/constituent_baostock"
OUT_DIR = Path(os.environ.get("BS_CONSTITUENT_OUT_DIR", DEFAULT_OUT_DIR))
DEFAULT_LOG_DIR = OUT_DIR / "logs"
_log_dir_env = os.environ.get("BS_CONSTITUENT_LOG_DIR")
LOG_DIR = Path(_log_dir_env) if _log_dir_env else DEFAULT_LOG_DIR

START_DATE_DEFAULT = "2006-01-01"  # 成分股历史起始
END_DATE_DEFAULT = ""  # 默认今天

BS_INTERVAL = 0.4             # 每次间隔秒 (防封 IP)
BS_MAX_RETRIES = 3            # 单次下载重试次数
PROGRESS_LOG_EVERY = 50       # 每 N 个快照打一次进度日志
RELOGIN_EVERY = 100           # 每 N 个快照定期重登
CIRCUIT_BREAKER = 10          # 连续失败 N 次 → 断路退出

# 指数配置 (baostock code, 显示名)
INDICES_CONFIG = {
    "hs300": {"code": "000300", "name": "HS300"},
    "zz500": {"code": "000905", "name": "ZZ500"},
    "sz50": {"code": "000016", "name": "SZ50"},
}


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

def setup_logging() -> Tuple[logging.Logger, Path]:
    """配置 root logger: 同时写 stdout + 文件。"""
    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_constituent_{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}")


def get_monday_dates(start_date: str, end_date: str) -> List[str]:
    """生成每周一日期列表 (包含 start_date 所在周, 含 end_date)。

    返回: YYYY-MM-DD 格式日期列表。
    """
    start = datetime.datetime.strptime(start_date, "%Y-%m-%d").date()
    end = datetime.datetime.strptime(end_date, "%Y-%m-%d").date()

    # 找到 start_date 后的第一个周一
    d = start
    while d.weekday() != 0:  # Monday=0
        d += timedelta(days=1)

    dates = []
    while d <= end:
        dates.append(d.strftime("%Y-%m-%d"))
        d += timedelta(days=7)  # 下一周

    return dates


def baostock_code_to_parts(bs_code: str) -> Tuple[str, str]:
    """sh.600519 → (600519, 'SH'); sz.000001 → (000001, 'SZ')。"""
    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}")


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

def _login_once() -> bool:
    """bs.login() — 已被 socket.setdefaulttimeout(30) 保护。"""
    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()


# ======================== 数据下载 ========================

def fetch_constituent_stocks(index_name: str, date_str: str) -> pd.DataFrame:
    """拉取指定指数在指定日期的成分股。

    返回 df 列: code (baostock格式, 如 sh.600519)
    空数据返回空 df。
    """
    index_info = INDICES_CONFIG.get(index_name)
    if not index_info:
        raise ValueError(f"未知指数: {index_name}")

    index_code = index_info["code"]

    if index_name == "hs300":
        rs = bs.query_hs300_stocks(date=date_str)
    elif index_name == "zz500":
        rs = bs.query_zz500_stocks(date=date_str)
    elif index_name == "sz50":
        rs = bs.query_sz50_stocks(date=date_str)
    else:
        raise ValueError(f"未实现的指数: {index_name}")

    if rs.error_code != "0":
        raise RuntimeError(
            f"query_{index_name}_stocks({date_str}) 错误: {rs.error_code} {rs.error_msg}"
        )

    if rs.error_code != "0":
        raise RuntimeError(
            f"query_{index_name}_stocks({date_str}) 错误: {rs.error_code} {rs.error_msg}"
        )

    # 获取字段列表
    fields = list(rs.fields)

    rows = []
    while rs.next():
        rows.append(rs.get_row_data())

    if not rows:
        # 合法空数据 (指数不存在于该日期)
        return pd.DataFrame(columns=fields)

    # 使用 baostock 实际返回的字段名
    return pd.DataFrame(rows, columns=fields)


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

def parquet_path_for(index_name: str, date_str: str) -> Path:
    """hs300 + 2020-01-06 → OUT_DIR / hs300_2020-01-06.parquet"""
    return OUT_DIR / f"{index_name}_{date_str}.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(indices: List[str]) -> Dict[str, set]:
    """扫 OUT_DIR 所有 marker 构造已完成集合。

    返回: {index_name: set(dates)} 字典。
    """
    done: Dict[str, set] = {idx: set() for idx in indices}
    if not OUT_DIR.exists():
        return done

    suffix = ".baostock"
    for marker in OUT_DIR.glob(f".*{suffix}"):
        name = marker.name
        if not name.startswith(".") or not name.endswith(suffix):
            continue
        stem = name[1:-len(suffix)]  # e.g. hs300_2020-01-06

        # 解析: <index_name>_<date>
        try:
            parts = stem.split("_", 1)
            if len(parts) != 2:
                continue
            index_name, date_str = parts
            if index_name in done:
                done[index_name].add(date_str)
        except ValueError:
            continue
    return done


# ======================== 单次下载 ========================

def download_one_snapshot(
    index_name: str,
    date_str: str,
    force: bool,
) -> Tuple[str, int]:
    """下载单个成分股快照 → 写 parquet + marker。

    返回 (status, rows): status ∈ {'ok', 'skipped', 'failed', 'empty'}。
    """
    parquet_path = parquet_path_for(index_name, date_str)
    marker_path = marker_path_for(parquet_path)

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

    # retry 循环
    df: Optional[pd.DataFrame] = None
    for attempt in range(BS_MAX_RETRIES):
        try:
            df = fetch_constituent_stocks(index_name, date_str)
            break
        except (socket.timeout, TimeoutError, OSError) as e:
            logger.warning(
                "%s %s socket 异常重试 %d/%d: %s — 强制重登",
                index_name, date_str, attempt + 1, BS_MAX_RETRIES, e,
            )
            if not _relogin():
                logger.error("重登失败, 放弃 %s %s", index_name, date_str)
                df = None
                break
        except Exception as e:
            logger.warning(
                "%s %s 下载异常重试 %d/%d: %s — 强制重登",
                index_name, date_str, attempt + 1, BS_MAX_RETRIES, e,
            )
            if not _relogin():
                logger.error("重登失败, 放弃 %s %s", index_name, date_str)
                df = None
                break

    if df is None:
        return "failed", 0

    if df.empty:
        # 合法空数据 (指数不存在于该日期) — 也写 marker
        try:
            OUT_DIR.mkdir(parents=True, exist_ok=True)
            df.to_parquet(parquet_path, index=False)
            marker_path.write_text(datetime.datetime.now().isoformat())
            return "empty", 0
        except Exception as e:
            logger.error("写入 %s 失败: %s", parquet_path, e)
            return "failed", 0

    # 写 parquet + marker
    try:
        OUT_DIR.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 parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(
        description="BaoStock 成分股历史下载 (只产 parquet)",
    )
    p.add_argument(
        "--start", default=START_DATE_DEFAULT,
        help=f"起始日期 YYYYMMDD 或 YYYY-MM-DD, 默认 {START_DATE_DEFAULT}",
    )
    p.add_argument("--end", default=END_DATE_DEFAULT, help="结束日期, 默认今天")
    p.add_argument(
        "--indices",
        default="hs300,zz500,sz50",
        help="逗号分隔指数名称, 默认 hs300,zz500,sz50",
    )
    p.add_argument("--limit", type=int, default=0, help="限制处理快照数 (测试用)")
    p.add_argument("--force", action="store_true", help="强制重下, 忽略 marker")
    return p.parse_args()


def resolve_indices(s: str) -> List[str]:
    """解析 --indices 字符串 → 去重保序的指数列表。"""
    parts = [p.strip() for p in s.split(",") if p.strip()]
    valid = set(INDICES_CONFIG.keys())
    bad = [p for p in parts if p not in valid]
    if bad:
        raise SystemExit(f"未知 --indices: {bad}, 可选 {list(valid)}")
    seen: set = set()
    out: List[str] = []
    for p in parts:
        if p not in seen:
            out.append(p)
            seen.add(p)
    return out


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

    logger.info("=" * 60)
    logger.info("BaoStock 成分股历史下载")
    logger.info("  输出目录: %s", OUT_DIR)
    logger.info("  日志文件: %s", LOG_FILE)
    logger.info("  日期范围: %s ~ %s", start_date, end_date)
    logger.info("  指数: %s", indices)
    logger.info("  频率: 每周一快照")
    logger.info("  socket.setdefaulttimeout(30)")
    logger.info("  当前时间: %s", datetime.datetime.now().isoformat())
    logger.info("=" * 60)

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

    # 生成日期列表 (每周一)
    monday_dates = get_monday_dates(start_date, end_date)
    logger.info("快照日期列表: %d 个周一 (%s ~ %s)", len(monday_dates),
                monday_dates[0] if monday_dates else "N/A",
                monday_dates[-1] if monday_dates else "N/A")

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

    # 总快照数
    total_snapshots = len(monday_dates) * len(indices)
    logger.info("总快照数: %d (日期%d × 指数%d)", total_snapshots,
                len(monday_dates), len(indices))

    # 断点续传
    done_set = load_done_set(indices)
    if args.force:
        todo_snapshots = [(idx, d) for idx in indices for d in monday_dates]
    else:
        todo_snapshots = [
            (idx, d) for idx in indices for d in monday_dates
            if d not in done_set.get(idx, set())
        ]
    logger.info("待处理: %d (已完成 %d)", len(todo_snapshots),
                total_snapshots - len(todo_snapshots))

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

    for i, (index_name, date_str) in enumerate(todo_snapshots):
        # 定期重登
        if i > 0 and i % RELOGIN_EVERY == 0:
            logger.info("定期重登 @ %d/%d", i, len(todo_snapshots))
            if not _relogin():
                logger.warning("定期重登失败, 继续")

        try:
            status, rows = download_one_snapshot(index_name, date_str, args.force)
        except Exception as e:
            status, rows = "failed", 0
            logger.debug("download_one_snapshot %s %s 异常: %s", index_name, date_str, e)

        stats[status] = stats.get(status, 0) + 1
        if status == "ok":
            stats["rows"] += rows
            consec_fail = 0
        elif status == "failed":
            consec_fail += 1

        # 断路器
        if consec_fail >= CIRCUIT_BREAKER:
            logger.error(
                "[FATAL] 断路器触发: 连续 %d 次失败, baostock 疑似不可达",
                consec_fail,
            )
            circuit_triggered = True
            break

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

        # 限速
        if i < len(todo_snapshots) - 1:
            time.sleep(BS_INTERVAL)

    # 登出
    try:
        bs.logout()
    except Exception:
        pass

    elapsed = time.time() - t_start
    logger.info("=" * 60)
    if circuit_triggered:
        logger.info("[RESULT] 断路器触发中止, 耗时 %.1f 秒", elapsed)
    else:
        logger.info("[RESULT] 完成, 耗时 %.1f 秒", elapsed)
    logger.info("统计: %s", json.dumps(stats, ensure_ascii=False))

    sys.exit(2 if circuit_triggered else 0)


if __name__ == "__main__":
    main()