#!/usr/bin/env python3
"""
BaoStock 15分钟线历史回补脚本 (v1.1)

功能：用BaoStock免费数据源全量重建全市场15分钟K线历史数据
策略：
  - 全量重建：BaoStock获取完整历史，完全替换旧parquet
  - 旧文件备份到 backup/ 目录
  - 增量模式：已有BaoStock重建过的文件自动跳过（检查标记）

BaoStock特点：
  - 无反爬限制，0.35s/只
  - 不复权数据（adjustflag=3）
  - 分钟线从1999年起，我们取2024-01-01起

用法：
  python3 backfill_15min_baostock.py                    # 全量回补
  python3 backfill_15min_baostock.py --start 20200101   # 指定起始日期
  python3 backfill_15min_baostock.py --limit 10         # 只处理前10只（测试）
  python3 backfill_15min_baostock.py --force             # 强制重建（覆盖已有BaoStock数据）

变更记录：
  v1.0 (2026-05-05) 赵云创建
  v1.1 (2026-05-05) 修复：旧新浪数据含后复权值，改为全量重建模式
  v1.2 (2026-07-13) 硬化：异常强制重登 / login SIGALRM 超时 / 断路器 / marker-based resume
"""

import argparse
import json
import logging
import os
import shutil
import signal
import sys
import time
from datetime import datetime
from pathlib import Path
from typing import List, Optional, Tuple

import baostock as bs
import pandas as pd

# ======================== 配置 ========================

NAS_ROOT = Path(os.environ.get("STOCK_ROOT", "/Volumes/stock"))
MINUTE_15_DIR = NAS_ROOT / "minute_kline" / "15min"
BACKUP_DIR = MINUTE_15_DIR / "backup_sina"  # 旧新浪数据备份
LOG_DIR = NAS_ROOT / "logs" / "daily_update"
PROGRESS_DIR = LOG_DIR / "progress"
ALL_STOCKS_FILE = NAS_ROOT / "A股数据" / "stock_info" / "stock_basic_info_raw_20260326_113530.csv"

# BaoStock配置
BS_START_DATE = os.environ.get("BS_START_DATE", "2024-01-01")
BS_INTERVAL = 0.4            # 请求间隔秒
BS_MAX_RETRIES = 3
PROGRESS_SAVE_EVERY = 500
BS_LOGIN_TIMEOUT = 30        # bs.login() 超时秒数（防永久挂起）
CIRCUIT_BREAKER_THRESHOLD = 30  # 连续失败N只 → 断路退出

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

def setup_logging():
    LOG_DIR.mkdir(parents=True, exist_ok=True)
    ts = datetime.now().strftime("%Y%m%d_%H%M%S")
    # 日志写本地避免SMB阻塞
    log_file = Path(f"/tmp/backfill_15min_{ts}.log")

    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s %(levelname)s %(message)s",
        handlers=[
            logging.FileHandler(log_file, encoding="utf-8"),
            logging.StreamHandler(),
        ],
    )
    return logging.getLogger(__name__), log_file

logger, LOG_FILE = setup_logging()


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

def get_all_codes() -> List[str]:
    df = pd.read_csv(ALL_STOCKS_FILE)
    for col in ["代码", "code", "股票代码"]:
        if col in df.columns:
            return [str(c).zfill(6) for c in df[col].tolist()]
    raise ValueError(f"找不到代码列: {list(df.columns)}")


def code_to_baostock(code: str) -> Tuple[str, str]:
    """6位代码 → (BaoStock格式, parquet前缀)"""
    if code.startswith(("6", "68", "51")):
        return f"sh.{code}", "sh"
    else:
        return f"sz.{code}", "sz"


def is_backfilled(parquet_path: Path) -> bool:
    """检查文件是否已经被BaoStock回补过"""
    # 用标记文件判断
    marker = parquet_path.parent / f".{parquet_path.stem}.baostock"
    return marker.exists()


def load_progress() -> set:
    """
    从 marker 文件构造已完成集合（真相源）。
    只有 .{prefix}{code}_15min.baostock marker 存在 = 该票已成功回填。
    failed 票无 marker，下次复跑会被重试，避免缺口永不重试。
    JSON 进度文件只用于观察，不再决定 todo。
    """
    done = set()
    if not MINUTE_15_DIR.exists():
        return done
    suffix = "_15min.baostock"
    for marker in MINUTE_15_DIR.glob(f".*{suffix}"):
        name = marker.name                       # .sh600000_15min.baostock
        stripped = name[1:-len(suffix)]          # sh600000
        if len(stripped) < 8:                    # sh/sz(2) + 6位代码
            continue
        code = stripped[2:]                      # 600000
        if len(code) == 6 and code.isdigit():
            done.add(code)
    return done


def save_progress(done_set: set):
    """写 JSON 进度文件（仅观察用，load_progress 不再读取）"""
    PROGRESS_DIR.mkdir(parents=True, exist_ok=True)
    progress_file = PROGRESS_DIR / "backfill_15min_progress.json"
    progress_file.write_text(json.dumps({
        "done": sorted(done_set),
        "ts": datetime.now().isoformat(),
    }, ensure_ascii=False))


# ======================== BaoStock 连接管理 ========================

class _LoginTimeout(Exception):
    """bs.login() SIGALRM 超时"""


def _login_timeout_handler(signum, frame):
    raise _LoginTimeout("bs.login() 超时")


def _login_with_timeout(timeout: int = BS_LOGIN_TIMEOUT) -> bool:
    """
    给 bs.login() 套 SIGALRM 超时，防止卡死时整个回填永久挂起。
    signal.alarm 只能在主线程用（本脚本主线程，OK）。
    返回 True=登录成功，False=失败/超时。
    """
    old_handler = signal.signal(signal.SIGALRM, _login_timeout_handler)
    signal.alarm(timeout)
    try:
        lg = bs.login()
        ok = lg.error_code == "0"
        if not ok:
            logger.error("bs.login() 失败: %s", lg.error_msg)
        return ok
    except _LoginTimeout:
        logger.error("bs.login() 超时 %d 秒（baostock 疑似限流冷却中）", timeout)
        return False
    except Exception as e:
        logger.error("bs.login() 异常: %s", e)
        return False
    finally:
        signal.alarm(0)
        signal.signal(signal.SIGALRM, old_handler)


def _relogin() -> bool:
    """
    强制重登：logout + login（强制建新 socket）。
    登录失败再试 1 次。baostock 是模块级单连接，本脚本单进程串行，重登安全。
    """
    try:
        bs.logout()
    except Exception:
        pass
    if _login_with_timeout():
        return True
    # 第一次失败，等2秒再试一次
    time.sleep(2)
    try:
        bs.logout()
    except Exception:
        pass
    return _login_with_timeout()


# ======================== 核心逻辑 ========================

def fetch_bs_15min(bs_code: str, start_date: str, end_date: str) -> Optional[pd.DataFrame]:
    """从BaoStock获取15min不复权数据"""
    rs = bs.query_history_k_data_plus(
        bs_code,
        "date,time,code,open,high,low,close,volume,amount,adjustflag",
        start_date=start_date,
        end_date=end_date,
        frequency="15",
        adjustflag="3",  # 不复权
    )

    if rs.error_code != "0":
        logger.debug("BaoStock %s 错误: %s %s", bs_code, rs.error_code, rs.error_msg)
        return None

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

    if not rows:
        return None

    df = pd.DataFrame(rows, columns=["date", "time", "code", "open", "high", "low",
                                      "close", "volume", "amount", "adjustflag"])

    # 转换时间格式: 20260428094500000 → 2026-04-28 09:45:00
    df["day"] = df["time"].apply(lambda t: f"{t[:4]}-{t[4:6]}-{t[6:8]} {t[8:10]}:{t[10:12]}:00")

    # 数值转换
    for col in ["open", "high", "low", "close"]:
        df[col] = pd.to_numeric(df[col], errors="coerce")
    df["volume"] = df["volume"].astype(str)
    df["amount"] = df["amount"].astype(str)

    # 保留与parquet一致的列
    df = df[["day", "open", "high", "low", "close", "volume", "amount"]]

    # 过滤无效数据
    df = df.dropna(subset=["open", "high", "low", "close"])
    bad_ohlc = (df["high"] < df[["open", "close"]].max(axis=1)) | \
               (df["low"] > df[["open", "close"]].min(axis=1))
    if bad_ohlc.any():
        df = df[~bad_ohlc]

    return df if not df.empty else None


def backfill_one(code: str, start_date: str, end_date: str, force: bool = False) -> Tuple[str, int]:
    """
    全量重建单只股票的15min历史
    返回: (status, total_rows)
    """
    bs_code, prefix = code_to_baostock(code)
    parquet_path = MINUTE_15_DIR / f"{prefix}{code}_15min.parquet"

    # 已回补过的跳过
    if not force and is_backfilled(parquet_path):
        return "skipped", 0

    # 获取BaoStock数据
    df_new = None
    for attempt in range(BS_MAX_RETRIES):
        try:
            df_new = fetch_bs_15min(bs_code, start_date, end_date)
            if df_new is not None and len(df_new) > 0:
                break
            # 空数据：可能是真的无数据（新股/停牌），也可能是查询异常
            logger.debug("backfill %s 空数据重试 %d/%d", code, attempt + 1, BS_MAX_RETRIES)
        except Exception as e:
            # socket 死了（如 [Errno 57] Socket is not connected），同 socket 重试必然再败
            # 强制重登建新 socket 再 retry
            logger.warning("backfill %s 异常重试 %d/%d: %s — 强制重登",
                           code, attempt + 1, BS_MAX_RETRIES, e)
            if not _relogin():
                logger.error("重登失败，放弃 %s", code)
                df_new = None
                break

    if df_new is None or df_new.empty:
        return "failed", 0

    # 备份旧文件
    if parquet_path.exists():
        BACKUP_DIR.mkdir(parents=True, exist_ok=True)
        backup_path = BACKUP_DIR / parquet_path.name
        if not backup_path.exists():  # 不覆盖已有备份
            shutil.copy2(str(parquet_path), str(backup_path))

    # 写入新文件
    try:
        df_new = df_new.sort_values("day").reset_index(drop=True)
        df_new.to_parquet(parquet_path, index=False)
        # 写标记文件
        marker = parquet_path.parent / f".{parquet_path.stem}.baostock"
        marker.write_text(datetime.now().isoformat())
        return "ok", len(df_new)
    except Exception as e:
        logger.error("写入 %s 失败: %s", code, e)
        return "failed", 0


# ======================== 主流程 ========================

def main():
    parser = argparse.ArgumentParser(description="BaoStock 15min历史回补（全量重建）")
    parser.add_argument("--start", default=BS_START_DATE, help="回补起始日期 (YYYYMMDD)")
    parser.add_argument("--end", default="", help="结束日期 (默认今天)")
    parser.add_argument("--codes", help="指定股票代码，逗号分隔")
    parser.add_argument("--limit", type=int, default=0, help="限制处理数量（测试用）")
    parser.add_argument("--force", action="store_true", help="强制重建（覆盖已有BaoStock数据）")
    parser.add_argument("--shard", type=int, default=None, help="分片号 (0-indexed，并行回填用)")
    parser.add_argument("--total", type=int, default=1, help="总分片数（与 --shard 配合）")
    args = parser.parse_args()

    start_date = f"{args.start[:4]}-{args.start[4:6]}-{args.start[6:8]}"
    end_date = f"{args.end[:4]}-{args.end[4:6]}-{args.end[6:8]}" if args.end else datetime.now().strftime("%Y-%m-%d")

    if not NAS_ROOT.exists():
        logger.error("❌ NAS未挂载")
        sys.exit(1)

    # 登录BaoStock（带超时，防永久挂起）
    if not _login_with_timeout():
        logger.error("❌ BaoStock登录失败（baostock 疑似冷却中），请稍后重试")
        sys.exit(1)
    logger.info("✅ BaoStock登录成功")

    # 获取股票列表
    if args.codes:
        codes = [c.strip() for c in args.codes.split(",")]
    else:
        codes = get_all_codes()

    if args.limit > 0:
        codes = codes[:args.limit]

    # 分片过滤（并行回填：每个 worker 处理 i % total == shard 的子集）
    if args.shard is not None and args.total > 1:
        codes = [s for i, s in enumerate(codes) if i % args.total == args.shard]

    logger.info("=" * 60)
    logger.info("BaoStock 15min全量重建开始")
    logger.info(f"  股票数: {len(codes)}")
    if args.shard is not None:
        logger.info(f"  分片: shard={args.shard}/{args.total}")
    logger.info(f"  日期范围: {start_date} ~ {end_date}")
    logger.info(f"  数据目录: {MINUTE_15_DIR}")
    logger.info(f"  旧数据备份: {BACKUP_DIR}")
    logger.info(f"  日志文件: {LOG_FILE}")

    # 进度恢复
    done_set = load_progress()
    if args.force:
        todo = codes  # 强制模式不跳过
    else:
        todo = [c for c in codes if c not in done_set]
    logger.info(f"  待处理: {len(todo)}（已完成: {len(done_set)}）")

    stats = {"ok": 0, "skipped": 0, "failed": 0, "rows": 0}
    t_start = time.time()
    RELOGIN_EVERY = 400  # 每400只重新登录BaoStock，防止连接断开
    consecutive_failures = 0  # 连续失败计数（断路器）：ok 重置，failed +1，skipped 中性
    circuit_triggered = False

    for i, code in enumerate(todo):
        # 定期重新登录保持连接（用带超时的 _relogin，防永久挂起）
        if i > 0 and i % RELOGIN_EVERY == 0:
            if not _relogin():
                logger.error("BaoStock 定期重连失败，等待30秒后再试")
                time.sleep(30)
                _relogin()
            logger.info("BaoStock重连 @ %d/%d", i, len(todo))

        try:
            status, total_rows = backfill_one(code, start_date, end_date, args.force)
        except Exception as e:
            status, total_rows = "failed", 0
            logger.debug("backfill %s 异常: %s", code, e)

        stats[status] = stats.get(status, 0) + 1
        if status == "ok":
            stats["rows"] += total_rows
            consecutive_failures = 0
        elif status == "failed":
            consecutive_failures += 1
        # skipped 中性：不重置（不证明 baostock 可用），不递增（不算失败）

        done_set.add(code)

        # 断路器：连续N只全 failed → baostock 疑似不可达，保存进度主动退出
        if consecutive_failures >= CIRCUIT_BREAKER_THRESHOLD:
            logger.error(
                "❌ 断路器触发：连续 %d 只失败，baostock 疑似不可达，"
                "保存进度后退出（等冷却后复跑，done_set 不含 failed 票会被重试）",
                consecutive_failures)
            save_progress(done_set)
            try:
                bs.logout()
            except Exception:
                pass
            circuit_triggered = True
            break

        if (i + 1) % PROGRESS_SAVE_EVERY == 0:
            save_progress(done_set)
            elapsed = time.time() - t_start
            logger.info("进度: %d/%d ok=%d skipped=%d failed=%d rows=%d (%.0f秒)",
                        i + 1, len(todo), stats["ok"], stats["skipped"],
                        stats["failed"], stats["rows"], elapsed)

        # 频率控制
        if i < len(todo) - 1:
            time.sleep(BS_INTERVAL)

    # 保存最终进度
    save_progress(done_set)

    try:
        bs.logout()
    except Exception:
        pass

    elapsed = time.time() - t_start
    logger.info("=" * 60)
    if circuit_triggered:
        logger.info("❌ 回补因断路器中止，耗时 %.1f 秒", elapsed)
    else:
        logger.info("✅ 回补完成，耗时 %.1f 秒", elapsed)
    logger.info("统计: %s", json.dumps(stats, ensure_ascii=False))

    # 断路器触发 → 非零退出码（cron/调度器可识别）
    if circuit_triggered:
        sys.exit(2)


if __name__ == "__main__":
    main()