#!/usr/bin/env python3
"""raw 日线重下（task #79）：akshare 新浪源 adjust="" 真实价。

根因：daily_dir 历史数据是 mixed-adjust（hfq bulk + akshare raw tail 拼接），
3-30 类单日 -94% 假跌。本脚本从头重下**单一 raw** 到 raw_dir，与 daily_dir 同构
（{prefix}{symbol}_daily.parquet，按 year 分目录），datareader 直接读。

**约束（用户反馈，见 memory/feedback-data-download-constraints）**：
- 直连不走代理（unset proxy env + NO_PROXY=*）
- 单线程 + SLEEP 间隔限速（不并发猛打，防数据源封 IP）

用法：
  # 验证单只/几只
  python3 raw_redownload.py --symbols 600000,000001 --start 2024-01-01
  # 全市场（读 STOCK_LIST csv，后台跑）
  python3 raw_redownload.py --all --start 2024-01-01

env:
  RAW_DIR      本地 raw 根（默认 /tmp/stock_dl/A股数据/日线数据/raw）
  SLEEP        每只请求间隔秒（默认 1.0，限速防封）
  STOCK_LIST   stock_basic_info csv 路径（--all 读，默认拉到本地的 csv）
"""
import argparse
import csv
import logging
import os
import socket
import sys
import time
from collections import deque

# 进程级 socket 超时（关键）：akshare 内部 requests 默认无 timeout，
# 遇源限速/慢响应会无限挂起整进程（实测 35 只后 socket 死等 → hang）。
# 15s 足够区分正常慢响应与挂死，超时即抛 → download_one 捕获 → fail → 继续下一只。
socket.setdefaulttimeout(15)

# 直连：进程级 unset 代理（用户约束）
for _k in ["HTTP_PROXY", "HTTPS_PROXY", "http_proxy", "https_proxy", "ALL_PROXY", "all_proxy"]:
    os.environ.pop(_k, None)
os.environ["NO_PROXY"] = "*"
os.environ["no_proxy"] = "*"

import pandas as pd  # noqa: E402

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

RAW_DIR = os.environ.get("RAW_DIR", "/tmp/stock_dl/A股数据/日线数据/raw")
SLEEP = float(os.environ.get("SLEEP", "1.0"))
DEFAULT_STOCK_LIST = "/tmp/stock_dl/A股数据/stock_info/stock_basic_info_raw_20260326_113530.csv"

# 断路器（借鉴 15min backfill 自愈范式）
CIRCUIT_BREAKER_WINDOW = 200       # 滚动窗口大小
CIRCUIT_BREAKER_FAIL_RATE = 0.30   # 触发阈值：失败率 > 30%
# 北交所码新浪不支持（必 fail，从断路器分母扣除，防误触发）
KNOWN_UNSUPPORTED_PREFIX = ("920", "921", "83", "87")


def prefix_for(code: str) -> str:
    """sh/sz 前缀（与 datareader.guess_exchange 一致）。"""
    return "sh" if code.startswith(("60", "68", "51", "56", "58")) else "sz"


def check_circuit_breaker(recent_results: list) -> bool:
    """检查滚动窗口内失败率是否超阈值（断路器触发判定）。

    recent_results: [(code, ok: bool), ...] 最近下载结果。
    北交所已知不支持码从分子分母同时扣除（它们必 fail，计入会误触发）。
    返回 True 表示应触发断路器。
    """
    counted = [(c, ok) for c, ok in recent_results
               if not c.startswith(KNOWN_UNSUPPORTED_PREFIX)]
    if len(counted) < CIRCUIT_BREAKER_WINDOW:
        return False
    fails = sum(1 for _, ok in counted if not ok)
    return (fails / len(counted)) > CIRCUIT_BREAKER_FAIL_RATE


def download_one(ak, code: str, start: str, end: str, adjust: str = ""):
    """新浪源拉日线（adjust="" raw / "qfq" 前复权），返回 (df, None) 或 (None, err)。"""
    sym = f"{prefix_for(code)}{code}"
    try:
        df = ak.stock_zh_a_daily(
            symbol=sym,
            start_date=start.replace("-", ""),
            end_date=end.replace("-", ""),
            adjust=adjust,
        )
    except Exception as e:  # noqa: BLE001
        return None, f"{type(e).__name__}: {str(e)[:100]}"
    if df is None or df.empty:
        return None, "empty"
    df["date"] = pd.to_datetime(df["date"])
    df["year"] = df["date"].dt.year
    return df, None


def save_one(code: str, df) -> int:
    """按 year 分组写 parquet（与 daily_dir 同构），返回写入行数。"""
    n = 0
    for year, g in df.groupby("year"):
        ydir = os.path.join(RAW_DIR, str(int(year)))
        os.makedirs(ydir, exist_ok=True)
        out = os.path.join(ydir, f"{prefix_for(code)}{code}_daily.parquet")
        g.drop(columns=["year"]).to_parquet(out)
        n += len(g)
    return n


def latest_trading_day(end: str) -> str:
    """end 往前最近的工作日（周一~周五）。忽略节假日——增量天天跑会自愈。"""
    d = pd.Timestamp(end)
    while d.weekday() >= 5:  # 5=Sat, 6=Sun
        d -= pd.Timedelta(days=1)
    return d.strftime("%Y-%m-%d")


def is_fresh(code: str, end: str, start_year: int) -> bool:
    """symbol 的 parquet 数据是否已到最新交易日（增量断点续传）。

    文件不存在 / 损坏 / 最新日期 < latest_trading_day(end) → 不 fresh → 重下。
    旧版 exists() 只查文件在不在，对"文件存在但数据旧"的增量场景失效
    （raw/2026 已有旧 parquet → 全 skip → 补数永远不写入）。
    """
    pref = prefix_for(code)
    f = os.path.join(RAW_DIR, str(start_year), f"{pref}{code}_daily.parquet")
    if not os.path.exists(f):
        return False
    try:
        maxd = pd.read_parquet(f, columns=["date"])["date"].max()
    except Exception:  # noqa: BLE001  损坏文件 → 重下
        return False
    return pd.Timestamp(maxd) >= pd.Timestamp(latest_trading_day(end))


def load_all_codes(stock_list: str) -> list[str]:
    """从 stock_basic_info csv 读代码列表（容错列名）。"""
    codes = []
    with open(stock_list, encoding="utf-8", errors="replace") as f:
        reader = csv.DictReader(f)
        for row in reader:
            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", help="逗号分隔代码，如 600000,000001")
    ap.add_argument("--all", action="store_true", help="全市场（读 STOCK_LIST csv）")
    ap.add_argument("--start", default="2024-01-01")
    ap.add_argument("--end", default=None, help="默认今天")
    ap.add_argument("--adjust", default="", help="复权: '' raw / 'qfq' 前复权（双源用）")
    ap.add_argument("--force", action="store_true", help="强制重下（默认 skip 已存在=断点续传）")
    args = ap.parse_args()

    end = args.end or time.strftime("%Y-%m-%d")

    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（先跑 run_daily_update.sh 拉取）", sl)
            sys.exit(1)
        codes = load_all_codes(sl)
        log.info("--all 从 %s 读到 %d 只", sl, len(codes))
    else:
        ap.error("需指定 --symbols 或 --all")

    import akshare as ak
    import warnings
    warnings.filterwarnings("ignore")

    start_year = int(args.start[:4])
    ok = fail = rows = skipped = 0
    circuit_recent: deque = deque(maxlen=CIRCUIT_BREAKER_WINDOW)
    circuit_count = 0  # 非北交所码下载计数（断路器窗口）
    for i, code in enumerate(codes, 1):
        if not args.force and is_fresh(code, end, start_year):
            skipped += 1
            if skipped % 500 == 0:
                log.info("[%d/%d] ... skipped %d 已最新", i, len(codes), skipped)
            continue
        df, err = download_one(ak, code, args.start, end, args.adjust)
        download_ok = False
        if df is None:
            fail += 1
            log.warning("[%d/%d] %s FAIL %s", i, len(codes), code, err)
        else:
            try:
                n = save_one(code, df)
                ok += 1
                rows += n
                download_ok = True
                log.info("[%d/%d] %s ok %d rows", i, len(codes), code, n)
            except Exception as e:  # noqa: BLE001
                fail += 1
                log.error("[%d/%d] %s SAVE FAIL %s", i, len(codes), code, e)
        # 断路器：北交所码不计入（必 fail 会误触发）
        if not code.startswith(KNOWN_UNSUPPORTED_PREFIX):
            circuit_recent.append((code, download_ok))
            circuit_count += 1
            if circuit_count % CIRCUIT_BREAKER_WINDOW == 0:
                if check_circuit_breaker(list(circuit_recent)):
                    fails = sum(1 for _, ok2 in circuit_recent if not ok2)
                    rate = fails / len(circuit_recent) * 100
                    log.error("断路器触发：最近 %d 只失败率 %.1f%%，abort",
                              len(circuit_recent), rate)
                    sys.exit(3)
        time.sleep(SLEEP)  # 限速（用户约束）

    log.info("=== 完成: ok=%d skip=%d fail=%d rows=%d，raw_dir=%s ===", ok, skipped, fail, rows, RAW_DIR)


if __name__ == "__main__":
    main()