#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""灌后校验：parquet 日线 ↔ quant_trading.db dbbardata 三项一致性检查。

独立可跑。读两个环境变量：
  VNPY_DB_PATH  默认 C:\\sanguo_vnpy_v2\\data\\quant_trading.db
  DAILY_DIR     默认 C:\\sanguo_vnpy_v2\\data\\raw

三项校验：
  A 最新日期相等: max(date) over raw parquet == max(datetime) where interval='d' from DB
  B 抽查 3 只股 OHLCV 逐值相等: 600519(SSE)/000001(SZSE)/000858(SZSE),
    各 parquet 最新 5 行 vs DB datetime 倒序 limit 5, date 对齐逐值比对 (abs<1e-4)
  C 当天数据合理性: DB 中 date==parquet_max 且 interval='d' 的所有行,
    无 close_price<=0、无 high_price<low_price、close_price 非 null

退出码: PASS=0 / FAIL=1。纯 stdlib + pandas (VPS 已装)。

Windows GBK 终端需 `python -X utf8 validate_import.py` 调用。
"""
import os
import sys
import sqlite3
from pathlib import Path

import pandas as pd

DB_PATH = os.environ.get('VNPY_DB_PATH', r'C:\sanguo_vnpy_v2\data\quant_trading.db')
DAILY_DIR = os.environ.get('DAILY_DIR', r'C:\sanguo_vnpy_v2\data\raw')

# (symbol_code, exchange, filename_prefix) 三只蓝筹股做抽查
SAMPLES = [
    ('600519', 'SSE', 'sh'),
    ('000001', 'SZSE', 'sz'),
    ('000858', 'SZSE', 'sz'),
]


def norm_date_str(s):
    """统一成 'YYYY-MM-DD' 字符串。接受 Timestamp / datetime / str。"""
    ts = pd.Timestamp(s)
    return ts.strftime('%Y-%m-%d')


def find_parquet_max_date():
    """扫最近年份目录下所有 raw parquet，取最大 date。
    只扫最新有数据的年份目录（~5000 文件 × 仅读 date 列 ≈ 10s）。
    返回 (max_date_str or None, n_files_scanned)。"""
    max_date = None
    n_files = 0
    year_dirs = sorted(
        [d for d in Path(DAILY_DIR).glob('*') if d.is_dir() and d.name.isdigit()],
        key=lambda d: d.name,
        reverse=True,
    )
    for year_dir in year_dirs[:2]:  # 最多看最近 2 个年份目录
        local_max = None
        for f in year_dir.glob('*.parquet'):
            n_files += 1
            try:
                df = pd.read_parquet(f, columns=['date'])
                if df.empty:
                    continue
                cur = df['date'].max()
                if local_max is None or cur > local_max:
                    local_max = cur
            except Exception:
                continue
        if local_max is not None:
            max_date = local_max
            break  # 当前年份有数据就不再看上一年
    return (norm_date_str(max_date) if max_date is not None else None, n_files)


def check_a(parquet_max, db_max):
    """A: 最新日期相等"""
    p = parquet_max
    d = db_max[:10] if db_max else None
    ok = (p is not None) and (d is not None) and (p == d)
    return ok, {'parquet_max': p, 'db_max': d}


def check_b(conn):
    """B: 3 只股 OHLCV 逐值比对"""
    detail = {}
    all_ok = True
    for code, exchange, prefix in SAMPLES:
        # 从最新年份目录往前找 parquet
        year_dirs = sorted(
            [d for d in Path(DAILY_DIR).glob('*') if d.is_dir() and d.name.isdigit()],
            key=lambda d: d.name,
            reverse=True,
        )
        parquet_path = None
        for yd in year_dirs:
            p = yd / f'{prefix}{code}_daily.parquet'
            if p.exists():
                parquet_path = p
                break
        if parquet_path is None:
            detail[code] = {'status': 'MISSING_PARQUET'}
            all_ok = False
            continue
        try:
            pdf = pd.read_parquet(parquet_path).sort_values('date').tail(5).reset_index(drop=True)
        except Exception as e:
            detail[code] = {'status': f'PARQUET_READ_ERR: {e}'}
            all_ok = False
            continue
        try:
            rows = conn.execute(
                "SELECT datetime, open_price, high_price, low_price, close_price, volume "
                "FROM dbbardata WHERE symbol=? AND exchange=? AND interval='d' "
                "ORDER BY datetime DESC LIMIT 5",
                (code, exchange),
            ).fetchall()
        except Exception as e:
            detail[code] = {'status': f'DB_QUERY_ERR: {e}'}
            all_ok = False
            continue
        if not rows:
            detail[code] = {'status': 'DB_EMPTY'}
            all_ok = False
            continue
        db_pdf = pd.DataFrame(rows, columns=['datetime', 'open', 'high', 'low', 'close', 'volume'])
        db_pdf = db_pdf.sort_values('datetime').reset_index(drop=True)
        db_pdf['date'] = db_pdf['datetime'].str[:10]
        pdf = pdf.copy()
        pdf['date'] = pdf['date'].astype(str).str[:10]

        # 取共同尾部 n 行
        n = min(len(pdf), len(db_pdf))
        if n == 0:
            detail[code] = {'status': 'EMPTY_AFTER_ALIGN'}
            all_ok = False
            continue
        p_tail = pdf.tail(n).reset_index(drop=True)
        d_tail = db_pdf.tail(n).reset_index(drop=True)

        dates_match = (p_tail['date'].values == d_tail['date'].values).all()
        cols = ['open', 'high', 'low', 'close', 'volume']
        vals_match = True
        bad_col = None
        for c in cols:
            try:
                pv = p_tail[c].astype(float).values
                dv = d_tail[c].astype(float).values
                max_diff = float((abs(pv - dv)).max())
                if max_diff > 1e-4:
                    vals_match = False
                    bad_col = f'{c}(max_diff={max_diff})'
                    break
            except Exception as e:
                vals_match = False
                bad_col = f'{c}: {e}'
                break

        ok = bool(dates_match and vals_match)
        if not ok:
            all_ok = False
        detail[code] = {
            'status': 'PASS' if ok else 'FAIL',
            'n_rows': n,
            'dates_match': bool(dates_match),
            'vals_match': bool(vals_match),
            'bad_col': bad_col,
            'latest_date': p_tail['date'].iloc[-1] if len(p_tail) else None,
            'latest_close': float(p_tail['close'].iloc[-1]) if len(p_tail) else None,
        }
    return all_ok, detail


def check_c(conn, parquet_max_str):
    """C: 当天数据合理性（DB 中 date==parquet_max 的所有 interval='d' 行）"""
    if not parquet_max_str:
        return False, {'reason': 'no parquet_max'}
    pattern = parquet_max_str + '%'
    try:
        total = conn.execute(
            "SELECT COUNT(*) FROM dbbardata WHERE interval='d' AND datetime LIKE ?",
            (pattern,),
        ).fetchone()[0]
    except Exception as e:
        return False, {'reason': f'DB_QUERY_ERR: {e}'}
    if total == 0:
        return False, {'reason': f'no rows for {parquet_max_str}'}
    try:
        bad_close = conn.execute(
            "SELECT COUNT(*) FROM dbbardata WHERE interval='d' AND datetime LIKE ? "
            "AND (close_price IS NULL OR close_price <= 0)",
            (pattern,),
        ).fetchone()[0]
        bad_hl = conn.execute(
            "SELECT COUNT(*) FROM dbbardata WHERE interval='d' AND datetime LIKE ? "
            "AND high_price < low_price",
            (pattern,),
        ).fetchone()[0]
    except Exception as e:
        return False, {'reason': f'DB_QUERY_ERR: {e}'}
    ok = (bad_close == 0) and (bad_hl == 0)
    return ok, {
        'date': parquet_max_str,
        'total_rows': total,
        'bad_close': bad_close,
        'bad_high_low': bad_hl,
    }


def main():
    print(f'[VALIDATE] DB={DB_PATH}')
    print(f'[VALIDATE] DAILY_DIR={DAILY_DIR}')

    if not os.path.exists(DB_PATH):
        print(f'[VALIDATE][FAIL] DB not found: {DB_PATH}')
        sys.exit(1)
    if not os.path.exists(DAILY_DIR):
        print(f'[VALIDATE][FAIL] DAILY_DIR not found: {DAILY_DIR}')
        sys.exit(1)

    parquet_max, n_files = find_parquet_max_date()
    if parquet_max is None:
        print(f'[VALIDATE][FAIL] no parquet under {DAILY_DIR} (scanned {n_files} files)')
        sys.exit(1)
    print(f'[VALIDATE] scanned {n_files} parquet files, parquet_max={parquet_max}')

    conn = sqlite3.connect(DB_PATH)

    try:
        row = conn.execute("SELECT MAX(datetime) FROM dbbardata WHERE interval='d'").fetchone()
        db_max = row[0] if row else None
    except Exception as e:
        print(f'[VALIDATE][FAIL] DB query max(datetime) err: {e}')
        conn.close()
        sys.exit(1)
    print(f'[VALIDATE] db_max={db_max}')

    # A
    a_ok, a_detail = check_a(parquet_max, db_max)
    print(f'[VALIDATE][A] {"PASS" if a_ok else "FAIL"} parquet_max={a_detail["parquet_max"]} db_max={a_detail["db_max"]}')

    # B
    b_ok, b_detail = check_b(conn)
    print(f'[VALIDATE][B] {"PASS" if b_ok else "FAIL"}')
    for code, info in b_detail.items():
        if info.get('status') == 'PASS':
            print(f'  {code}: PASS n={info["n_rows"]} latest={info["latest_date"]} close={info["latest_close"]}')
        else:
            print(f'  {code}: {info}')

    # C
    c_ok, c_detail = check_c(conn, parquet_max)
    print(f'[VALIDATE][C] {"PASS" if c_ok else "FAIL"} {c_detail}')

    conn.close()

    overall = a_ok and b_ok and c_ok
    print(f'[VALIDATE] {"PASS" if overall else "FAIL"} '
          f'(A={"PASS" if a_ok else "FAIL"} B={"PASS" if b_ok else "FAIL"} C={"PASS" if c_ok else "FAIL"})')
    sys.exit(0 if overall else 1)


if __name__ == '__main__':
    main()