#!/usr/bin/env python3
"""batch_generate_summaries — 批量为 session log 生成 LLM 摘要。

从双索引中找出所有缺少 description 的主对话 entry，
并发调用百炼千问 API 生成摘要，写回索引。

用法：
  python -B scripts/ops/batch_generate_summaries.py                  # 全量
  python -B scripts/ops/batch_generate_summaries.py --limit 10       # 只处理 10 条
  python -B scripts/ops/batch_generate_summaries.py --concurrency 5  # 并发 5
  python -B scripts/ops/batch_generate_summaries.py --dry-run        # 预览
"""

import asyncio
import json
import os
import re
import sys
import time

from dotenv import load_dotenv

load_dotenv()

import dashscope
from dashscope import Application

DASHSCOPE_API_KEY = os.environ.get("DASHSCOPE_API_KEY", "")
if not DASHSCOPE_API_KEY:
    raise RuntimeError("DASHSCOPE_API_KEY 未设置，请检查 .env 文件")
SUMMARY_APP_ID = os.environ.get("DASHSCOPE_APP_ID_SUMMARY", "")
if not SUMMARY_APP_ID:
    raise RuntimeError("DASHSCOPE_APP_ID_SUMMARY 未设置，请检查 .env 文件")
dashscope.api_key = DASHSCOPE_API_KEY

# 导入索引管理函数
sys.path.insert(0, os.path.join(os.path.dirname(__file__)))
from extract_kiro_session import (
    load_index, save_index, load_full_index, save_full_index,
)

SYSTEM_PROMPT = """你是一个专业的技术对话分析师。你的任务是为 AI 编程助手的一轮执行（execution）生成简洁的中文摘要。

背景：一个对话（chatSession）包含多轮执行（execution）。每轮执行 = 用户发一条消息 → AI 完成响应。你收到的是单轮执行的完整记录。

摘要规则：
1. 只描述本轮执行实际完成的工作，不要描述历史背景
2. 列出完成的功能点/任务（一轮可能完成多个）
3. 包含关键技术细节：文件路径、模块名、数据库表、API 端点等
4. bug 修复要说明原因和方案
5. 不写过程性描述（"用户说..."），只写结果
6. 内容太短或无实质内容的，写"无实质内容"
7. 不限字数，信息完整优先，避免截断失真

重要：
- "执行摘要"（📋）是最可靠的信息源，优先基于它判断本轮做了什么
- 如果"用户输入"包含 CONTEXT TRANSFER，那是之前多轮的历史摘要，不是本轮工作
- 对话记录中的实际工具调用和文件变更才是本轮的真实操作

请直接输出摘要，不要添加任何前缀或解释。"""


def _extract_summary_content(md_content: str) -> str:
    """检测 CONTEXT TRANSFER，替换用户输入为简短标注。"""
    ct_pattern = re.compile(
        r"## 2\. 用户输入\s*\n```\s*\n.*?CONTEXT TRANSFER", re.DOTALL
    )
    if ct_pattern.search(md_content):
        md_content = re.sub(
            r"(## 2\. 用户输入)\s*\n```[\s\S]*?```\s*\n(?=## 3\.)",
            r"\1\n\n[本轮为 Context Transfer 续接，已省略。]\n\n",
            md_content,
        )
    return md_content


PLACEHOLDER_PREFIX = "[待生成摘要]"


def collect_targets(index: dict, include_placeholder: bool = False) -> list[tuple[str, dict]]:
    """收集所有缺少 description 的主对话 entry。

    include_placeholder=True 时，也包含占位标记的 entry（用于覆盖生成真实摘要）。
    """
    targets = []
    for eid, ent in index.get("entries", {}).items():
        if ent.get("is_sub"):
            continue
        # 跳过已替代或无日志的占位条目
        if ent.get("superseded_by") or ent.get("no_log"):
            continue
        desc = ent.get("description", "")
        if not desc:
            targets.append((eid, ent))
        elif include_placeholder and desc.startswith(PLACEHOLDER_PREFIX):
            targets.append((eid, ent))
    # 按 startTime 排序（旧的先处理）
    targets.sort(key=lambda t: t[1].get("startTime", ""))
    return targets


def load_md_content(eid: str, entry: dict) -> str | None:
    """加载 entry 对应的 main_*.md 内容。"""
    out_dir = entry.get("output_dir", "")
    if not out_dir or not os.path.isdir(out_dir):
        return None
    eid_short = eid[:8]
    main_files = sorted(
        f for f in os.listdir(out_dir)
        if f.startswith("main_") and f.endswith(".md") and eid_short in f
    )
    if not main_files:
        main_files = sorted(
            f for f in os.listdir(out_dir)
            if f.startswith("main_") and f.endswith(".md")
        )
    if not main_files:
        return None
    parts = []
    for mf in main_files:
        try:
            with open(os.path.join(out_dir, mf), "r", encoding="utf-8") as fh:
                parts.append(fh.read())
        except Exception:
            continue
    if not parts:
        return None
    content = "\n\n---\n\n".join(parts)
    content = _extract_summary_content(content)
    if len(content) > 60000:
        content = content[:60000] + "\n\n[TRUNCATED]"
    return content


async def generate_one(
    content: str,
    semaphore: asyncio.Semaphore,
    max_retries: int = 3,
) -> str:
    """调用百炼 Application API 生成摘要，带限流和指数退避。"""
    # CHANGE 2026-03-22 | openai SDK → dashscope Application API（P14 迁移收尾）
    prompt = f"请为以下单轮执行记录生成摘要：\n\n{content}"
    async with semaphore:
        for attempt in range(max_retries):
            try:
                resp = await asyncio.to_thread(
                    Application.call,
                    app_id=SUMMARY_APP_ID,
                    prompt=prompt,
                )
                if resp.status_code == 200:
                    return (resp.output.get("text", "") or "").strip()
                else:
                    raise RuntimeError(f"Application.call 失败: {resp.message}")
            except Exception as e:
                if attempt < max_retries - 1:
                    wait = 2 ** attempt
                    print(f"  ⏳ 重试 {attempt+1}: {e}", file=sys.stderr)
                    await asyncio.sleep(wait)
                else:
                    return ""
    return ""


async def process_target(
    eid: str,
    entry: dict,
    semaphore: asyncio.Semaphore,
) -> tuple[str, str]:
    """处理单个 target，返回 (eid, description)。"""
    content = load_md_content(eid, entry)
    if not content:
        return (eid, "")
    desc = await generate_one(content, semaphore)
    return (eid, desc)



async def main():
    import argparse
    parser = argparse.ArgumentParser(description="批量生成 session 摘要")
    parser.add_argument("--limit", type=int, default=0,
                        help="只处理前 N 条（0=全量）")
    parser.add_argument("--concurrency", type=int, default=10,
                        help="并发数（默认 10）")
    parser.add_argument("--batch-size", type=int, default=20,
                        help="每批处理条数，每批完成后保存 checkpoint（默认 20）")
    parser.add_argument("--replace-placeholder", action="store_true",
                        help="也处理占位标记 '[待生成摘要]' 的 entry")
    parser.add_argument("--dry-run", action="store_true",
                        help="预览待处理列表")
    args = parser.parse_args()

    index = load_index()
    full_index = load_full_index()
    targets = collect_targets(index, include_placeholder=args.replace_placeholder)

    if args.limit > 0:
        targets = targets[:args.limit]

    print(f"待处理: {len(targets)} 条缺少 description 的主对话 entry")

    if not targets:
        print("全部已有摘要，无需处理")
        return

    if args.dry_run:
        for eid, ent in targets[:20]:
            st = ent.get("startTime", "?")
            od = ent.get("output_dir", "?")
            print(f"  {eid[:8]} | {st} | {od}")
        if len(targets) > 20:
            print(f"  ... 还有 {len(targets) - 20} 条")
        return

    semaphore = asyncio.Semaphore(args.concurrency)

    t0 = time.time()
    total_generated = 0
    total_failed = 0
    batch_size = args.batch_size

    # 分批处理 + 每批 checkpoint 保存
    for batch_start in range(0, len(targets), batch_size):
        batch = targets[batch_start:batch_start + batch_size]
        batch_num = batch_start // batch_size + 1
        total_batches = (len(targets) + batch_size - 1) // batch_size
        print(f"\n📦 批次 {batch_num}/{total_batches}（{len(batch)} 条）...")

        tasks = [process_target(eid, ent, semaphore)
                 for eid, ent in batch]
        results = await asyncio.gather(*tasks)

        # 写回索引
        batch_generated = 0
        full_entries = full_index.get("entries", {})
        idx_entries = index.get("entries", {})
        for eid, desc in results:
            if not desc:
                total_failed += 1
                continue
            if eid in idx_entries:
                idx_entries[eid]["description"] = desc
            if eid in full_entries:
                full_entries[eid]["description"] = desc
            batch_generated += 1

        total_generated += batch_generated

        # checkpoint：每批完成后立即保存
        if batch_generated > 0:
            save_index(index)
            save_full_index(full_index)
            elapsed = time.time() - t0
            print(f"  ✅ 本批 {batch_generated} 条已保存（累计 {total_generated}，耗时 {elapsed:.1f}s）")

    elapsed = time.time() - t0
    print(f"\n完成: {total_generated}/{len(targets)} 生成成功，"
          f"{total_failed} 失败，耗时 {elapsed:.1f}s")




if __name__ == "__main__":
    asyncio.run(main())
