#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""bs_eod.py — sanguo-bs-eod (方案A schtask 18:05): baostock 个股 EOD 增量。

每日收盘后跑(VPS, baostock 日终更新就绪):
- 个股日线(含退市, LOOKBACK 7 天) -> dbbardata('d') INSERT OR REPLACE (治幸存者偏差)
- 个股 15min(LOOKBACK 7) -> dbbardata('15m')
- 日线 pe/pb/turn/pctChg/isST -> data/valuation_baostock/<year>.parquet 追加
- DAILY_LIMIT=48000 单进程单登录, sleep 0.3s, login 探针 graceful skip

预算: 5537股 × (1日线+1 15min) ≈ 11000 query/天 = 48000 的 23%, 安全。
退出码: 0=完成; 1=致命; 2=黑名单 graceful skip; 3=query 超限 graceful stop
"""
import argparse
import datetime as dt
import logging
import os
import socket
import sys
import time
from pathlib import Path

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
import pandas as pd

from dbbardata_utils import normalize_daily_dt

BASE = Path(r"C:\sanguo_vnpy_v2")
DB = BASE / "data" / "quant_trading.db"
VAL_DIR = BASE / "data" / "valuation_baostock"
LOOKBACK = int(os.environ.get("LOOKBACK_DAYS", "7"))
DAILY_LIMIT = int(os.environ.get("BS_DAILY_LIMIT", "48000"))
BS_INTERVAL = 0.3
QUERY_COUNT = 0
EXC_MAP = {"sh": "SSE", "sz": "SZSE"}

DAILY_FIELDS = ("date,code,open,high,low,close,volume,amount,turn,"
                "pctChg,peTTM,psTTM,pcfNcfTTM,pbMRQ,isST")
M15_FIELDS = "date,time,code,open,high,low,close,volume,amount"

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


def login_once():
    try:
        lg = bs.login()
        if lg.error_code == "0":
            return True
        log.error("login fail: %s %s", lg.error_code, lg.error_msg)
        return False
    except Exception as e:
        log.error("login exc: %s", e)
        return False


def relogin():
    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_all_stocks():
    """query_stock_basic() 无参 -> 全 A 含退市 (type=1)。返回 [(code, 'sh'/'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()
        if r[idx["type"]] != "1":
            continue
        bc = r[idx["code"]]
        if "." not in bc:
            continue
        prefix, num = bc.split(".", 1)
        if prefix in ("sh", "sz") and len(num) == 6 and num.isdigit():
            out.append((num, prefix))
    return out


def fetch_k(bs_code, fields, freq, start, end):
    global QUERY_COUNT
    QUERY_COUNT += 1
    rs = bs.query_history_k_data_plus(bs_code, fields, start_date=start,
                                      end_date=end, frequency=freq, 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())
    return rows


def upsert_daily(conn, code, prefix, rows):
    """日线 rows -> dbbardata('d') + valuation_baostock 当年 parquet 追加。"""
    if not rows:
        return 0
    df = pd.DataFrame(rows, columns=DAILY_FIELDS.split(","))
    for c in ["open", "high", "low", "close", "volume", "amount",
              "turn", "pctChg", "peTTM", "psTTM", "pcfNcfTTM", "pbMRQ"]:
        df[c] = pd.to_numeric(df[c], errors="coerce")
    exc = EXC_MAP[prefix]
    # OHLCV -> dbbardata('d') — datetime 归一纯日期 (dbbardata 双行根治方案A)
    db = pd.DataFrame({
        "symbol": code, "exchange": exc,
        "datetime": df["date"].astype(str).map(normalize_daily_dt),
        "interval": "d", "volume": df["volume"], "turnover": df["amount"],
        "open_interest": 0.0,
        "open_price": df["open"], "high_price": df["high"],
        "low_price": df["low"], "close_price": df["close"],
    })
    conn.executemany(
        "INSERT OR REPLACE INTO dbbardata "
        "(symbol,exchange,datetime,interval,volume,turnover,open_interest,"
        "open_price,high_price,low_price,close_price) VALUES (?,?,?,?,?,?,?,?,?,?,?)",
        db.itertuples(index=False, name=None))
    # pe/pb -> parquet 追加 (isST->int, 修 pyarrow ArrowTypeError)
    vdf = df[["date", "peTTM", "psTTM", "pcfNcfTTM", "pbMRQ", "turn", "pctChg", "isST"]].copy()
    vdf["isST"] = pd.to_numeric(vdf["isST"], errors="coerce").fillna(0).astype(int)
    vdf.insert(0, "symbol", code)
    vdf.insert(1, "exchange", exc)
    yr = dt.date.today().year
    p = VAL_DIR / f"{yr}.parquet"
    if p.exists():
        try:
            old = pd.read_parquet(p)
            vdf = pd.concat([old, vdf]).drop_duplicates(["symbol", "date"], keep="last")
        except Exception:
            pass
    vdf.sort_values(["symbol", "date"]).to_parquet(p, index=False)
    return len(db)


def _build_15m_dt(date_series, time_series):
    """baostock 15min datetime 拼接: date="2026-07-21"(带 -) + time="20260721094500000"(17 位)。

    从 17 位 time 第 8-12 位提取 HHMM, date 直连(带 -)。
    产出 'YYYY-MM-DD HH:MM:00' (符合 dbbardata GLOB 模式, 不被清理误删)。

    bug 根因(2026-07-25): 原代码假设 date 纯数字 + time[:6] 取年月,
    但 baostock 实测 date 带 -, time[:6]=YYYYMM, 产乱 datetime 致 15min 全市场停 7-17。
    """
    _t = time_series.astype(str)
    return (date_series.astype(str) + " "
            + _t.str.slice(8, 10) + ":" + _t.str.slice(10, 12) + ":00")


def upsert_15m(conn, code, prefix, rows):
    if not rows:
        return 0
    df = pd.DataFrame(rows, columns=M15_FIELDS.split(","))
    for c in ["open", "high", "low", "close", "volume", "amount"]:
        df[c] = pd.to_numeric(df[c], errors="coerce")
    exc = EXC_MAP[prefix]
    dt_col = _build_15m_dt(df["date"], df["time"])
    db = pd.DataFrame({
        "symbol": code, "exchange": exc, "datetime": dt_col,
        "interval": "15m", "volume": df["volume"], "turnover": df["amount"],
        "open_interest": 0.0,
        "open_price": df["open"], "high_price": df["high"],
        "low_price": df["low"], "close_price": df["close"],
    })
    conn.executemany(
        "INSERT OR REPLACE INTO dbbardata "
        "(symbol,exchange,datetime,interval,volume,turnover,open_interest,"
        "open_price,high_price,low_price,close_price) VALUES (?,?,?,?,?,?,?,?,?,?,?)",
        db.itertuples(index=False, name=None))
    return len(db)


def main():
    global QUERY_COUNT
    ap = argparse.ArgumentParser()
    ap.add_argument("--limit", type=int, default=0)
    ap.add_argument("--no-15m", action="store_true")
    ap.add_argument("--no-daily", action="store_true",
                    help="跳日线, 只跑 15min(用于 15min 重灌快)")
    args = ap.parse_args()

    today = dt.date.today()
    end = today.strftime("%Y-%m-%d")
    start = (today - dt.timedelta(days=LOOKBACK)).strftime("%Y-%m-%d")
    log.info("bs_eod start window=%s~%s LOOKBACK=%d limit=%s", start, end, LOOKBACK, args.limit or "无")

    if not login_once():
        log.error("[SKIP] baostock 黑名单/冷却, graceful exit 2")
        sys.exit(2)

    try:
        stocks = fetch_all_stocks()
    except Exception as e:
        log.error("[FATAL] fetch_all: %s", e)
        sys.exit(1)
    log.info("全 A 含退市: %d 只", len(stocks))
    if args.limit:
        stocks = stocks[:args.limit]

    import sqlite3
    VAL_DIR.mkdir(parents=True, exist_ok=True)
    conn = sqlite3.connect(str(DB), timeout=60)
    conn.execute("PRAGMA busy_timeout = 60000")
    conn.execute("PRAGMA journal_mode = WAL")

    stats = {"ok": 0, "empty": 0, "failed": 0, "db_rows": 0}
    limit_reached = False
    t0 = time.time()
    conn.execute("BEGIN")
    try:
        for i, (code, prefix) in enumerate(stocks):
            if QUERY_COUNT >= DAILY_LIMIT:
                log.warning("query %d 达防线 %d, graceful stop", QUERY_COUNT, DAILY_LIMIT)
                limit_reached = True
                break
            bs_code = f"{prefix}.{code}"
            try:
                n1 = 0
                if not args.no_daily:
                    d_rows = fetch_k(bs_code, DAILY_FIELDS, "d", start, end)
                    n1 = upsert_daily(conn, code, prefix, d_rows)
                n2 = 0
                if not args.no_15m:
                    m_rows = fetch_k(bs_code, M15_FIELDS, "15", start, end)
                    n2 = upsert_15m(conn, code, prefix, m_rows)
                stats["db_rows"] += n1 + n2
                if n1 + n2:
                    stats["ok"] += 1
                else:
                    stats["empty"] += 1
            except Exception as e:
                stats["failed"] += 1
                if stats["failed"] <= 5:
                    log.warning("%s err: %s", code, e)
                if not relogin():
                    log.error("%s relogin 失败, 跳过", code)
            if (i + 1) % 500 == 0:
                log.info("进度 %d/%d ok=%d empty=%d failed=%d q=%d (%.0fs)",
                         i + 1, len(stocks), stats["ok"], stats["empty"],
                         stats["failed"], QUERY_COUNT, time.time() - t0)
            if i < len(stocks) - 1:
                time.sleep(BS_INTERVAL)
        conn.execute("COMMIT")
    except Exception as e:
        conn.execute("ROLLBACK")
        log.error("[FATAL] rollback: %s", e)
        sys.exit(1)
    finally:
        conn.close()
        try:
            bs.logout()
        except Exception:
            pass

    log.info("[DONE] ok=%d empty=%d failed=%d db_rows=%d query=%d 耗时%.0fs",
             stats["ok"], stats["empty"], stats["failed"], stats["db_rows"],
             QUERY_COUNT, time.time() - t0)
    sys.exit(3 if limit_reached else 0)


if __name__ == "__main__":
    main()