#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
修复 daily 表历史覆盖度缺口（2012-2013 沪市+创业板缺失）

背景：daily 表 2012-2013 年只有深市主板/中小板（2012 年 1383 只、2013 年 677 只），
     缺全部沪市 + 创业板。原因：当年全量导入沪市请求失败/中断，但深市成功，
     导致断点续传的"日期级"缺失检测（NOT EXISTS trade_date）认为该日已有数据，
     永不补拉。

修复策略（覆盖度级校验）：
1. 对指定日期范围，对比 daily 与 daily_basic 的当日去重股票数
   （daily_basic 同期数据完整，作为覆盖度基准）
2. daily 当日股票数 < daily_basic 当日股票数 × (1 - tolerance) 的日期 → 判定为"部分缺失"
3. 部分缺失的日期按全市场重新拉取（Tushare daily 按 trade_date 返回全市场）
4. 用 INSERT ON CONFLICT DO NOTHING 幂等写入，可重复执行

用法：
    python3 repair_daily_backfill.py [--start 2012-01-01] [--end 2013-12-31] [--tolerance 0.05] [--dry-run]
"""
import os
import sys
import time
import argparse
import logging
from datetime import datetime

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)

# 显式加载 .env（config.py 会加载，但确保顺序正确）
try:
    from dotenv import load_dotenv
    env_path = os.path.join(HERE, ".env")
    if os.path.exists(env_path):
        load_dotenv(env_path, override=False)
except ImportError:
    pass

import importer
from importer import (
    get_pg_connection, get_ts_pro, batch_insert,
    fetch_with_retry, _normalize_daily_df,
)

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    handlers=[
        logging.FileHandler(os.path.join(HERE, "repair_daily_backfill.log"), encoding="utf-8"),
        logging.StreamHandler(),
    ],
)
logger = logging.getLogger("repair_daily_backfill")


def get_coverage_ratio(conn, trade_date: str) -> tuple:
    """返回 (daily_stocks, daily_basic_stocks)。无参考数据时 daily_basic_stocks=None"""
    cur = conn.cursor()
    try:
        cur.execute("SELECT COUNT(DISTINCT ts_code) FROM daily WHERE trade_date=%s", (trade_date,))
        d_cnt = cur.fetchone()[0]
        cur.execute("SELECT COUNT(DISTINCT ts_code) FROM daily_basic WHERE trade_date=%s", (trade_date,))
        db_cnt = cur.fetchone()[0]
        return d_cnt, db_cnt
    finally:
        cur.close()


def find_partial_dates(conn, start_date: str, end_date: str, tolerance: float) -> list:
    """
    找出 daily 覆盖度不足的交易日（单条 SQL 聚合，避免逐日查询）。
    返回 [(trade_date, daily_cnt, daily_basic_cnt), ...]
    判据：daily_basic 有数据且 daily 股票数 < daily_basic × (1 - tolerance)
    """
    cur = conn.cursor()
    try:
        cur.execute(
            """
            SELECT db.trade_date,
                   COALESCE(d.cnt, 0) AS daily_cnt,
                   db.cnt AS db_cnt
            FROM (SELECT trade_date, COUNT(DISTINCT ts_code) AS cnt
                  FROM daily_basic
                  WHERE trade_date BETWEEN %s AND %s
                  GROUP BY trade_date) db
            LEFT JOIN (SELECT trade_date, COUNT(DISTINCT ts_code) AS cnt
                       FROM daily
                       WHERE trade_date BETWEEN %s AND %s
                       GROUP BY trade_date) d
              ON d.trade_date = db.trade_date
            WHERE COALESCE(d.cnt, 0) < db.cnt * (1 - %s)
            ORDER BY db.trade_date
            """,
            (start_date, end_date, start_date, end_date, tolerance),
        )
        partial = [(r[0].strftime("%Y-%m-%d"), r[1], r[2]) for r in cur.fetchall()]
    finally:
        cur.close()
    return partial


def backfill_date(pro, conn, td_str: str) -> bool:
    """拉取单个交易日全市场 daily 数据并写入。成功返回 True"""
    td_compact = td_str.replace("-", "")

    def fetch():
        return pro.daily(trade_date=td_compact)

    df = fetch_with_retry(fetch, max_retries=4, delay=3)
    if df is None or df.empty:
        logger.warning(f"  [{td_str}] 返回空数据，跳过")
        return False

    df = _normalize_daily_df(df)
    # 幂等写入：已存在的行跳过（ON CONFLICT DO NOTHING）
    batch_insert("daily", df, conn, ["ts_code", "trade_date"])
    return True


def main():
    parser = argparse.ArgumentParser(description="修复 daily 表历史覆盖度缺口")
    parser.add_argument("--start", default="2012-01-01", help="起始日期 YYYY-MM-DD")
    parser.add_argument("--end", default="2013-12-31", help="结束日期 YYYY-MM-DD")
    parser.add_argument("--tolerance", type=float, default=0.05,
                        help="覆盖度容差，默认 0.05（daily 少于 daily_basic 的 95% 即判定缺失）")
    parser.add_argument("--dry-run", action="store_true", help="只扫描不导入")
    args = parser.parse_args()

    conn = get_pg_connection()
    logger.info("=" * 60)
    logger.info(f"[覆盖度扫描] {args.start} ~ {args.end} (容差 {args.tolerance:.0%})")

    partial = find_partial_dates(conn, args.start, args.end, args.tolerance)
    if not partial:
        logger.info("  未发现覆盖度不足的交易日 ✅")
        conn.close()
        return

    logger.info(f"  发现 {len(partial)} 个覆盖度不足的交易日:")
    # 按年份统计
    years = {}
    for td, d_cnt, db_cnt in partial:
        y = td[:4]
        years.setdefault(y, []).append((td, d_cnt, db_cnt))
    for y in sorted(years):
        lst = years[y]
        logger.info(f"    {y} 年: {len(lst)} 个交易日 | 样例 {lst[0][0]}(daily={lst[0][1]}/db={lst[0][2]})")

    if args.dry_run:
        logger.info("[dry-run] 不执行导入，以上为待补拉清单")
        conn.close()
        return

    pro = get_ts_pro()
    success = 0
    fail_list = []
    for i, (td, d_cnt, db_cnt) in enumerate(partial, 1):
        ok = backfill_date(pro, conn, td)
        if ok:
            success += 1
        else:
            fail_list.append(td)
        if i % 20 == 0 or i == len(partial):
            logger.info(f"  进度: {i}/{len(partial)} 成功={success} 失败={len(fail_list)}")
        time.sleep(0.35)  # 限流保护

    conn.close()
    logger.info(f"  补拉完成: 成功 {success}/{len(partial)}")
    if fail_list:
        logger.warning(f"  失败日期({len(fail_list)}): {fail_list[:20]}...")
        # 写失败清单供重试
        with open(os.path.join(HERE, "repair_failed_dates.txt"), "w") as f:
            f.write("\n".join(fail_list))


if __name__ == "__main__":
    main()