#!/usr/bin/env python3
"""历史股本回填脚本。

目标：把 2023-01-01 ~ 2026-07-15 期间每只 A 股的股本，按“每周一条”
写入 market_data.share 表。

数据源：
  1. akshare stock_share_change_cninfo（巨潮资讯-公司股本变动）
     覆盖 2023 及更早的历史股本变动记录，单位：万股。
  2. 现有 market_data.share 快照（2026 年雪球 quote_detail 已写入）
     作为 2024 年底之后最新股本的补充。

策略：
  - 每周五作为本周代表日（若周五休市则仍取周五，因为股本变动是事件日，
    不依赖当天是否交易；最后一周如果结束日不是周五，额外补一个结束日）。
  - 对每只股票，把历史股本变动 + 现有快照合并成一条时间线，按日期前向填充，
    得到每个采样日的 total_share / float_share。
  - 最终批量 upsert 到 share 表。
  - 默认使用巨潮“总股本 / 已流通股份”口径；加 --a-share-only 时只取 A 股部分
    （total_share = 人民币普通股，float_share = 人民币普通股 - 流通受限股份），
    与雪球 quote_detail 的 float_shares 口径更一致。

用法：
  .venv/bin/python bin/backfill_share_history.py
  .venv/bin/python bin/backfill_share_history.py --dry-run --limit 50
  .venv/bin/python bin/backfill_share_history.py --max-workers 10 --start 2023-01-01 --end 2026-07-15
  .venv/bin/python bin/backfill_share_history.py --a-share-only
"""
from __future__ import annotations

import argparse
import os
import sys
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import date, datetime, timedelta
from typing import Any, Optional

import pandas as pd
from sqlalchemy import text

# 项目根目录导入
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from app.core.db import ops as db_ops
from app.core.db.orm import get_session
from app.core.datasource.utils import is_a_share_code, to_hermes
from app.core.utils.logging import get_logger

logger = get_logger("backfill_share_history")

# 默认回填区间
DEFAULT_START = "2023-01-01"
DEFAULT_END = "2026-07-15"


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="回填历史股本到 market_data.share")
    parser.add_argument("--start", default=DEFAULT_START, help="开始日期 (YYYY-MM-DD)")
    parser.add_argument("--end", default=DEFAULT_END, help="结束日期 (YYYY-MM-DD)")
    parser.add_argument("--max-workers", type=int, default=5, help="并发 worker 数")
    parser.add_argument("--limit", type=int, default=0, help="仅处理前 N 只股票（测试用）")
    parser.add_argument(
        "--codes",
        default="",
        help="指定股票代码，逗号分隔（如 000001,600519）；空则处理全市场",
    )
    parser.add_argument("--dry-run", action="store_true", help="只统计不写入数据库")
    parser.add_argument(
        "--a-share-only",
        action="store_true",
        help="只取 A 股股本口径（去掉 H 股 / B 股）",
    )
    parser.add_argument(
        "--upsert-chunk",
        type=int,
        default=2000,
        help="每次 upsert 行数",
    )
    return parser.parse_args()


def load_stock_codes(limit: int = 0, codes_str: str = "") -> list[str]:
    """返回 6 位代码列表（已过滤 A 股）。"""
    if codes_str:
        raw = [c.strip() for c in codes_str.split(",") if c.strip()]
        return [c for c in raw if is_a_share_code(c)]

    codes = [
        c
        for c in db_ops.iter_stock_codes(active_only=True)
        if is_a_share_code(c)
    ]
    if limit > 0:
        codes = codes[:limit]
    return codes


def load_existing_share_snapshots(code6s: list[str]) -> dict[str, dict[str, float]]:
    """读取每只股票的最新 share 快照（ trade_date -> total_share/float_share ）。

    返回：{code6: {trade_date_str: {"total_share": float, "float_share": float}}}
    """
    if not code6s:
        return {}
    hermes_codes = [to_hermes(c) for c in code6s]
    out: dict[str, dict[str, dict[str, float]]] = {c: {} for c in code6s}

    # 按 hermes 前缀分组查询，避免 SQL IN 子句过长
    batch_size = 500
    for i in range(0, len(hermes_codes), batch_size):
        batch = hermes_codes[i : i + batch_size]
        placeholders = ",".join([f"'{c}'" for c in batch])
        sql = f"""
            SELECT stock_code, trade_date, total_share, float_share
            FROM market_data.share
            WHERE stock_code IN ({placeholders})
              AND trade_date >= '2023-01-01'
            ORDER BY stock_code, trade_date
        """
        with get_session() as s:
            rows = s.execute(text(sql)).fetchall()
        for r in rows:
            hermes = r[0]
            code6 = hermes[2:] if hermes[:2] in ("SH", "SZ", "BJ") else hermes
            if code6 not in out:
                out[code6] = {}
            out[code6][r[1].strftime("%Y-%m-%d")] = {
                "total_share": float(r[2] or 0),
                "float_share": float(r[3] or 0),
            }
    return out


def fetch_cninfo_changes(code6: str, a_share_only: bool = False) -> pd.DataFrame:
    """调用 akshare 巨潮接口，返回标准化后的股本变动 DataFrame。

    列：change_date, total_share, float_share（单位：亿股）

    Args:
        a_share_only: True 时只取 A 股口径：
            total_share = 人民币普通股（缺失时用 总股本 - H股 - B股）
            float_share = 人民币普通股 - 流通受限股份
    """
    try:
        import akshare as ak

        df = ak.stock_share_change_cninfo(symbol=code6)
    except Exception as e:
        logger.warning("[%s] 巨潮接口调用失败: %s", code6, e)
        return pd.DataFrame()

    if df is None or df.empty:
        return pd.DataFrame()

    # 只保留我们需要的列
    if a_share_only:
        rename_map = {
            "变动日期": "change_date",
            "总股本": "total_share_all",
            "人民币普通股": "a_share_total",
            "流通受限股份": "restricted",
            "境外上市外资股-H股": "h_share",
            "境内上市外资股-B股": "b_share",
        }
    else:
        rename_map = {
            "变动日期": "change_date",
            "总股本": "total_share",
            "已流通股份": "float_share",
        }

    missing = [k for k in rename_map if k not in df.columns]
    if missing:
        logger.warning("[%s] 巨潮返回缺少字段: %s", code6, missing)
        return pd.DataFrame()

    df = df[list(rename_map.keys())].rename(columns=rename_map)
    df["change_date"] = pd.to_datetime(df["change_date"], errors="coerce")
    for col in rename_map.values():
        if col != "change_date":
            df[col] = pd.to_numeric(df[col], errors="coerce")

    # 万股 -> 亿股
    for col in rename_map.values():
        if col != "change_date":
            df[col] = df[col] / 10000.0

    if a_share_only:
        # A 股总股本：优先用“人民币普通股”，缺失时从总股本中扣除 H 股 / B 股
        a_total = df["a_share_total"].where(
            df["a_share_total"].notna(),
            df["total_share_all"] - df["h_share"].fillna(0) - df["b_share"].fillna(0),
        )
        # A 股流通股本：A 股总股本 - 流通受限股份
        a_float = a_total - df["restricted"].fillna(0)
        df["total_share"] = a_total.round(4)
        df["float_share"] = a_float.round(4)

    df = df.dropna(subset=["change_date", "total_share", "float_share"])
    if df.empty:
        return pd.DataFrame()

    # 剔除异常：A 股流通股本不能为负，也不能超过 A 股总股本
    df = df[df["float_share"] >= 0]
    df = df[df["float_share"] <= df["total_share"] * 1.001]
    if df.empty:
        return pd.DataFrame()

    df = df[["change_date", "total_share", "float_share"]].sort_values("change_date").reset_index(drop=True)
    return df


def build_weekly_snapshots(
    code6: str,
    cninfo_df: pd.DataFrame,
    existing: dict[str, dict[str, float]],
    sample_dates: list[date],
) -> list[dict[str, Any]]:
    """合并历史变动 + 现有快照，生成该股票每周股本记录。"""
    # 构建时间线：change_date -> 最新股本
    timeline: dict[str, dict[str, float]] = {}

    # 1) 巨潮历史变动
    for _, row in cninfo_df.iterrows():
        d = row["change_date"].strftime("%Y-%m-%d")
        timeline[d] = {
            "total_share": float(row["total_share"]),
            "float_share": float(row["float_share"]),
        }

    # 2) 现有 share 快照（仅用于补充巨潮未覆盖的 2025 年及之后最新日期，
    #    避免旧口径数据覆盖 2024 及更早的历史回填）
    if existing:
        for d, v in existing.items():
            if d >= "2025-01-01":
                timeline[d] = v

    if not timeline:
        return []

    # 按日期排序并做前向填充
    dates = sorted(timeline.keys())
    rows = []
    for sd in sample_dates:
        sd_str = sd.strftime("%Y-%m-%d")
        # 找到 <= sd 的最新一条
        latest = None
        for d in dates:
            if d > sd_str:
                break
            latest = d
        if latest is None:
            continue
        v = timeline[latest]
        if v["total_share"] <= 0 or v["float_share"] <= 0:
            continue
        rows.append(
            {
                "stock_code": to_hermes(code6),
                "trade_date": sd_str,
                "total_share": v["total_share"],
                "float_share": v["float_share"],
            }
        )
    return rows


def generate_sample_dates(start: str, end: str) -> list[date]:
    """生成每周五采样日，若结束日不是周五则额外包含结束日。"""
    start_dt = datetime.strptime(start, "%Y-%m-%d").date()
    end_dt = datetime.strptime(end, "%Y-%m-%d").date()

    # pandas 每周五
    fridays = pd.date_range(start=start, end=end, freq="W-FRI")
    dates = [d.date() for d in fridays]

    # 若结束日不是周五，补一个结束日本身
    if end_dt not in dates:
        dates.append(end_dt)
        dates.sort()

    return dates


def process_one_stock(
    code6: str,
    existing_map: dict[str, dict[str, float]],
    sample_dates: list[date],
    a_share_only: bool = False,
) -> list[dict[str, Any]]:
    cninfo_df = fetch_cninfo_changes(code6, a_share_only=a_share_only)
    existing = existing_map.get(code6, {})
    return build_weekly_snapshots(code6, cninfo_df, existing, sample_dates)


def main() -> int:
    args = parse_args()

    logger.info("开始回填历史股本: %s ~ %s", args.start, args.end)
    logger.info(
        "参数: workers=%s, dry_run=%s, limit=%s, a_share_only=%s",
        args.max_workers, args.dry_run, args.limit, args.a_share_only,
    )

    # 1. 股票代码
    code6s = load_stock_codes(limit=args.limit, codes_str=args.codes)
    if not code6s:
        logger.error("无可用股票代码")
        return 1
    logger.info("共 %s 只股票待处理", len(code6s))

    # 2. 采样日期
    sample_dates = generate_sample_dates(args.start, args.end)
    logger.info("采样日期: %s 个（%s ~ %s）", len(sample_dates), sample_dates[0], sample_dates[-1])

    # 3. 预读现有 share 快照
    logger.info("预读现有 share 快照...")
    t0 = time.time()
    existing_map = load_existing_share_snapshots(code6s)
    logger.info("预读完成，耗时 %.1fs", time.time() - t0)

    # 4. 多线程拉取 + 合并
    all_rows: list[dict[str, Any]] = []
    ok = fail = 0
    t0 = time.time()
    with ThreadPoolExecutor(max_workers=args.max_workers) as pool:
        futures = {
            pool.submit(
                process_one_stock, c, existing_map, sample_dates, args.a_share_only
            ): c
            for c in code6s
        }
        for i, future in enumerate(as_completed(futures), 1):
            c6 = futures[future]
            try:
                rows = future.result()
            except Exception as e:
                fail += 1
                logger.warning("[%s] 处理异常: %s", c6, e)
                continue
            if rows:
                all_rows.extend(rows)
                ok += 1
            else:
                fail += 1
            if i % 100 == 0 or i == len(code6s):
                logger.info(
                    "进度 %s/%s, OK=%s, FAIL=%s, rows=%s",
                    i,
                    len(code6s),
                    ok,
                    fail,
                    len(all_rows),
                )

    elapsed = time.time() - t0
    logger.info(
        "拉取完成: OK=%s, FAIL=%s, 总条数=%s, 耗时=%.1fs",
        ok,
        fail,
        len(all_rows),
        elapsed,
    )

    if not all_rows:
        logger.warning("无数据可写入")
        return 0

    # 5. 写入
    if args.dry_run:
        logger.info("DRY RUN: 本应写入 %s 行", len(all_rows))
        # 打印几只样本（首尾各 3 行，方便看后段是否被现有快照覆盖）
        df = pd.DataFrame(all_rows)
        sample_codes = df["stock_code"].unique()[:3]
        for sc in sample_codes:
            sub = df[df["stock_code"] == sc]
            preview = pd.concat([sub.head(3), sub.tail(3)])
            logger.info("样本 %s:\n%s", sc, preview.to_string(index=False))
        return 0

    # 5. 写入前先清理目标区间内非采样日的旧记录，保证“每周一条”
    keep_dates = {d.strftime("%Y-%m-%d") for d in sample_dates}
    keep_str = ",".join([f"'{d}'" for d in keep_dates])
    delete_sql = f"""
        DELETE FROM market_data.share
        WHERE trade_date BETWEEN '{args.start}' AND '{args.end}'
          AND trade_date NOT IN ({keep_str})
    """
    logger.info("清理非采样日旧记录...")
    with get_session() as s:
        deleted = s.execute(text(delete_sql)).rowcount
        s.commit()
    logger.info("清理完成: 删除 %s 行", deleted)

    logger.info("开始写入 share 表，chunk=%s", args.upsert_chunk)
    t0 = time.time()
    total_written = 0
    for i in range(0, len(all_rows), args.upsert_chunk):
        chunk = all_rows[i : i + args.upsert_chunk]
        total_written += db_ops.upsert_share(chunk)
        logger.info("已写入 %s/%s", total_written, len(all_rows))
    logger.info("写入完成: %s 行, 耗时 %.1fs", total_written, time.time() - t0)

    return 0


if __name__ == "__main__":
    sys.exit(main())