"""一次性回填脚本:为 Qdrant 存量**多源** point 的 payload 补充 sources 字段。

背景:
    M5/M6 旧链路未把 M3 多源记录写入 EmbeddingResult / Qdrant payload,
    导致知识库检索只能看到单源。新链路修复后,此脚本为存量数据补齐。

范围(关键设计):
    本地文件模式 Qdrant 写入逐点成本 ≈0.4~0.65s,全量 3.3 万条回填需数小时,
    且本地模式为单进程锁,长时间占用会阻塞后续定时调度。
    但**只有多源 point 需要 sources 字段**——单源 point 的展示层
    (CLI/MCP/日报)会在 sources 为空时兜底显示 source_id。
    因此本脚本仅回填指纹库中 source_ids 长度 > 1 的多源记录,分钟级完成。

数据来源:
    指纹库 data/dedup/fingerprints.sqlite3 的 source_ids 列(权威多源记录)。

用法:
    uv run python -m scripts.backfill_qdrant_sources          # 实际执行
    uv run python -m scripts.backfill_qdrant_sources --dry-run  # 只统计不写入
"""

from __future__ import annotations

import argparse
import json
import sqlite3
import sys
from collections import defaultdict
from pathlib import Path

from loguru import logger

from vectorstore import DEFAULT_COLLECTION, make_qdrant_client
from vectorstore.client import url_hash_to_uuid


def _setup_logger(level: str) -> None:
    logger.remove()
    logger.add(
        sys.stderr,
        level=level,
        format="{time:YYYY-MM-DD HH:mm:ss} | {level} | {name} | {message}",
    )


def _load_multi_source(db_path: Path) -> dict[str, list[str]]:
    """从指纹库读多源记录: url_hash -> sources(仅 source_ids 长度 > 1)。"""
    conn = sqlite3.connect(db_path)
    conn.row_factory = sqlite3.Row
    rows = conn.execute(
        "SELECT url_hash, source_id, source_ids FROM fingerprints"
    ).fetchall()
    conn.close()

    result: dict[str, list[str]] = {}
    for r in rows:
        if not r["source_ids"]:
            continue
        try:
            val = json.loads(r["source_ids"])
        except (TypeError, ValueError):
            continue
        if not isinstance(val, list):
            continue
        sources = [s for s in val if s]
        if len(sources) <= 1:
            continue
        # 保证主源居首
        if r["source_id"] and r["source_id"] in sources:
            sources = [r["source_id"], *[s for s in sources if s != r["source_id"]]]
        result[r["url_hash"]] = sources
    return result


def main() -> int:
    parser = argparse.ArgumentParser(description="Qdrant 存量多源 payload 回填 sources 字段")
    parser.add_argument("--db", default="data/dedup/fingerprints.sqlite3",
                        help="指纹库路径")
    parser.add_argument("--collection", default=DEFAULT_COLLECTION)
    parser.add_argument("--dry-run", action="store_true", help="只统计不写入")
    parser.add_argument("--log-level", default="INFO")
    args = parser.parse_args()

    _setup_logger(args.log_level)

    multi = _load_multi_source(Path(args.db))
    logger.info("指纹库多源记录: {} 条", len(multi))
    if not multi:
        logger.info("无多源记录,退出")
        return 0

    # 按 sources 值分组(同值一次 set_payload)
    groups: dict[str, list[str]] = defaultdict(list)
    for url_hash, sources in multi.items():
        groups[json.dumps(sources, ensure_ascii=False)].append(url_hash)

    if args.dry_run:
        logger.info("待回填: {} 组 / {} 条多源 point (--dry-run 不写入)",
                    len(groups), len(multi))
        return 0

    client = make_qdrant_client()
    updated = 0
    missing = 0
    for sources_json, url_hashes in groups.items():
        payload = {"sources": json.loads(sources_json)}
        ids = [url_hash_to_uuid(h) for h in url_hashes]
        try:
            client.set_payload(
                collection_name=args.collection,
                payload=payload,
                points=ids,
            )
            updated += len(ids)
        except KeyError:
            # 部分 point 已不在集合中(如指纹库保留但向量库已清理),逐点重试
            for uid in ids:
                try:
                    client.set_payload(
                        collection_name=args.collection,
                        payload=payload,
                        points=[uid],
                    )
                    updated += 1
                except KeyError:
                    missing += 1
    logger.info("回填完成: 更新 {} 条多源 point,{} 条不存在于集合",
                updated, missing)

    client.close()
    return 0


if __name__ == "__main__":
    raise SystemExit(main())