#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""baostock 日增量 → daily_baostock_full (VPS 直跑, schtask 17:00)

设计 (用户决策 A: VPS 直跑增量, 守 48000/天):
  - 跑前 login 探针: 黑名单(10001011)则 graceful skip (exit 2), 不硬跑
  - LOOKBACK_DAYS=7 覆盖周末/节假日 (baostock 日终更新, 17:00 跑当日 bar 已就绪)
  - query_stock_basic() 无参 1 拉全 A (含退市, 与 Day1 同口径) → 每股 1 query
  - 每股 query_history_k_data_plus(LOOKBACK 窗口) → staging parquet (审计) → executemany INSERT OR REPLACE daily_baostock_full
  - 幂等: PK(symbol,exchange,date) + INSERT OR REPLACE, 重复跑安全
  - QUERY_COUNT 计数器 + DAILY_LIMIT 防线 (实际 ~5538/天 = 11.5%, 远低于 48000)
  - 单进程单登录, sleep 0.3s 限速, socket timeout 30s

预算 (实测): 1(list) + 5537(每股) = ~5538 query/天, 占 48000 的 11.5%, 安全。
耗时: 5537 × (sleep 0.3 + query ~0.3) ≈ 55min, 17:00 跑 18:00 前完。

复用模式:
  - baostock_daily_fullmarket_download.py: login/relogin/socket timeout/QUERY_COUNT/fetch_all_stocks
  - import_delisted_to_db.py / import_baostock_to_db.py: parse_baostock_code/executemany/WAL

用法 (VPS):
  C:\\Python310\\python.exe -X utf8 daily_update_static.py            # 正式跑
  python daily_update_static.py --limit 10                           # 小样本测
  python daily_update_static.py --dry-run                            # 只下不灌
环境变量: LOOKBACK_DAYS (默认7), BS_INCREMENT_OUT_DIR (staging, 默认 data/daily_baostock_increment/{YYYYMMDD})
退出码: 0=完成; 1=致命错误; 2=黑名单 graceful skip; 3=query 超限 graceful stop
"""
import argparse
import datetime
import json
import logging
import os
import socket
import sys
import time
from pathlib import Path
from typing import List, Optional, Tuple

# ======================== 硬约束: unset proxy + socket timeout ========================
# 必须在 import baostock 之前清理 (baostock 底层 urllib 读 proxy env)
for _k in ("http_proxy", "https_proxy", "HTTP_PROXY", "HTTPS_PROXY", "all_proxy", "ALL_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

# ======================== 配置 (VPS 路径) ========================
BASE = Path(r"C:\sanguo_vnpy_v2")
DB_PATH = BASE / "data" / "quant_trading.db"

LOOKBACK_DAYS = int(os.environ.get("LOOKBACK_DAYS", "7"))
_date_tag = datetime.date.today().strftime("%Y%m%d")
STAGING = Path(os.environ.get(
    "BS_INCREMENT_OUT_DIR",
    str(BASE / "data" / "daily_baostock_increment" / _date_tag),
))

# baostock 18 字段 (与 daily_baostock_full 同口径, adjustflag=3 raw)
BS_FIELDS = ("date,code,open,high,low,close,preclose,volume,amount,adjustflag,"
             "turn,tradestatus,pctChg,peTTM,psTTM,pcfNcfTTM,pbMRQ,isST")
OUT_COLS = BS_FIELDS.split(",")
NUMERIC_COLS = ["open", "high", "low", "close", "preclose", "volume", "amount",
                "turn", "pctChg", "peTTM", "psTTM", "pcfNcfTTM", "pbMRQ"]

BS_INTERVAL = 0.3
BS_MAX_RETRIES = 3
DAILY_LIMIT = 40000  # 防线 (实际 ~5538); 接近则 graceful stop
QUERY_COUNT = 0

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s %(levelname)s %(message)s",
    handlers=[logging.StreamHandler(sys.stdout)],
)
logger = logging.getLogger(__name__)


# ======================== 登录 ========================
def login_once() -> bool:
    """login 探针。黑名单/超时返回 False (调用方 graceful skip)。"""
    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", e)
        return False
    except Exception as e:
        logger.error("bs.login() 异常: %s", e)
        return False


def relogin() -> bool:
    try:
        bs.logout()
    except Exception:
        pass
    if login_once():
        return True
    time.sleep(2)
    try:
        bs.logout()
    except Exception:
        pass
    return login_once()


# ======================== 股票列表 (query_stock_basic 无参, 含退市) ========================
def fetch_all_stocks() -> List[Tuple[str, str]]:
    """query_stock_basic() 无参 → 全 A (type=1, 含退市 status=0), 与 Day1 同口径。
    返回 [(code, exchange), ...], 例 [('600519', 'SH'), ('000001', 'SZ')]。"""
    global QUERY_COUNT
    QUERY_COUNT += 1
    rs = bs.query_stock_basic()
    if rs.error_code != "0":
        raise RuntimeError(f"query_stock_basic 失败: {rs.error_code} {rs.error_msg}")
    fields = list(rs.fields)
    idx = {n: i for i, n in enumerate(fields)}
    out = []
    while rs.next():
        r = rs.get_row_data()
        type_ = r[idx["type"]] if "type" in idx else ""
        if type_ != "1":  # 只要股票
            continue
        bs_code = r[idx["code"]] if "code" in idx else ""
        if "." not in bs_code:
            continue
        prefix, num = bs_code.split(".", 1)
        if prefix not in ("sh", "sz") or len(num) != 6 or not num.isdigit():
            continue
        out.append((num, "SH" if prefix == "sh" else "SZ"))
    return out


# ======================== 单股增量下载 ========================
def fetch_one_increment(bs_code: str, start: str, end: str) -> Optional[pd.DataFrame]:
    """query_history_k_data_plus LOOKBACK 窗口。失败抛异常 (上层 retry); 空返回空 df。"""
    global QUERY_COUNT
    QUERY_COUNT += 1
    rs = bs.query_history_k_data_plus(
        bs_code, BS_FIELDS, start_date=start, end_date=end,
        frequency="d", adjustflag="3",
    )
    if rs.error_code != "0":
        raise RuntimeError(f"{bs_code} 查询错误: {rs.error_code} {rs.error_msg}")
    rows = []
    while rs.next():
        rows.append(rs.get_row_data())
    if not rows:
        return pd.DataFrame(columns=OUT_COLS)
    df = pd.DataFrame(rows, columns=OUT_COLS)
    for col in NUMERIC_COLS:
        if col in df.columns:
            df[col] = pd.to_numeric(df[col], errors="coerce")  # 空字符串→NaN→None
    df = df.dropna(subset=["open", "high", "low", "close"])
    return df if not df.empty else pd.DataFrame(columns=OUT_COLS)


def parse_baostock_code(code: str) -> Tuple[str, str]:
    """sh.600519 → ('600519', 'SH')。"""
    prefix, symbol = code.split(".", 1)
    return symbol, prefix.upper()


# ======================== 灌库 ========================
def upsert_to_db(conn, df: pd.DataFrame) -> int:
    """df → INSERT OR REPLACE daily_baostock_full (19 列)。返回写入行数。"""
    if df.empty:
        return 0
    work = df.copy()
    work[["symbol", "exchange"]] = work["code"].apply(
        lambda x: pd.Series(parse_baostock_code(x))
    )
    cols = ["symbol", "exchange", "date", "open", "high", "low", "close",
            "preclose", "volume", "amount", "adjustflag", "turn", "tradestatus",
            "pctChg", "peTTM", "psTTM", "pcfNcfTTM", "pbMRQ", "isST"]
    work = work[cols]
    work["adjustflag"] = work["adjustflag"].astype(str).replace({"nan": None, "None": None})
    # NaN → None (sqlite 写 NULL)
    work = work.where(pd.notnull(work), None)
    conn.executemany(
        "INSERT OR REPLACE INTO daily_baostock_full VALUES "
        "(?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)",
        work.itertuples(index=False, name=None),
    )
    return len(work)


# ======================== 主流程 ========================
def parse_args():
    p = argparse.ArgumentParser(description="baostock 日增量 → daily_baostock_full (VPS 直跑)")
    p.add_argument("--limit", type=int, default=0, help="限制处理股数 (测试)")
    p.add_argument("--dry-run", action="store_true", help="只下不灌库")
    return p.parse_args()


def main():
    global QUERY_COUNT
    args = parse_args()
    today = datetime.date.today()
    end_date = today.strftime("%Y-%m-%d")
    start_date = (today - datetime.timedelta(days=LOOKBACK_DAYS)).strftime("%Y-%m-%d")

    logger.info("=" * 60)
    logger.info("baostock 日增量 → daily_baostock_full (VPS 直跑)")
    logger.info("  DB: %s", DB_PATH)
    logger.info("  窗口: %s ~ %s (LOOKBACK=%d)", start_date, end_date, LOOKBACK_DAYS)
    logger.info("  staging: %s", STAGING)
    logger.info("  dry_run=%s  limit=%s", args.dry_run, args.limit or "无")
    logger.info("=" * 60)

    # 1. login 探针 — 黑名单 graceful skip
    if not login_once():
        logger.error("[SKIP] baostock 登录失败 (疑似黑名单/冷却), graceful exit 2, 不硬跑")
        sys.exit(2)
    logger.info("baostock 登录成功 (IP 未被封)")

    # 2. 全 A 列表
    try:
        all_codes = fetch_all_stocks()
    except Exception as e:
        logger.error("[FATAL] 取股票列表失败: %s", e)
        try:
            bs.logout()
        except Exception:
            pass
        sys.exit(1)
    logger.info("全 A (含退市): %d 只", len(all_codes))
    if args.limit > 0:
        all_codes = all_codes[: args.limit]
        logger.info("limit=%d 截断", args.limit)

    # 3. 连 DB (WAL + busy_timeout)
    import sqlite3
    conn = sqlite3.connect(DB_PATH, timeout=60)
    conn.execute("PRAGMA busy_timeout = 60000")
    conn.execute("PRAGMA journal_mode = WAL")
    rows_before = conn.execute("SELECT COUNT(*) FROM daily_baostock_full").fetchone()[0]
    logger.info("灌库前 baseline rows=%d", rows_before)

    # 4. 主循环
    STAGING.mkdir(parents=True, exist_ok=True)
    stats = {"ok": 0, "empty": 0, "failed": 0, "rows_db": 0, "rows_staged": 0}
    limit_reached = False
    sample_latest = []
    t0 = time.time()

    if not args.dry_run:
        conn.execute("BEGIN")
    try:
        for i, (code, exc) in enumerate(all_codes):
            if QUERY_COUNT >= DAILY_LIMIT:
                logger.warning("query 计数 %d 达防线 %d, graceful stop", QUERY_COUNT, DAILY_LIMIT)
                limit_reached = True
                break
            bs_code = f"{'sh' if exc == 'SH' else 'sz'}.{code}"
            df = None
            for attempt in range(BS_MAX_RETRIES):
                try:
                    df = fetch_one_increment(bs_code, start_date, end_date)
                    break
                except (socket.timeout, TimeoutError, OSError, Exception) as e:
                    logger.warning("%s 重试 %d/%d: %s", code, attempt + 1, BS_MAX_RETRIES, e)
                    if not relogin():
                        df = None
                        break
            if df is None:
                stats["failed"] += 1
            elif df.empty:
                stats["empty"] += 1  # 退市/停牌/区间外, 合法空
            else:
                # staging (审计)
                try:
                    df.to_parquet(STAGING / f"{code}.{exc}_incr.parquet", index=False)
                    stats["rows_staged"] += len(df)
                except Exception as e:
                    logger.debug("staging 写失败 %s: %s", code, e)
                # 灌库
                if not args.dry_run:
                    n = upsert_to_db(conn, df)
                    stats["rows_db"] += n
                stats["ok"] += 1
                if len(sample_latest) < 5:
                    sample_latest.append((code, exc, len(df), df["date"].max()))
            if (i + 1) % 500 == 0:
                logger.info("进度 %d/%d ok=%d empty=%d failed=%d query=%d (%.0fs)",
                            i + 1, len(all_codes), stats["ok"], stats["empty"],
                            stats["failed"], QUERY_COUNT, time.time() - t0)
            if i < len(all_codes) - 1:
                time.sleep(BS_INTERVAL)
        if args.dry_run:
            logger.info("[DRY-RUN] 不灌库")
        else:
            conn.execute("COMMIT")
    except Exception as e:
        if not args.dry_run:
            conn.execute("ROLLBACK")
        logger.error("[FATAL] 事务回滚: %s", e)
        try:
            bs.logout()
        except Exception:
            pass
        sys.exit(1)

    # 5. 统计
    rows_after = conn.execute("SELECT COUNT(*) FROM daily_baostock_full").fetchone()[0]
    conn.close()
    try:
        bs.logout()
    except Exception:
        pass

    logger.info("=" * 60)
    logger.info("[RESULT] dry_run=%s  耗时 %.0f 秒", args.dry_run, time.time() - t0)
    logger.info("  ok=%d empty=%d failed=%d  limit_reached=%s",
                stats["ok"], stats["empty"], stats["failed"], limit_reached)
    logger.info("  staging 行=%d  灌库行=%d  query 总计=%d",
                stats["rows_staged"], stats["rows_db"], QUERY_COUNT)
    logger.info("  daily_baostock_full: %d → %d (+%d)", rows_before, rows_after, rows_after - rows_before)
    logger.info("  样本 latest: %s", sample_latest)
    logger.info("=" * 60)
    sys.exit(3 if limit_reached else 0)


if __name__ == "__main__":
    main()