"""
故事提取流程 v4 - 逐步测试脚本
用法：python test_pipeline.py <docx文件路径> [步骤编号]
  步骤编号：1=预处理, 2=判断故事, 3=切片提取, 4=重组, 5=生成故事, all=全部执行
"""
import sys
import os
import json
import asyncio
import time

sys.path.insert(0, os.path.join(os.path.dirname(__file__), "backend"))

from app.services.docx_parser import parse_transcript
from app.services.llm_client import llm_client, _parse_json

TEST_DIR = os.path.join(os.path.dirname(__file__), "backend", "uploads", "test")


def _save_json(filepath, data):
    """保存 JSON 到文件"""
    os.makedirs(os.path.dirname(filepath), exist_ok=True)
    with open(filepath, "w", encoding="utf-8") as f:
        json.dump(data, f, ensure_ascii=False, indent=2)


def _save_text(filepath, text):
    """保存文本到文件"""
    os.makedirs(os.path.dirname(filepath), exist_ok=True)
    with open(filepath, "w", encoding="utf-8") as f:
        f.write(text)


# ============================================================
# 第一步：代码预处理（v3：时间范围截取，允许重叠）
# ============================================================
def step1_preprocess(lines, output_dir, save=True):
    """代码预处理：识别老师 → 过滤杂音 → 按时间范围截取每个学员的全部内容"""
    print("\n" + "=" * 60)
    print("第一步：代码预处理（v3：时间范围截取）")
    print("=" * 60)

    # 1.1 统计每个说话人
    print(f"总行数: {len(lines)}")
    speaker_stats = {}
    for i, line in enumerate(lines):
        speaker = line.get("speaker", "")
        if not speaker:
            continue
        if speaker not in speaker_stats:
            speaker_stats[speaker] = {"lines": 0, "chars": 0, "first": i, "last": i}
        speaker_stats[speaker]["lines"] += 1
        speaker_stats[speaker]["chars"] += len(line.get("content", ""))
        speaker_stats[speaker]["last"] = i

    print(f"\n说话人统计:")
    for s, stats in sorted(speaker_stats.items(), key=lambda x: -x[1]["lines"]):
        print(f"  {s}: {stats['lines']}行, {stats['chars']}字, 行范围[{stats['first']}-{stats['last']}]")

    # 1.2 识别老师（发言最多的人）
    teacher = max(speaker_stats, key=lambda s: speaker_stats[s]["lines"])
    print(f"\n识别老师: {teacher} ({speaker_stats[teacher]['lines']}行)")

    # 1.3 过滤杂音（< 20句的说话人）+ 排除老师
    students = []
    noise_speakers = []
    for s, stats in speaker_stats.items():
        if s == teacher:
            continue
        if stats["lines"] < 20:
            noise_speakers.append(f"{s} ({stats['lines']}句)")
            continue
        students.append(s)

    if noise_speakers:
        print(f"\n杂音过滤: {', '.join(noise_speakers)}")
    print(f"\n有效学员: {students}")

    # 1.4 按时间范围截取：每个学员从第一句到最后一句，中间所有行全部截取
    segments = []
    for speaker in students:
        stats = speaker_stats[speaker]
        first_idx = stats["first"]
        last_idx = stats["last"]
        seg = lines[first_idx:last_idx + 1]

        # 统计该学员在此范围内的发言
        student_lines = [l for l in seg if l.get("speaker", "") == speaker]
        student_chars = sum(len(l.get("content", "")) for l in student_lines)

        segments.append((speaker, seg))
        print(f"  {speaker}: 截取行[{first_idx}-{last_idx}], 共{len(seg)}行, 学员{len(student_lines)}句/{student_chars}字")

    # 1.5 过滤：短段落 + 当事人占比过低（非主讲述人）
    final_segments = []
    filtered = []
    for speaker, seg in segments:
        student_lines = [l for l in seg if l.get("speaker", "") == speaker]
        if len(student_lines) < 15:
            filtered.append(f"{speaker} ({len(student_lines)}句<15)")
            continue
        student_chars = sum(len(l.get("content", "")) for l in student_lines)
        if student_chars < 80:
            filtered.append(f"{speaker} ({student_chars}字<80)")
            continue
        # 当事人说话占比 < 1/10 → 非主讲述人，丢弃
        ratio = len(student_lines) / len(seg) if seg else 0
        if ratio < 0.1:
            filtered.append(f"{speaker} (占比{ratio:.1%}<10%)")
            continue
        final_segments.append((speaker, seg))

    if filtered:
        print(f"\n过滤: {', '.join(filtered)}")

    print(f"\n最终: {len(final_segments)} 个学员段落")
    for i, (speaker, seg) in enumerate(final_segments):
        student_lines = [l for l in seg if l.get("speaker", "") == speaker]
        student_chars = sum(len(l.get("content", "")) for l in student_lines)
        ratio = len(student_lines) / len(seg) if seg else 0
        print(f"  段落{i}: {speaker}, 截取{len(seg)}行, 学员{len(student_lines)}句/{student_chars}字, 占比{ratio:.1%}")

    # 保存结果
    if save:
        step1_dir = os.path.join(output_dir, "step1_预处理")
        # 保存每个学员段落
        for i, (speaker, seg) in enumerate(final_segments):
            text = _segment_to_text(seg)
            safe_name = speaker.replace("/", "_")
            _save_text(os.path.join(step1_dir, f"segment_{i}_{safe_name}.txt"), text)
            print(f"  保存: segment_{i}_{safe_name}.txt ({len(text)}字)")
        # 保存统计摘要
        summary = {
            "total_lines": len(lines),
            "teacher": teacher,
            "noise_filtered": noise_speakers,
            "final_filtered": filtered,
            "segments": [
                {
                    "speaker": s,
                    "total_lines": len(seg),
                    "student_lines": len([l for l in seg if l.get("speaker", "") == s]),
                    "student_chars": sum(len(l.get("content", "")) for l in seg if l.get("speaker", "") == s),
                    "ratio": round(len([l for l in seg if l.get("speaker", "") == s]) / len(seg), 3) if seg else 0,
                }
                for s, seg in final_segments
            ]
        }
        _save_json(os.path.join(step1_dir, "summary.json"), summary)
        print(f"  保存: summary.json")

    return final_segments, teacher


def _get_main_speaker(student_lines):
    if not student_lines:
        return "unknown"
    counts = {}
    for l in student_lines:
        s = l.get("speaker", "")
        counts[s] = counts.get(s, 0) + 1
    return max(counts, key=counts.get)


def _segment_to_text(seg):
    parts = []
    for line in seg:
        speaker = line.get("speaker", "")
        content = line.get("content", "")
        if speaker:
            parts.append(f"{speaker}: {content}")
        else:
            parts.append(content)
    return "\n".join(parts)


# ============================================================
# 第二步：判断是否包含故事 + 提取事实（短段落合并处理）
# ============================================================
async def step2_identify(segments, output_dir, save=True):
    """短段落直接提取事实，长段落先判断是否有故事再提取"""
    print("\n" + "=" * 60)
    print("第二步：判断故事 + 提取事实")
    print("=" * 60)

    SHORT_THRESHOLD = 4000  # 短段落阈值（字数）
    all_facts = []
    step2_results = []  # 记录每个段落的处理结果
    step2_t0 = time.time()

    # 构建所有段落的异步任务，并行执行
    async def process_segment(i, speaker, seg):
        text = _segment_to_text(seg)
        safe_name = speaker.replace("/", "_")
        prefix = f"seg{i}_{safe_name}_"
        print(f"\n--- 段落 {i}: {speaker} ({len(text)} 字) ---")

        result_record = {
            "index": i,
            "speaker": speaker,
            "text_length": len(text),
            "is_short": len(text) <= SHORT_THRESHOLD,
        }

        if len(text) <= SHORT_THRESHOLD:
            print(f"  短段落，跳过判断，直接提取事实")
            facts = await _extract_facts_from_chunk(text, "", output_dir, prefix)
            if facts:
                hints = facts.get("story_hints", "")
                if not hints:
                    hints = f"{speaker}的故事"
                print(f"  提取完成: hints={hints}")
                result_record["status"] = "extracted"
                result_record["hints"] = hints
                return (speaker, seg, hints, [facts]), result_record
            else:
                print(f"  提取失败，跳过")
                result_record["status"] = "extract_failed"
                return None, result_record
        else:
            is_story, hints = await _identify_story(text, output_dir, prefix)
            print(f"  判断结果: is_story={is_story}, hints={hints}")
            result_record["is_story"] = is_story
            result_record["hints"] = hints

            if is_story:
                story_segments_for_step3 = [(speaker, seg, hints)]
                step3_result = await step3_extract_facts(story_segments_for_step3, output_dir, save)
                result_record["status"] = "has_story"
                return step3_result, result_record
            else:
                print(f"  无故事，跳过")
                result_record["status"] = "no_story"
                return None, result_record

    tasks = [process_segment(i, speaker, seg) for i, (speaker, seg) in enumerate(segments)]
    results = await asyncio.gather(*tasks)

    step2_elapsed = round(time.time() - step2_t0, 2)
    print(f"\n⏱ 第二步总耗时: {step2_elapsed}s")

    # 收集结果
    for i, (r, record) in enumerate(results):
        if r is not None:
            if isinstance(r, list):
                all_facts.extend(r)
            else:
                all_facts.append(r)
        step2_results.append(record)

    print(f"\n共 {len(all_facts)} 个故事段落完成事实提取")

    # 保存结果
    if save:
        step2_dir = os.path.join(output_dir, "step2_判断提取")
        _save_json(os.path.join(step2_dir, "results.json"), step2_results)
        print(f"  保存: results.json")

    return all_facts


async def _identify_story(text, output_dir=None, log_prefix="", do_save=True):
    prompt = """判断这段对话是否包含学员的个人故事。直接输出JSON，不要分析。

个人故事特征（满足至少2条）：学员讲述具体经历、包含时间地点人物细节、表达深层情感、老师深入引导疗愈、揭示心理或家庭问题。

非故事内容：纯课堂讲解、简单问答寒暄、课堂管理、纯疗愈引导词。

直接输出JSON：
{"is_story": true或false, "confidence": 0.0到1.0, "story_hints": "一句话描述故事主题和讲述者"}

对话内容：
```
""" + text[:4000] + """
```"""

    try:
        detail = await llm_client.chat_detail([{"role": "user", "content": prompt}], temperature=0.2)
        api_log = {
            "type": "identify_story",
            "prompt": prompt,
            "content": detail.get("content", ""),
            "reasoning_content": detail.get("reasoning_content", ""),
            "finish_reason": detail.get("finish_reason", ""),
            "usage": detail.get("usage", {}),
            "elapsed": detail.get("elapsed", 0),
        }
        print(f"  {log_prefix}API耗时: {api_log['elapsed']}s, finish_reason={api_log['finish_reason']}, "
              f"tokens: {api_log['usage'].get('total_tokens', '?')} "
              f"(reasoning={api_log['usage'].get('reasoning_tokens', 0)})")
        if output_dir and do_save:
            _save_json(os.path.join(output_dir, "api_logs", f"{log_prefix}identify.json"), api_log)
        result = _parse_json(detail["content"])
        if result:
            return bool(result.get("is_story", False)) and result.get("confidence", 0) >= 0.5, result.get("story_hints", "")
    except Exception as e:
        print(f"  {log_prefix}错误: {e}")
    return False, ""


# ============================================================
# 第三步：切片提取事实（Map）
# ============================================================
async def step3_extract_facts(story_segments, output_dir, save=True):
    print("\n" + "=" * 60)
    print("第三步：切片提取事实（Map）")
    print("=" * 60)

    all_facts = []
    step3_t0 = time.time()
    for i, (speaker, seg, hints) in enumerate(story_segments):
        text = _segment_to_text(seg)
        safe_name = speaker.replace("/", "_")
        print(f"\n--- 故事 {i}: {speaker} - {hints[:50]}... ({len(text)} 字) ---")

        chunks = _chunk_text_by_paragraphs(text, max_chars=3000, overlap_paragraphs=5)
        print(f"  切成 {len(chunks)} 块")

        # 保存切片信息
        chunk_records = []
        for j, chunk in enumerate(chunks):
            chunk_records.append({"chunk_index": j, "chars": len(chunk), "preview": chunk[:100] + "..."})

        tasks = [_extract_facts_from_chunk(chunk, hints, output_dir, f"story{i}_{safe_name}_chunk{j}_") for j, chunk in enumerate(chunks)]
        results = await asyncio.gather(*tasks, return_exceptions=True)

        facts = []
        for j, r in enumerate(results):
            if isinstance(r, Exception):
                print(f"  块{j} 错误: {r}")
                chunk_records[j]["status"] = "error"
                chunk_records[j]["error"] = str(r)
            else:
                print(f"  块{j}: {json.dumps(r, ensure_ascii=False)[:200]}")
                facts.append(r)
                chunk_records[j]["status"] = "ok"
                chunk_records[j]["facts"] = r

        all_facts.append((speaker, seg, hints, facts))

        # 保存切片和提取结果
        if save:
            step3_dir = os.path.join(output_dir, "step3_切片提取", f"story_{i}_{safe_name}")
            for j, chunk in enumerate(chunks):
                _save_text(os.path.join(step3_dir, f"chunk_{j}.txt"), chunk)
            _save_json(os.path.join(step3_dir, "chunks_info.json"), chunk_records)
            print(f"  保存: {len(chunks)} 个切片文件 + chunks_info.json")

    step3_elapsed = round(time.time() - step3_t0, 2)
    print(f"\n⏱ 第三步总耗时: {step3_elapsed}s")

    return all_facts


def _chunk_text_by_paragraphs(text, max_chars=3000, overlap_paragraphs=5):
    """按自然段切分，在 max_chars 附近找段落结尾切，重叠 overlap_paragraphs 个自然段"""
    paragraphs = [p for p in text.split("\n") if p.strip()]
    if not paragraphs:
        return [text]
    if len(text) <= max_chars:
        return [text]

    chunks = []
    i = 0
    while i < len(paragraphs):
        # 从当前位置开始，累加段落直到超过 max_chars
        current = []
        current_len = 0
        j = i
        while j < len(paragraphs):
            para = paragraphs[j]
            # 超过 max_chars 且已有内容 → 在这个段落前切断（段落结尾切）
            if current_len + len(para) > max_chars and current:
                break
            current.append(para)
            current_len += len(para)
            j += 1

        chunks.append("\n".join(current))

        # 到末尾了就结束
        if j >= len(paragraphs):
            break

        # 下一次起点：当前切到的位置 - 重叠段落数
        next_i = j - overlap_paragraphs
        # 防止死循环：确保至少前进1个段落
        if next_i <= i:
            next_i = i + 1
        i = next_i

    return chunks


async def _extract_facts_from_chunk(text, hints, output_dir=None, log_prefix="", do_save=True):
    prompt = f"""从这段对话中提取关键信息。只输出JSON。

故事主题：{hints}

提取以下内容：
- events: 发生的事件（尽量详细，保留具体细节）
- emotions: 表达的情感
- quotes: 讲述者的原话（保留原话，不要改写，尽量多提取；完全相同的重复语句只保留一条）
- key_people: 提到的人物

{{"events":["事件1","事件2"],"emotions":["情感1","情感2"],"quotes":["原话1","原话2","原话3"],"key_people":["人物1"]}}

对话内容：
```
{text}
```"""

    detail = await llm_client.chat_detail([{"role": "user", "content": prompt}], temperature=0.1)
    api_log = {
        "type": "extract_facts",
        "prompt": prompt,
        "content": detail.get("content", ""),
        "reasoning_content": detail.get("reasoning_content", ""),
        "finish_reason": detail.get("finish_reason", ""),
        "usage": detail.get("usage", {}),
        "elapsed": detail.get("elapsed", 0),
    }
    print(f"  {log_prefix}API耗时: {api_log['elapsed']}s, finish_reason={api_log['finish_reason']}, "
          f"tokens: {api_log['usage'].get('total_tokens', '?')} "
          f"(reasoning={api_log['usage'].get('reasoning_tokens', 0)})")
    if output_dir and do_save:
        _save_json(os.path.join(output_dir, "api_logs", f"{log_prefix}extract.json"), api_log)
    result = _parse_json(detail["content"])
    return result


# ============================================================
# 模糊去重工具函数
# ============================================================
def _fuzzy_dedup_quotes(quotes: list[str]) -> list[str]:
    """
    模糊去重：对高度相似的语句只保留第一条。
    策略：提取每条语句的核心部分（去掉常见前缀后缀），相似度超过阈值的视为重复。
    """
    import re

    def normalize(s: str) -> str:
        s = s.strip()
        if s.startswith("「") and s.endswith("」"):
            s = s[1:-1]
        for prefix in ["嗯", "啊", "呃", "就是", "然后", "那个", "所以"]:
            if s.startswith(prefix):
                s = s[len(prefix):].strip()
        for suffix in ["啊", "吧", "呢", "嘛", "哦", "呀"]:
            if s.endswith(suffix) and len(s) > 4:
                s = s[:-1]
        s = re.sub(r'(.)\1{2,}', r'\1', s)
        return s

    def similarity(a: str, b: str) -> float:
        if not a or not b:
            return 0.0
        a_bigrams = set(a[i:i+2] for i in range(len(a)-1))
        b_bigrams = set(b[i:i+2] for i in range(len(b)-1))
        if not a_bigrams or not b_bigrams:
            return 0.0
        intersection = a_bigrams & b_bigrams
        union = a_bigrams | b_bigrams
        return len(intersection) / len(union)

    kept = []
    kept_normalized = []
    for q in quotes:
        nq = normalize(q)
        is_dup = False
        for kn in kept_normalized:
            threshold = 0.8 if len(nq) < 15 else 0.7
            if similarity(nq, kn) >= threshold:
                is_dup = True
                break
        if not is_dup:
            kept.append(q)
            kept_normalized.append(nq)

    removed = len(quotes) - len(kept)
    if removed > 0:
        print(f"    模糊去重: {len(quotes)} -> {len(kept)} (去除 {removed} 条相似重复)")

    return kept


# ============================================================
# 第四步：重组（Reduce）
# ============================================================
def step4_merge(all_facts, output_dir, save=True):
    print("\n" + "=" * 60)
    print("第四步：重组（Reduce）")
    print("=" * 60)

    merged = []
    for i, (speaker, seg, hints, facts_list) in enumerate(all_facts):
        print(f"\n--- 故事 {i}: {speaker} ---")

        all_events, all_emotions, all_quotes, all_people = [], [], [], []
        for facts in facts_list:
            if isinstance(facts, dict):
                all_events.extend(facts.get("events", []))
                all_emotions.extend(facts.get("emotions", []))
                all_quotes.extend(facts.get("quotes", []))
                all_people.extend(facts.get("key_people", []))

        all_emotions = list(dict.fromkeys(all_emotions))
        all_people = list(dict.fromkeys(all_people))
        seen = set()
        unique_quotes = [q for q in all_quotes if q not in seen and not seen.add(q)]
        # 模糊去重：去除相似重复语句（如多次释放语句只保留一条）
        unique_quotes = _fuzzy_dedup_quotes(unique_quotes)

        material = {
            "events": all_events,
            "emotions": all_emotions,
            "quotes": unique_quotes,
            "key_people": all_people,
            "story_hints": hints,
        }

        print(f"  事件: {len(all_events)}, 情感: {len(all_emotions)}, 原话: {len(unique_quotes)}, 人物: {len(all_people)}")
        print(f"  素材: {json.dumps(material, ensure_ascii=False)[:500]}")

        merged.append((speaker, seg, material))

    # 保存结果
    if save:
        step4_dir = os.path.join(output_dir, "step4_重组")
        for i, (speaker, seg, material) in enumerate(merged):
            safe_name = speaker.replace("/", "_")
            _save_json(os.path.join(step4_dir, f"material_{i}_{safe_name}.json"), material)
            print(f"  保存: material_{i}_{safe_name}.json")

    return merged


# ============================================================
# 第五步：生成最终故事
# ============================================================
async def step5_generate(merged, output_dir, save=True):
    print("\n" + "=" * 60)
    print("第五步：生成最终故事（并行）")
    print("=" * 60)

    step5_t0 = time.time()

    async def generate_one(i, speaker, seg, material):
        print(f"\n--- 生成故事 {i}: {speaker} ---")

        prompt = f"""基于素材生成第三人称故事。只输出JSON，不要其他内容。

素材：
主题：{material['story_hints']}
事件：{json.dumps(material['events'], ensure_ascii=False)}
情感：{json.dumps(material['emotions'], ensure_ascii=False)}
原话：{json.dumps(material['quotes'], ensure_ascii=False)}
人物：{json.dumps(material['key_people'], ensure_ascii=False)}

写作手法（非常重要）：

【双时空结构】
故事有两个时空层：
- "当下时空"：学员在课堂上与卢慧老师的对话——这是故事的框架，用对话体呈现，保留师生互动的完整脉络
- "过往时空"：学员讲述的过去经历——这是故事的核心，用"场景重建"手法呈现，让读者穿越到那个时空

【场景重建手法（用于"过往的故事"部分）】
当学员讲述过去的经历时，不要写成"她说她小时候…"这种转述体，而是直接把读者拉进那个时空：
- 直接切入场景，用过去时态叙述。比如不要写"文晔回忆说她的儿子15岁去了美国"，而要写"十五岁那年，文晔的儿子独自登上了飞往美国的航班"
- 用五感描写构建画面：她看到了什么、听到了什么、身体感受到了什么（比如"手心出汗"、"脚发凉"这些身体反应要融入场景中）
- 用环境细节暗示情绪：不要直接说"她很痛苦"，而是通过场景让读者感受到痛苦
- 关键情感爆发处用直接引语（「」包裹原话），其他地方用间接引语和叙述自然融合
- 场景之间要有过渡，比如"然而，命运的转折发生在前年冬天——"

【对话体手法（用于"当下的困境"和"转变"部分）】
- 保留师生对话的现场感，写出谁说了什么、对方如何回应
- 在对话中穿插学员的内心活动、身体反应、情绪变化
- 老师的引导语要完整保留，这是疗愈过程的关键

【通用要求】
- 每个事件都要展开写，写出具体过程和细节，不要一句话概括
- 事件之间要有过渡和因果逻辑，不能跳跃
- 情感变化要写出过程：从什么情感转变到什么情感，中间经历了什么
- 不编造素材中没有的内容，但可以根据素材合理补充场景细节（如环境、氛围）
- 每个部分充分展开，不要限制长度，越详细越好
- 原文金句数组中的每条金句必须用英文双引号包裹，如 ["金句1", "金句2"]，不要用「」替代双引号

禁止事项：
- 禁止用一句话概括多个事件（如"她倾诉了挣扎"、"她经历了疗愈过程"）
- 禁止省略素材中的具体细节
- 禁止省略老师的引导和回应
- 禁止事件之间跳跃，要保持时间顺序和因果逻辑
- 禁止过度总结或抽象化描述
- 禁止在"过往的故事"部分全程使用转述体（"她说…她回忆…"），必须用场景重建手法让读者身临其境
- 禁止为了简洁而牺牲内容的丰富性

输出格式：{{"title":"标题","summary":"80字摘要","tags":["标签"],"speaker_nickname":"昵称","content":{{"讲述者":"...","当下的困境":"...","过往的故事":"...","转变":"...","原文金句":["原话1","原话2","原话3"]}}}}"""

        try:
            detail = await llm_client.chat_detail([{"role": "user", "content": prompt}], temperature=0.7, max_tokens=16384)
            raw = detail["content"]
            api_log = {
                "type": "generate_story",
                "speaker": speaker,
                "prompt": prompt,
                "content": raw,
                "reasoning_content": detail.get("reasoning_content", ""),
                "finish_reason": detail.get("finish_reason", ""),
                "usage": detail.get("usage", {}),
                "elapsed": detail.get("elapsed", 0),
            }
            print(f"  API耗时: {api_log['elapsed']}s, finish_reason={api_log['finish_reason']}, "
                  f"tokens: {api_log['usage'].get('total_tokens', '?')} "
                  f"(reasoning={api_log['usage'].get('reasoning_tokens', 0)})")
            print(f"  原始输出长度: {len(raw)} 字")
            if save:
                step5_dir = os.path.join(output_dir, "step5_生成故事")
                os.makedirs(step5_dir, exist_ok=True)
                _save_text(os.path.join(step5_dir, f"raw_{i}_{speaker}.txt"), raw)
                _save_json(os.path.join(output_dir, "api_logs", f"story{i}_{speaker}_generate.json"), api_log)
            result = _parse_json(raw)
        except Exception as e:
            print(f"  生成失败: {e}")
            result = None

        if result:
            print(f"  标题: {result.get('title', '')}")
            print(f"  摘要: {result.get('summary', '')[:100]}")
            content = result.get('content', {})
            if isinstance(content, dict):
                for key, val in content.items():
                    print(f"  {key}: {str(val)[:80]}...")
            print(f"\n  完整JSON:")
            print(json.dumps(result, ensure_ascii=False, indent=2))
        else:
            print("  生成失败！")

        return result

    tasks = [generate_one(i, speaker, seg, material) for i, (speaker, seg, material) in enumerate(merged)]
    results = await asyncio.gather(*tasks)
    stories = [r for r in results if r is not None]

    step5_elapsed = round(time.time() - step5_t0, 2)
    print(f"\n⏱ 第五步总耗时: {step5_elapsed}s")

    # 保存结果
    if save:
        step5_dir = os.path.join(output_dir, "step5_生成故事")
        for i, story in enumerate(stories):
            safe_name = story.get("speaker_nickname", f"story_{i}").replace("/", "_")
            _save_json(os.path.join(step5_dir, f"story_{i}_{safe_name}.json"), story)
            print(f"  保存: story_{i}_{safe_name}.json")

    return stories


# ============================================================
# 主函数
# ============================================================
async def main():
    if len(sys.argv) < 2:
        print("用法: python test_pipeline.py <docx文件路径> [步骤编号]")
        print("  步骤编号: 1=预处理, 2=判断故事, 3=切片提取, 4=重组, 5=生成故事, all=全部")
        sys.exit(1)

    file_path = sys.argv[1]
    step = sys.argv[2] if len(sys.argv) > 2 else "all"

    if not os.path.exists(file_path):
        print(f"文件不存在: {file_path}")
        sys.exit(1)

    # 根据文件名创建输出目录
    basename = os.path.splitext(os.path.basename(file_path))[0]
    output_dir = os.path.join(TEST_DIR, basename)
    os.makedirs(output_dir, exist_ok=True)

    print(f"文件: {file_path}")
    print(f"输出目录: {output_dir}")

    total_t0 = time.time()

    lines = parse_transcript(file_path)
    print(f"解析到 {len(lines)} 行")

    # 如果只跑第5步，从 step4 产物恢复
    if step == "5":
        step4_dir = os.path.join(output_dir, "step4_重组")
        if os.path.exists(step4_dir):
            print("从 step4 产物恢复...")
            merged = []
            for f in sorted(os.listdir(step4_dir)):
                if f.endswith(".json"):
                    with open(os.path.join(step4_dir, f), "r", encoding="utf-8") as fp:
                        material = json.load(fp)
                    parts = f.replace("material_", "").replace(".json", "").split("_", 1)
                    speaker = parts[1] if len(parts) > 1 else f
                    merged.append((speaker, [], material))
            print(f"恢复到 {len(merged)} 个故事")
            stories = await step5_generate(merged, output_dir)

            total_elapsed = round(time.time() - total_t0, 2)
            print(f"\n{'=' * 60}")
            print(f"✅ 全流程完成！总耗时: {total_elapsed}s")
            print(f"   生成故事: {len(stories)} 个")
            print(f"{'=' * 60}")
            return
        else:
            print("step4 产物不存在，请先跑前4步")
            return

    segments, teacher = step1_preprocess(lines, output_dir)
    if step == "1":
        return

    all_facts = await step2_identify(segments, output_dir)
    if step == "2":
        return

    if not all_facts:
        print("没有发现故事段落，结束。")
        return

    merged = step4_merge(all_facts, output_dir)
    if step == "4":
        return

    stories = await step5_generate(merged, output_dir)


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