#!/usr/bin/env python3
"""判断今天是否为 A 股交易日。

依据：上证指数(secid 1.000001)最近一根日 K 线的日期是否等于今天（北京时间）。
- 相等  -> exit 0（交易日，继续跑流水线）
- 不等  -> exit 1（休市 / 当日数据尚未更新，跳过）
"""
import sys
from datetime import datetime, timezone, timedelta

try:
    import requests
except ImportError:
    print("ERROR: 缺少 requests，请先执行 pip install requests")
    sys.exit(1)

try:
    from zoneinfo import ZoneInfo
    CN_TZ = ZoneInfo("Asia/Shanghai")
except Exception:
    CN_TZ = timezone(timedelta(hours=8))

URL = "https://push2his.eastmoney.com/api/qt/stock/kline/get"
PARAMS = {
    "secid": "1.000001",  # 上证指数
    "fields1": "f1,f2,f3,f4,f5,f6",
    "fields2": "f51,f52,f53,f54,f55,f56,f57,f58",
    "klt": "101",          # 日 K
    "fqt": "1",
    "beg": "0",
    "end": "20500101",
    "lmt": "5",
}


def main():
    today = datetime.now(CN_TZ).strftime("%Y-%m-%d")
    try:
        r = requests.get(URL, params=PARAMS, headers={"User-Agent": "Mozilla/5.0"}, timeout=10)
        r.raise_for_status()
        klines = (r.json().get("data") or {}).get("klines") or []
        if not klines:
            print("ERROR: 未取到上证指数 K 线数据，无法判断交易日")
            sys.exit(1)
        last_date = klines[-1].split(",")[0]
        if last_date == today:
            print(f"交易日确认：上证指数最近 K 线 {last_date} == 今天 {today}，是 A 股交易日")
            sys.exit(0)
        print(f"非交易日：上证指数最近 K 线 {last_date} != 今天 {today}，跳过本次扫描")
        sys.exit(1)
    except Exception as e:
        print(f"ERROR: 交易日判断异常：{e}")
        sys.exit(1)


if __name__ == "__main__":
    main()