#!/usr/bin/env python3
"""合并 staging 增量到主库（task: 截断 bug 修复）。

**核心不变量（截断 bug 回归测试核心）**：
    合并后 main 的行数 >= 合并前 main 的行数（绝不截断）。

raw_redownload.py 的 save_one() 用"本次拉的几天"覆盖写整年文件 → 整年数据被截
（NAS 2026 从 121 行截成 4 行就是这 bug）。本脚本负责按 symbol-year 安全合并：
    pd.concat([main, staging]).drop_duplicates(subset=['date'], keep='last')
staging 的新数据/修订值优先（keep='last'），main 已有的历史保留。

用法：
  python3 merge_increment.py --staging data_cache/daily_update/raw --main data_cache/raw
  python3 merge_increment.py --staging ... --main ... --dry-run

约定：
- staging 与 main 同构：`{root}/{year}/{sh|sz}{code}_daily.parquet`
- staging 文件不修改/不删除（保留供排查），只写 main
"""
from __future__ import annotations

import argparse
import json
import logging
import os
import sys

import pandas as pd

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

# drop_duplicates 的去重键——日线按 date 唯一
DEDUP_KEY = ["date"]


def symbol_from_filename(fname: str) -> str:
    """`sh600000_daily.parquet` → `600000`（剥 exchange 前缀和 _daily.parquet 后缀）。"""
    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 merge_one(staging_file: str, main_file: str, dry_run: bool = False) -> dict:
    """合并一个 staging parquet 到 main。

    返回 stats: {symbol, before, after, added, action}。
    若 main 不存在 → action="created"（直接搬 staging 过去）。
    若 main 存在 → action="merged"，按 DEDUP_KEY 去重，staging 值优先（keep='last'）。
    """
    sym = symbol_from_filename(os.path.basename(staging_file))
    staging_df = pd.read_parquet(staging_file)

    if not os.path.exists(main_file):
        if not dry_run:
            os.makedirs(os.path.dirname(main_file), exist_ok=True)
            staging_df.to_parquet(main_file, index=False)
        return {"symbol": sym, "before": 0, "after": len(staging_df),
                "added": len(staging_df), "action": "created"}

    main_df = pd.read_parquet(main_file)
    before = len(main_df)

    # concat → drop_duplicates(keep='last' 让 staging 的新/修订值覆盖 main) → sort
    merged = pd.concat([main_df, staging_df], ignore_index=True)
    merged = merged.drop_duplicates(subset=DEDUP_KEY, keep="last")
    merged = merged.sort_values(by="date").reset_index(drop=True)
    after = len(merged)

    # 核心不变量：绝不能让 main 变少
    if after < before:
        raise RuntimeError(
            f"INVARIANT VIOLATED: {main_file} {before}→{after} (staging={staging_file})"
        )

    if not dry_run:
        merged.to_parquet(main_file, index=False)

    return {"symbol": sym, "before": before, "after": after,
            "added": after - before, "action": "merged"}


def walk_staging(staging_root: str) -> list[tuple[str, str]]:
    """收集 staging 下所有 `{year}/{sym}_daily.parquet`，返回 [(staging_file, main_file_relpath)]。

    main_file_relpath 是相对 staging_root 的路径（如 `2026/sh600000_daily.parquet`），
    拼到 main_root 即得 main 文件全路径，保持两边同构。
    """
    pairs: list[tuple[str, str]] = []
    for year in sorted(os.listdir(staging_root)):
        ydir = os.path.join(staging_root, year)
        if not os.path.isdir(ydir):
            continue
        for fname in sorted(os.listdir(ydir)):
            if not fname.endswith("_daily.parquet"):
                continue
            rel = os.path.join(year, fname)
            pairs.append((os.path.join(ydir, fname), rel))
    return pairs


def run_merge(staging_root: str, main_root: str, dry_run: bool = False) -> dict:
    """合并 staging → main，返回汇总统计。"""
    if not os.path.isdir(staging_root):
        raise FileNotFoundError(f"staging dir 不存在: {staging_root}")

    pairs = walk_staging(staging_root)
    if not pairs:
        log.warning("staging 无 parquet: %s", staging_root)
        return {"merged": 0, "created": 0, "skipped": 0, "total_new_rows": 0,
                "details": [], "dry_run": dry_run, "ok": True,
                "invariant_violations": []}

    os.makedirs(main_root, exist_ok=True)
    details: list[dict] = []
    merged_n = created_n = skipped_n = total_new = 0
    invariant_violations: list[str] = []

    for i, (staging_file, rel) in enumerate(pairs, 1):
        main_file = os.path.join(main_root, rel)
        try:
            stat = merge_one(staging_file, main_file, dry_run=dry_run)
        except RuntimeError as e:
            # 不变式违反：立刻停（绝不能继续写入更小的 main）
            invariant_violations.append(str(e))
            log.error("[%d/%d] INVARIANT %s", i, len(pairs), e)
            continue
        except Exception as e:  # noqa: BLE001
            log.error("[%d/%d] %s ERR %s", i, len(pairs), rel, e)
            skipped_n += 1
            continue

        details.append(stat)
        if stat["action"] == "created":
            created_n += 1
        else:
            merged_n += 1
        total_new += stat["added"]

        if i % 1000 == 0 or i == len(pairs):
            log.info("[%d/%d] %s %s +%d (main %d→%d)",
                     i, len(pairs), stat["symbol"], stat["action"],
                     stat["added"], stat["before"], stat["after"])

    summary = {
        "merged": merged_n,
        "created": created_n,
        "skipped": skipped_n,
        "total_new_rows": total_new,
        "details": details,
        "dry_run": dry_run,
        "invariant_violations": invariant_violations,
    }
    if invariant_violations:
        # 不变式违反致命——即使个别合并成功也判失败
        summary["ok"] = False
    else:
        summary["ok"] = True
    return summary


def main():
    ap = argparse.ArgumentParser(description="合并 staging 增量到主库（绝不截断）")
    ap.add_argument("--staging", required=True, help="staging 根（如 data_cache/daily_update/raw）")
    ap.add_argument("--main", required=True, help="主库根（如 data_cache/raw）")
    ap.add_argument("--dry-run", action="store_true", help="只报不写")
    ap.add_argument("--summary-json", default=None, help="把汇总写到该 JSON 文件")
    args = ap.parse_args()

    summary = run_merge(args.staging, args.main, dry_run=args.dry_run)

    mode = "[DRY-RUN] " if args.dry_run else ""
    log.info("=== %s合并完成: merged=%d created=%d skipped=%d 新增行=%d ok=%s ===",
             mode, summary["merged"], summary["created"], summary["skipped"],
             summary["total_new_rows"], summary["ok"])
    if summary["invariant_violations"]:
        log.error("不变式违反 %d 条（main 被截），详情见上", len(summary["invariant_violations"]))

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

    # 不变式违反 → 退 1（即便部分成功，也提示 main 可能已损坏需排查）
    sys.exit(0 if summary["ok"] else 1)


if __name__ == "__main__":
    main()