#!/usr/bin/env python3
"""验证 staging 增量质量（task: 截断 bug 修复的安全闸门）。

在 merge_increment 之前跑：校验 staging 的 raw/qfq 增量数据是否合格，
不合格则 run_daily_update.sh 不合并、不推 NAS（staging 留存供排查）。

**复用 scripts/data_platform/validator.py 的 DataValidator（七条 fatal）**：
  D1 价格>0 / D2 OHLC 一致 / D3 volume≥0 / D6 日期不重复 / D7 非未来日期 / ...

阈值（Main Agent 已定，硬编码常量）：
  MIN_SUCCESS_RATE = 0.95      成功率（分母扣除北交所已知不支持码）
  MIN_FRESH_RATE = 0.95        最大日期 >= 最近交易日的 symbol 占比
  KNOWN_UNSUPPORTED_PREFIX     新浪 stock_zh_a_daily 不支持的北交所码（从分母扣）

用法：
  python3 verify_increment.py --staging data_cache/daily_update/raw \\
      --stock-list data_cache/stock_info/stock_basic_info_raw_*.csv \\
      --start 2026-07-01

退出码：passed=0，否则 1。输出 JSON 到 stdout。
"""
from __future__ import annotations

import argparse
import json
import logging
import os
import sys
from datetime import datetime, timedelta

import pandas as pd

# 同目录 import（脚本运行目录 = scripts/data_platform/）
_SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
if _SCRIPT_DIR not in sys.path:
    sys.path.insert(0, _SCRIPT_DIR)

from validator import DataValidator  # noqa: E402

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

# 阈值（Main Agent 已定）
MIN_SUCCESS_RATE = 0.95
MIN_FRESH_RATE = 0.95
KNOWN_UNSUPPORTED_PREFIX = ("920", "921", "83", "87")
MAX_FATAL_SAMPLES = 5  # 输出里只放前 N 个失败样本（避免输出爆掉）


def _latest_available_trading_day(now: datetime | None = None) -> str:
    """A 股收盘感知：返回 now 时点最近的可获取交易日（YYYY-MM-DD）。

    catch-up 跨夜场景：下载发生在昨晚、verify 跑在今天盘前，
    今天数据尚未生成（15:00 才收盘）→ 目标应是昨天而非今天。

    - 工作日 >= 15:00 → 今天（已收盘，数据可获取）
    - 工作日 < 15:00 → 上一交易日（今天未收盘，回退到工作日）
    - 周末 → 上周五

    注意：与 raw_redownload.latest_trading_day 不同——那个只回退周末、
    用于 is_fresh 决定是否重下；本函数额外感知收盘时间，用于 verify 闸门。
    """
    if now is None:
        now = datetime.now()
    d = now
    if d.weekday() >= 5:  # 周末 → 回退到周五
        d = d - timedelta(days=(d.weekday() - 4))
    elif d.hour < 15:  # 工作日盘前 → 上一交易日
        d = d - timedelta(days=1)
        while d.weekday() >= 5:  # 回退周末
            d = d - timedelta(days=1)
    return d.strftime("%Y-%m-%d")


def symbol_from_filename(fname: str) -> str:
    """`sh600000_daily.parquet` → `600000`。"""
    base = fname
    if base.endswith("_daily.parquet"):
        base = base[: -len("_daily.parquet")]
    if base[:2] in ("sh", "sz", "bj"):
        base = base[2:]
    return base


def verify(staging_root: str, start: str, _now: datetime | None = None) -> dict:
    """校验 staging 下 start_year 目录的所有 parquet。

    返回 dict: {passed, success_rate, fresh_rate, total, unsupported_skipped,
                success, fresh, failed_symbols, fatal_samples, latest_trading_day}。

    _now 仅用于测试注入固定时间；生产留空取 datetime.now()。
    """
    start_year = int(start[:4])
    year_dir = os.path.join(staging_root, str(start_year))
    if not os.path.isdir(year_dir):
        raise FileNotFoundError(f"staging year dir 不存在: {year_dir}")

    now = _now or datetime.now()
    latest = _latest_available_trading_day(now)
    latest_ts = pd.Timestamp(latest)
    log.info("最近可获取交易日=%s（now=%s）", latest, now.strftime("%Y-%m-%d %H:%M"))

    validator = DataValidator()

    total = 0
    unsupported_skipped = 0
    success = 0
    fresh = 0
    failed_symbols: list[str] = []
    fatal_samples: list[dict] = []

    files = sorted(f for f in os.listdir(year_dir) if f.endswith("_daily.parquet"))
    for fname in files:
        total += 1
        code = symbol_from_filename(fname)
        # 北交所码新浪不支持（log 大量 920xxx KeyError:'date' 确认）→ 从分母扣
        if code.startswith(KNOWN_UNSUPPORTED_PREFIX):
            unsupported_skipped += 1
            continue

        fpath = os.path.join(year_dir, fname)
        try:
            df = pd.read_parquet(fpath)
        except Exception as e:  # noqa: BLE001  损坏文件记 fail
            failed_symbols.append(code)
            if len(fatal_samples) < MAX_FATAL_SAMPLES:
                fatal_samples.append({"symbol": code, "errors": [f"read_parquet: {type(e).__name__}: {str(e)[:80]}"]})
            continue

        result = validator.validate(df, "daily")
        if not result.passed:
            failed_symbols.append(code)
            if len(fatal_samples) < MAX_FATAL_SAMPLES:
                fatal_samples.append({"symbol": code, "errors": result.fatal_errors[:3]})
            continue

        success += 1
        # 新鲜度：最大日期 >= 最近交易日（staging 是增量，多数会 == latest）
        try:
            maxd = pd.Timestamp(df["date"].max())
        except Exception:  # noqa: BLE001  无 date 列记为不 fresh（但已过校验，理论上不会）
            maxd = pd.Timestamp("1970-01-01")
        if maxd >= latest_ts:
            fresh += 1

    denom = total - unsupported_skipped
    success_rate = (success / denom) if denom > 0 else 0.0
    fresh_rate = (fresh / success) if success > 0 else 0.0
    passed = (success_rate >= MIN_SUCCESS_RATE) and (fresh_rate >= MIN_FRESH_RATE) and (denom > 0)

    return {
        "passed": passed,
        "success_rate": round(success_rate, 4),
        "fresh_rate": round(fresh_rate, 4),
        "total": total,
        "unsupported_skipped": unsupported_skipped,
        "success": success,
        "fresh": fresh,
        "denom": denom,
        "failed_count": len(failed_symbols),
        "failed_symbols": failed_symbols,
        "fatal_samples": fatal_samples,
        "latest_trading_day": latest,
        "thresholds": {
            "MIN_SUCCESS_RATE": MIN_SUCCESS_RATE,
            "MIN_FRESH_RATE": MIN_FRESH_RATE,
            "KNOWN_UNSUPPORTED_PREFIX": list(KNOWN_UNSUPPORTED_PREFIX),
        },
    }


def main():
    ap = argparse.ArgumentParser(description="验证 staging 增量质量")
    ap.add_argument("--staging", required=True, help="staging 根（如 data_cache/daily_update/raw）")
    ap.add_argument("--stock-list", default=None, help="stock_basic_info csv（保留接口，本版未强用）")
    ap.add_argument("--start", required=True, help="增量起点 YYYY-MM-DD（决定看哪个 year 目录）")
    ap.add_argument("--summary-json", default=None, help="把结果写到该 JSON 文件")
    args = ap.parse_args()

    try:
        result = verify(args.staging, args.start)
    except FileNotFoundError as e:
        log.error("%s", e)
        sys.exit(2)

    # 输出
    print(json.dumps(result, ensure_ascii=False, indent=2, default=str))

    if args.summary_json:
        with open(args.summary_json, "w") as f:
            json.dump(result, f, ensure_ascii=False, indent=2, default=str)

    verdict = "PASSED" if result["passed"] else "FAILED"
    log.info("=== verify %s: success_rate=%.2f fresh_rate=%.2f total=%d skipped=%d failed=%d ===",
             verdict, result["success_rate"], result["fresh_rate"],
             result["total"], result["unsupported_skipped"], result["failed_count"])

    sys.exit(0 if result["passed"] else 1)


if __name__ == "__main__":
    main()