"""全量回补历史复权因子（adj_factor 表）。

用法（在 backend 目录下）：
    uv run python scripts/backfill_adj_factor.py                # 从 candles 最早日期回补到今天
    uv run python scripts/backfill_adj_factor.py --start 20180101
    uv run python scripts/backfill_adj_factor.py --force        # 已有日期也重拉

- 按交易日逐日拉取全市场因子（pro.adj_factor(trade_date=...)），幂等可断点续跑；
- 交易日取自本地 trade_calendar（缓存不到的区间自动刷新一次日历）；
- Tushare 每分钟限频由 _call_retry 自动等待 62s 重试。
"""
from __future__ import annotations

import argparse
import asyncio
import sys
import time
from datetime import datetime
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from sqlalchemy import delete, func, insert, select

from app.db import async_session
from app.models import AdjFactor, Candle, TradeCalendar
from app.screener.market_sync import _call_retry, _get_pro, _norm_date, _parse_d

_INTERVAL_MSG = 20  # 每完成 N 个交易日打印一次进度


async def _calendar_dates(start: str, end: str) -> list[str]:
    """[start, end] 交易日（升序）。本地日历覆盖不足时直接拉宽范围日历并回写缓存。"""
    async with async_session() as session:
        all_cached = set((await session.execute(select(TradeCalendar.trade_date))).scalars().all())
    cached = sorted(d for d in all_cached if start <= d <= end)
    if cached and min(cached) <= start:
        return cached

    # 覆盖不到起点：按需拉宽范围日历（trade_cal 低积分限频 1 次/小时，失败沿用缓存）
    pro = _get_pro()
    try:
        cal = await asyncio.to_thread(
            _call_retry, pro.trade_cal, exchange="SSE", start_date=start, end_date=end, is_open="1"
        )
        dates = sorted(cal["cal_date"].tolist())
    except Exception as e:  # noqa: BLE001
        if not cached:
            raise
        print(f"交易日历拉取受限（{str(e)[:100]}），沿用本地缓存")
        return cached
    fresh = [d for d in dates if d not in all_cached]
    if fresh:
        async with async_session() as session:
            await session.execute(insert(TradeCalendar), [{"trade_date": d} for d in fresh])
            await session.commit()
    return dates


async def _existing_dates() -> set[str]:
    async with async_session() as session:
        res = await session.execute(select(func.distinct(AdjFactor.trade_date)))
        return {_norm_date(r[0]) for r in res}


async def main(start: str, end: str, force: bool) -> None:
    # 默认起点：candles 最早日线（因子只需覆盖有 K 线的区间）
    if start is None:
        async with async_session() as session:
            first = await session.scalar(select(func.min(Candle.ts)).where(Candle.timeframe == "1d"))
        start = first.strftime("%Y%m%d") if first else "20050101"
    if end is None:
        end = datetime.now().strftime("%Y%m%d")

    dates = await _calendar_dates(start, end)
    have = set() if force else await _existing_dates()
    todo = [d for d in dates if d not in have]
    print(f"区间 {start}~{end} 共 {len(dates)} 个交易日，待回补 {len(todo)} 个（已有 {len(dates) - len(todo)}）")
    if not todo:
        return

    pro = _get_pro()
    done = 0
    for d in todo:
        time.sleep(0.15)  # 轻微控频；分钟级限频由 _call_retry 自动等待重试
        df = None
        for attempt in range(5):  # 网络抖动（超时/断连）也重试，_call_retry 只兜限频
            try:
                df = _call_retry(pro.adj_factor, trade_date=d)  # noqa: 线性脚本直接同步调用
                break
            except Exception as e:  # noqa: BLE001
                wait = min(30 * (attempt + 1), 120)
                print(f"  {d} 拉取异常（{str(e)[:80]}），{wait}s 后重试 {attempt + 1}/5")
                time.sleep(wait)
        if df is None:
            print(f"  {d} 连续 5 次失败，跳过（断点续跑可补）")
            continue
        if df is None or df.empty:
            print(f"  {d} 无数据（非交易日或未生成），跳过")
            continue
        rows = [
            {"trade_date": _parse_d(d), "ts_code": r["ts_code"], "adj_factor": float(r["adj_factor"])}
            for _, r in df.iterrows()
        ]
        async with async_session() as session:
            dt = _parse_d(d)
            await session.execute(delete(AdjFactor).where(AdjFactor.trade_date == dt))
            await session.execute(insert(AdjFactor), rows)
            await session.commit()
        done += 1
        if done % _INTERVAL_MSG == 0 or done == len(todo):
            print(f"  进度 {done}/{len(todo)}（{d}，+{len(rows)} 行）")
    print(f"回补完成：{done} 个交易日")


if __name__ == "__main__":
    ap = argparse.ArgumentParser(description="全量回补历史复权因子")
    ap.add_argument("--start", default=None, help="YYYYMMDD，默认 candles 最早日期")
    ap.add_argument("--end", default=None, help="YYYYMMDD，默认今天")
    ap.add_argument("--force", action="store_true", help="已有日期也重拉")
    a = ap.parse_args()
    asyncio.run(main(a.start, a.end, a.force))