#!/usr/bin/env python3
import argparse
import sys
import os
import re
import json
import yaml
import time
from openai import OpenAI
from pathlib import Path

# 加载配置
BASE_DIR = Path(__file__).parent.parent
ASSETS_DIR = BASE_DIR / "assets"

# 加载外部配置文件（所有可调优参数全部在assets目录下的yaml文件，无需改代码）
try:
    # 基础配置
    with open(ASSETS_DIR / "sci_fi_map.yaml", "r", encoding="utf-8") as f:
        SCI_FI_WORD_MAP = yaml.safe_load(f)
    with open(ASSETS_DIR / "stage_config.yaml", "r", encoding="utf-8") as f:
        STAGE_CONFIG = yaml.safe_load(f)
    # 调优配置
    with open(ASSETS_DIR / "expression_map.yaml", "r", encoding="utf-8") as f:
        EXPRESSION_MAP = yaml.safe_load(f)
    with open(ASSETS_DIR / "prompt_config.yaml", "r", encoding="utf-8") as f:
        PROMPT_CONFIG = yaml.safe_load(f)
    with open(ASSETS_DIR / "validation_config.yaml", "r", encoding="utf-8") as f:
        VALIDATION_CONFIG = yaml.safe_load(f)
    # 词表配置
    with open(BASE_DIR / "references" / "l1_word_list.json", "r", encoding="utf-8") as f:
        word_data = json.load(f)
        L1_WORD_LIST = set([w['word'].lower() for w in word_data.get('words', []) if w.get('enabled', True)])
except Exception as e:
    print(f"❌ 配置文件加载失败，请检查yaml格式是否正确: {str(e)}")
    sys.exit(1)

# 初始化LLM客户端，配置从环境变量读取
try:
    client = OpenAI(
        api_key=os.getenv("OPENAI_API_KEY", "your-api-key"),
        base_url=os.getenv("OPENAI_BASE_URL", "https://ark.cn-beijing.volces.com/api/v3")
    )
    MODEL = os.getenv("OPENAI_MODEL", "volcengine/doubao-seed-2-0-pro-260215")
except Exception as e:
    print(f"❌ LLM客户端初始化失败: {str(e)}")
    sys.exit(1)

def load_input(input_path):
    """加载输入内容，支持单个文件或目录批量加载"""
    input_path = Path(input_path)
    if not input_path.exists():
        print(f"❌ 输入路径不存在: {input_path}")
        sys.exit(1)
    
    if input_path.is_file():
        with open(input_path, "r", encoding="utf-8") as f:
            return [(input_path.name, f.read())]
    elif input_path.is_dir():
        # 批量加载目录下所有txt文件
        script_files = list(input_path.glob("*.txt"))
        if not script_files:
            print(f"❌ 目录下没有找到txt格式的剧本文件: {input_path}")
            sys.exit(1)
        results = []
        for f in script_files:
            with open(f, "r", encoding="utf-8") as fp:
                results.append((f.name, fp.read()))
        return results
    else:
        print(f"❌ 不支持的输入类型: {input_path}")
        sys.exit(1)

# ============================================================
# Pipeline 节点日志工具
# ============================================================
class PipelineLogger:
    """流程节点日志记录器，跟踪每个节点的输入/输出/耗时/状态"""
    def __init__(self, verbose=True):
        self.verbose = verbose
        self.nodes = []  # 全部节点记录
        self._current = None

    def start_node(self, name, description=""):
        self._current = {
            'name': name,
            'description': description,
            'start_time': time.time(),
            'input_summary': '',
            'output_summary': '',
            'output_full': '',
            'status': 'running',
            'warnings': [],
            'duration': 0,
        }
        if self.verbose:
            print(f"\n{'='*60}")
            print(f"▶ 节点 {len(self.nodes)+1}: {name}")
            if description:
                print(f"  {description}")
            print(f"{'='*60}")

    def set_input(self, summary):
        if self._current:
            self._current['input_summary'] = summary
            if self.verbose:
                print(f"📥 输入: {summary[:200]}{'...' if len(summary)>200 else ''}")

    def set_output(self, summary, full_content=""):
        if self._current:
            self._current['output_summary'] = summary
            self._current['output_full'] = full_content or summary
            if self.verbose:
                print(f"📤 输出: {summary[:300]}{'...' if len(summary)>300 else ''}")

    def add_warning(self, msg):
        if self._current:
            self._current['warnings'].append(msg)
            if self.verbose:
                print(f"⚠️  {msg}")

    def end_node(self, status='pass'):
        if self._current:
            self._current['duration'] = time.time() - self._current['start_time']
            self._current['status'] = status
            icon = {'pass': '✅', 'warn': '⚠️', 'fail': '❌', 'skip': '⏭️'}.get(status, '❓')
            if self.verbose:
                print(f"{icon} 状态: {status.upper()} | 耗时: {self._current['duration']:.1f}s")
                if self._current['warnings']:
                    print(f"   共 {len(self._current['warnings'])} 条警告")
            self.nodes.append(self._current)
            self._current = None

    def get_trace_report(self):
        """生成完整流程追踪报告"""
        lines = []
        lines.append("\n" + "="*60)
        lines.append("📋 全流程追踪报告")
        lines.append("="*60)
        total_time = sum(n['duration'] for n in self.nodes)
        for i, n in enumerate(self.nodes, 1):
            icon = {'pass': '✅', 'warn': '⚠️', 'fail': '❌', 'skip': '⏭️'}.get(n['status'], '❓')
            lines.append(f"\n{icon} 节点{i}: {n['name']} [{n['status'].upper()}] ({n['duration']:.1f}s)")
            if n['input_summary']:
                lines.append(f"   输入: {n['input_summary'][:80]}{'...' if len(n['input_summary'])>80 else ''}")
            if n['output_summary']:
                lines.append(f"   输出: {n['output_summary'][:80]}{'...' if len(n['output_summary'])>80 else ''}")
            for w in n['warnings']:
                lines.append(f"   ⚠️ {w}")
        lines.append(f"\n⏱️ 总耗时: {total_time:.1f}s | 节点数: {len(self.nodes)}")
        pass_count = sum(1 for n in self.nodes if n['status'] == 'pass')
        warn_count = sum(1 for n in self.nodes if n['status'] == 'warn')
        fail_count = sum(1 for n in self.nodes if n['status'] == 'fail')
        lines.append(f"📊 结果: {pass_count}通过 / {warn_count}警告 / {fail_count}失败")
        lines.append("="*60)
        return "\n".join(lines)

    def get_intermediate_outputs(self):
        """获取所有中间产物，用于保存"""
        outputs = {}
        for n in self.nodes:
            outputs[n['name']] = {
                'status': n['status'],
                'duration': n['duration'],
                'content': n['output_full'],
                'warnings': n['warnings'],
            }
        return outputs


# ============================================================
# LLM 调用封装（带重试）
# ============================================================
SYSTEM_MSG = (
    "你是专为4-8岁儿童打造的英文台词生产专家。"
    "你必须严格遵守所有生产规则，100%忠于原剧本内容，"
    "禁止新增或删除原剧本没有的信息。"
    "输出必须完全匹配目标Stage的词汇、句法、句长要求。"
)

def call_llm(prompt, stage, max_retries=3):
    """通用LLM调用，带重试和错误处理"""
    temp = PROMPT_CONFIG['temperature'].get(stage, 0.3)
    for attempt in range(1, max_retries + 1):
        try:
            response = client.chat.completions.create(
                model=MODEL,
                messages=[
                    {"role": "system", "content": SYSTEM_MSG},
                    {"role": "user", "content": prompt}
                ],
                temperature=temp,
                max_tokens=2000,
                timeout=180,
                extra_body={"thinking": {"type": "disabled"}}
            )
            return response.choices[0].message.content, None
        except Exception as e:
            if attempt < max_retries:
                wait = 2 ** attempt
                print(f"  ⚠️ 第{attempt}次调用失败，{wait}秒后重试: {str(e)}")
                time.sleep(wait)
            else:
                return None, f"LLM调用失败（已重试{max_retries}次）: {str(e)}"


# ============================================================
# Node 1: 输入解析（纯代码）
# ============================================================
def node_input_parse(raw_text, logger):
    """解析输入文本，检测类型/角色/行数"""
    logger.start_node("输入解析", "检测输入类型、角色、行数，纯代码节点")
    logger.set_input(raw_text[:100])

    lines = [l.strip() for l in raw_text.strip().split('\n') if l.strip()]
    # 检测输入语言类型
    cn_chars = len(re.findall(r'[\u4e00-\u9fff]', raw_text))
    en_chars = len(re.findall(r'[a-zA-Z]', raw_text))
    if cn_chars > 0 and en_chars > 0:
        lang_type = "中英混合"
    elif cn_chars > 0:
        lang_type = "纯中文"
    elif en_chars > 0:
        lang_type = "纯英文"
    else:
        lang_type = "未知"

    # 提取角色名
    roles = set()
    for line in lines:
        for sep in ['：', ':']:
            if sep in line:
                role = line.split(sep, 1)[0].strip()
                if role and len(role) < 20:
                    roles.add(role)
                break

    result = {
        'raw_text': raw_text,
        'line_count': len(lines),
        'lang_type': lang_type,
        'cn_chars': cn_chars,
        'en_chars': en_chars,
        'roles': list(roles),
        'lines': lines,
    }
    summary = f"{lang_type} | {len(lines)}行 | 角色: {', '.join(roles) if roles else '未检测到'} | 中文{cn_chars}字/英文{en_chars}字"
    logger.set_output(summary)
    logger.end_node('pass')
    return result


# ============================================================
# Node 2: 输入归一（LLM #1，按语言类型分流）
# ============================================================
def node_normalize(parsed_input, logger):
    """根据输入语言类型分流处理：
    - 纯中文：直接保留
    - 中英混合：归一为标准中文
    - 纯英文：只做格式标准化，保留英文（避免英→中→英的语义折损）
    """
    raw_text = parsed_input['raw_text']
    lang_type = parsed_input['lang_type']

    # ---- 纯中文：已是标准格式，直接返回 ----
    if lang_type == '纯中文' and parsed_input['roles']:
        logger.start_node("输入归一", "纯中文输入，已是标准格式，跳过")
        logger.set_input(f"纯中文 | {parsed_input['line_count']}行")
        logger.set_output(raw_text, raw_text)
        logger.end_node('pass')
        return raw_text

    # ---- 纯英文：只做格式标准化，不转中文 ----
    if lang_type == '纯英文':
        logger.start_node("输入归一", "纯英文输入 → 格式标准化（保留英文，不转中文）")
        logger.set_input(f"纯英文 | {parsed_input['line_count']}行")

        if parsed_input['roles']:
            # 已有角色格式，直接保留
            logger.set_output(raw_text, raw_text)
            logger.end_node('pass')
            return raw_text

        # 没有角色格式，用LLM标准化格式（保持英文）
        prompt = f"""将以下英文输入标准化为「Role: dialogue」格式。
要求：
1. 保持英文原文，不要翻译成中文
2. 完整保留所有剧情、动作、角色关系、道具信息
3. 每行格式：RoleName: dialogue content
4. 只输出标准化结果，不要其他说明

输入：
{raw_text}"""
        result, error = call_llm(prompt, 'S2')
        if error:
            logger.add_warning(error)
            logger.end_node('fail')
            return None
        logger.set_output(result, result)
        logger.end_node('pass')
        return result

    # ---- 中英混合：归一为标准中文 ----
    logger.start_node("输入归一", "LLM调用 → 中英混合输入统一转成标准中文")
    logger.set_input(f"中英混合 | {parsed_input['line_count']}行")

    prompt = f"""将以下中英混合输入统一转成标准中文「角色: 台词」格式。
要求：
1. 完整保留所有剧情、动作、角色关系、道具、事件触发点信息
2. 不能丢失任何核心内容，不能新增内容
3. 每行格式：角色名: 台词内容
4. 保留原始英文角色名（如Ben、Amy等），不要翻译角色名
5. 只输出转换结果，不要其他说明

输入：
{raw_text}"""

    result, error = call_llm(prompt, 'S2')
    if error:
        logger.add_warning(error)
        logger.end_node('fail')
        return None

    logger.set_output(result, result)

    # 简单校验：行数不应该差太多
    result_lines = [l for l in result.strip().split('\n') if l.strip()]
    if len(result_lines) < len(parsed_input['lines']) * 0.5:
        logger.add_warning(f"归一后行数({len(result_lines)})远少于原始行数({parsed_input['line_count']})，可能丢失内容")
        logger.end_node('warn')
    else:
        logger.end_node('pass')
    return result


# ============================================================
# Node 3: AR预处理（LLM #2，按语言类型分流）
# ============================================================
def node_ar_preprocess(normalized_text, stage, lang_type, logger):
    """拆句 + 打AR等级，根据语言类型走不同路径"""
    is_english = (lang_type == '纯英文')
    desc = "LLM调用 → 英文句子难度标注+拆分" if is_english else "LLM调用 → 中文意群拆解（为英语思维重构做准备）"
    logger.start_node("AR预处理", desc)
    logger.set_input(f"{len(normalized_text)}字 | 语言={'EN' if is_english else 'CN'}")

    fidelity_rule = "100%忠于原剧本内容：禁止新增任何原剧本没有的信息、禁止删除任何原剧本已有的信息" if PROMPT_CONFIG['script_fidelity']['strictly_no_add'] and PROMPT_CONFIG['script_fidelity']['strictly_no_delete'] else "允许适当调整细节"

    if is_english:
        # 纯英文路径：直接在英文上标注难度和拆分
        prompt = f"""你是英文台词AR预处理专家。将下面的英文台词拆分为简单短句并标注AR等级。

### 核心规则
{fidelity_rule}
保持英文原文，不要翻译成中文。

### 拆分规则
1. 复合句拆成简单句
2. 多步动作拆成单动作句
3. 因果/条件句拆分，保留全部信息
4. 情绪与事实拆分
5. 保留完整事件链、角色关系、道具信息
6. 上下文连贯：前文已提到的事物，后文用代词(them/it/these)指代，不重复同一名词

### 输出格式
每行格式：Role: dialogue [AR等级]
AR1=最简单(4-7词) AR2=简单(7-10词) AR3=中等(10-15词) AR4=复杂(15+词)

### 待处理台词
{normalized_text}"""
    else:
        # 中文路径：意群拆解（不是简单拆句，而是为英语思维重构做准备）
        prompt = f"""你是儿童台词意群拆解专家。你的任务不是简单拆句，而是将中文台词拆解为**最小语义意群**，为后续用英语思维重构做准备。

### 核心理念
目标不是翻译，而是让每个意群都能被英语母语者用最自然的方式独立表达。
拆解时要想："一个4-8岁英语母语孩子会怎么一句句说出这些意思？"

### 忠实度（最高优先级）
{fidelity_rule}

### 4个保留机制（绝对不能改）
1. 保留完整事件动词链
2. 保留所有事件触发点
3. 保留完整道具逻辑链
4. 保留原有角色关系

### 意群拆解规则
1. 一个意群 = 一个完整的最小语义单元（一个动作/一个状态/一个情绪）
2. 因果关系拆成独立事实意群，去掉"因为/所以"等中文逻辑连接词（英语孩子不这样说话）
3. 多步动作拆成单动作意群
4. 条件+结果拆成独立意群，去掉"如果/假设"等假设框架
5. 情绪和事实分开成独立意群（"我好开心，终于找到了" → "我好开心！" + "终于找到了！"）
6. 去掉中文特有的修饰堆叠，只留核心可感知信息
7. 保留语气和情绪强度，用标点标注（！表强烈 ？表疑问）

### 上下文连贯性（重要！）
台词是连续的对话，拆解时必须考虑上下文语境：
- 前文已经明确提到的事物，后文可以用代词/简称指代，不需要重复完整名词
- 例如："收集了很多衣物准备义卖" → 下一句"小朋友一定会喜欢这些东西的" → 意群保留"这些"而不是重复"衣物"
- 英语孩子说话时自然会用 them/these/it/this stuff 等代词，而不会反复重复同一个名词

### 输出格式
每行格式：角色: 意群内容 [AR等级]
AR等级说明：AR1=最简单(对应4-7英文词) AR2=简单(7-10词) AR3=中等(10-15词) AR4=复杂(15+词)

### 待拆解台词
{normalized_text}"""

    result, error = call_llm(prompt, stage)
    if error:
        logger.add_warning(error)
        logger.end_node('fail')
        return None

    logger.set_output(result, result)

    # 校验AR等级分布
    ar_counts = {f'AR{i}': len(re.findall(f'\\[AR{i}\\]', result)) for i in range(1, 5)}
    total_ar = sum(ar_counts.values())
    if total_ar > 0:
        ar_summary = " | ".join([f"{k}:{v}({v/total_ar:.0%})" for k, v in ar_counts.items() if v > 0])
        logger.set_output(f"{ar_summary}\n{result}", result)

        if stage == 'S1' and (ar_counts.get('AR3', 0) + ar_counts.get('AR4', 0)) > 0:
            ratio = (ar_counts['AR3'] + ar_counts['AR4']) / total_ar
            logger.add_warning(f"S1阶段存在AR3/AR4句子({ratio:.0%})，生成时需降级")
        if stage == 'S2' and ar_counts.get('AR4', 0) > 0:
            logger.add_warning(f"S2阶段存在AR4句子({ar_counts['AR4']}条)，生成时需降级")

    status = 'warn' if (logger._current and logger._current['warnings']) else 'pass'
    logger.end_node(status)
    return result


# ============================================================
# Node 4: 分级英文生成 / 英文难度适配（LLM #3，按语言类型分流）
# ============================================================
def _build_naturalization_rules():
    """构建自然化规则文本（中英文路径共用）"""
    sci_fi_map_str = "\n".join([f"{k} → {v}" for k, v in SCI_FI_WORD_MAP.items()])
    emotion_map_rule = "优先使用以下映射匹配情绪词：" + "、".join([f"{k}→{v}" for k,v in EXPRESSION_MAP['emotion_map'].items()]) if PROMPT_CONFIG['naturalization']['enable_emotion_word'] else "不使用自定义情绪词映射"
    synonym_replace_rule = "可使用以下同义口语替换（不改变原意）：" + "、".join([f"{k}→{v}" for k,v in EXPRESSION_MAP['synonym_replace'].items()]) if PROMPT_CONFIG['naturalization']['enable_synonym_replace'] else "不使用同义替换"
    split_rule = "包含2个及以上信息的句子拆成单信息短句" if PROMPT_CONFIG['naturalization']['enable_long_sentence_split'] else "不拆分长句"
    repeat_rule = "允许自然重复（比如It is dirty. Very dirty.）" if PROMPT_CONFIG['naturalization']['allow_repeat_expression'] else "不允许重复表达"
    exclamation_rule = "情绪强烈的句子可用感叹号" if PROMPT_CONFIG['naturalization']['enable_exclamation_mark'] else "统一使用句号"
    return sci_fi_map_str, emotion_map_rule, synonym_replace_rule, split_rule, repeat_rule, exclamation_rule

def node_generate_english(ar_processed_text, stage, lang_type, logger):
    """按Stage要求生成/适配英文台词，纯英文走难度适配，中文走翻译"""
    is_english = (lang_type == '纯英文')
    task_desc = f"英文难度适配 → {stage}" if is_english else f"英语思维重构意群 → {stage}"
    logger.start_node("分级英文生成", f"LLM调用 → {task_desc}({STAGE_CONFIG[stage]['age']})")
    logger.set_input(f"Stage={stage} | 蓝思值={STAGE_CONFIG[stage]['lexile']} | 模式={'EN适配' if is_english else '意群重构'}")

    sci_fi_map_str, emotion_map_rule, synonym_replace_rule, split_rule, repeat_rule, exclamation_rule = _build_naturalization_rules()

    if is_english:
        # 纯英文路径：直接在英文上做难度适配，不经过中文
        prompt = f"""将下面的英文台词调整为{stage}级别的儿童英文。
注意：这是英文难度适配，不是翻译！保持原文语义和语气，只调整词汇/句法/句长难度。

### 目标Stage：{stage}
年龄段：{STAGE_CONFIG[stage]['age']}
要求：{STAGE_CONFIG[stage]['rules']}
蓝思值：{STAGE_CONFIG[stage]['lexile']}

### 适配规则
1. 用更简单的同义词替换难词（保持原意）
2. 拆分复合句为简单句
3. 调整句长匹配目标Stage要求
4. 同义替换：{synonym_replace_rule}
5. 情绪词：{emotion_map_rule}
6. 标点：{exclamation_rule}
7. 绝对禁止成人化连接词（actually/in fact/however等）
8. 完全符合母语小朋友说话习惯
9. 科幻词汇替换：
{sci_fi_map_str}

### 上下文连贯性
台词是连续对话，前文已提到的事物，后文用代词(them/it/these)指代，不重复同一名词。

### 忠实度要求
100%保留原文语义，禁止新增/删除任何原剧本没有的信息。

### 输出格式（只输出台词，不要其他内容）
Role: adapted dialogue

### 待适配台词
{ar_processed_text}"""
    else:
        # 中文路径：英语思维重构意群（不是翻译！）
        prompt = f"""你的任务是用英语母语儿童的思维方式，将下面的中文意群重构为地道的{stage}级别英文台词。

### 核心理念（最重要！）
**这不是翻译！** 你要做的是：
1. 理解每个中文意群要表达的核心意思和情绪
2. 忘掉中文的表达方式和句式结构
3. 想象一个{STAGE_CONFIG[stage]['age']}的英语母语孩子在同样情境下会怎么说
4. 用那个孩子的方式说出来

判断标准：如果一句话听起来像是"从中文翻译过来的"，那就是失败的。

### 上下文连贯性（重要！）
台词是连续对话，重构时必须通读全文语境：
- 前文已提到的事物，后文用代词指代（them/it/these/this stuff），不要反复重复同一个名词
- 英语母语孩子的对话习惯：第一次提到用具体词，后续自然切换为代词/简称
- 例如："我们收集了很多衣物" → "We got lots of clothes!" → 下一句"小朋友会喜欢这些" → "Kids will love them!" 而不是 "Kids will love these clothes!"
- 同理，上一句刚说完某个动作，下一句不需要重复主语/对象，直接连接即可

### 目标Stage：{stage}
年龄段：{STAGE_CONFIG[stage]['age']}
要求：{STAGE_CONFIG[stage]['rules']}
蓝思值：{STAGE_CONFIG[stage]['lexile']}

### 自然表达规则
1. 情绪词：{emotion_map_rule}
2. 口语替换：{synonym_replace_rule}
3. 句子拆分：{split_rule}
4. 自然重复：{repeat_rule}
5. 标点：{exclamation_rule}
6. 绝对禁止成人化连接词（actually/in fact/however等）
7. 优先用祈使句、感叹句、短问句等儿童高频句式
8. 科幻词汇替换：
{sci_fi_map_str}

### 忠实度要求
100%保留原始意群的语义，禁止新增/删除任何原剧本没有的信息。

### 输出格式（只输出台词，不要其他内容）
角色: 英文台词

### 待重构意群
{ar_processed_text}"""

    result, error = call_llm(prompt, stage)
    if error:
        logger.add_warning(error)
        logger.end_node('fail')
        return None

    logger.set_output(result, result)
    logger.end_node('pass')
    return result


# ============================================================
# Node 5: 代码侧校验（纯代码）
# ============================================================
def node_code_validation(english_text, stage, logger):
    """代码侧校验：句长、连接词、从句、时态、OOV、重复"""
    logger.start_node("代码侧校验", "纯代码 → 句长/连接词/从句/时态/超纲词/重复")
    logger.set_input(f"Stage={stage}")

    validation = validate_output(english_text, stage)
    oov_words = check_out_of_vocab(english_text, stage)

    all_warnings = validation['warnings'][:]
    if oov_words:
        all_warnings.append(f"超纲词: {', '.join(oov_words)}")

    for w in all_warnings:
        logger.add_warning(w)

    status = 'fail' if len(all_warnings) > 3 else ('warn' if all_warnings else 'pass')
    summary = f"{len(all_warnings)}条问题" if all_warnings else "全部通过"
    logger.set_output(summary)
    logger.end_node(status)
    return {'warnings': all_warnings, 'oov_words': oov_words, 'status': status}


# ============================================================
# Node 6: 自动修正（可选，LLM #4）
# ============================================================
def node_auto_fix(english_text, validation_result, stage, logger):
    """仅在校验不通过时触发，根据校验报告自动修正"""
    if validation_result['status'] == 'pass':
        logger.start_node("自动修正", "校验通过，跳过修正")
        logger.set_output("无需修正")
        logger.end_node('skip')
        return english_text

    logger.start_node("自动修正", "LLM调用 → 根据校验报告修正问题")
    warnings_text = "\n".join([f"- {w}" for w in validation_result['warnings']])
    logger.set_input(f"{len(validation_result['warnings'])}条待修正问题")

    prompt = f"""以下英文台词存在问题，请按要求修正。

### 当前台词
{english_text}

### 需要修正的问题
{warnings_text}

### 修正要求
目标Stage：{stage}（{STAGE_CONFIG[stage]['age']}）
句长要求：{STAGE_CONFIG[stage]['rules']}
1. 只修正上面列出的具体问题
2. 保持原有台词内容和角色不变
3. 不要新增/删除台词行

### 输出格式（只输出修正后的台词，不要其他说明）
角色: 修正后的英文台词"""

    result, error = call_llm(prompt, stage)
    if error:
        logger.add_warning(f"修正失败: {error}")
        logger.set_output("修正失败，保留原始结果", english_text)
        logger.end_node('fail')
        return english_text

    logger.set_output(result, result)
    logger.end_node('pass')
    return result


# ============================================================
# Pipeline 主流程
# ============================================================
def generate_single_script(input_text, stage, verbose=True):
    """完整Pipeline：6个可观测节点，生成单个剧本的台词"""
    logger = PipelineLogger(verbose=verbose)

    # Node 1: 输入解析
    parsed = node_input_parse(input_text, logger)
    lang_type = parsed['lang_type']

    # Node 2: 输入归一（纯英文保留英文，中文/混合归一为中文）
    normalized = node_normalize(parsed, logger)
    if normalized is None:
        return logger.get_trace_report()

    # Node 3: AR预处理（纯英文在英文上拆分，中文在中文上拆分）
    ar_processed = node_ar_preprocess(normalized, stage, lang_type, logger)
    if ar_processed is None:
        return logger.get_trace_report()

    # Node 4: 分级英文生成（纯英文做难度适配，中文做意群重构）
    english = node_generate_english(ar_processed, stage, lang_type, logger)
    if english is None:
        return logger.get_trace_report()

    # Node 5: 代码侧校验
    validation = node_code_validation(english, stage, logger)

    # Node 6: 自动修正（仅校验不通过时触发）
    final = node_auto_fix(english, validation, stage, logger)

    # 拼接最终输出
    output_lines = []
    output_lines.append(f"【Stage {stage} 英文台词（适配{STAGE_CONFIG[stage]['age']}）】")
    output_lines.append(final)

    # 附加追踪报告
    output_lines.append(logger.get_trace_report())

    result = "\n".join(output_lines)
    # 附加 logger 引用供外部保存中间产物
    result = type('PipelineResult', (), {'text': result, 'logger': logger, '__str__': lambda self: self.text})()
    return result

def extract_script_lines(content):
    """从 LLM 输出中提取纯台词行，排除元数据标签"""
    lines = content.strip().split('\n')
    script_lines = []
    for line in lines:
        line = line.strip()
        # 跳过元数据行
        if not line or line.startswith('【') or line.startswith('---') or line.startswith('#'):
            continue
        # 台词行格式: "角色: 台词"
        if ':' in line or '：' in line:
            # 提取冒号后的台词部分
            sep = '：' if '：' in line else ':'
            parts = line.split(sep, 1)
            if len(parts) == 2:
                script_lines.append(parts[1].strip())
    return script_lines

def simple_lemmatize(word):
    """简单词形还原，覆盖常见词尾变化"""
    w = word.lower()
    # 不规则过去式常见补充
    irregulars = {
        'went': 'go', 'got': 'get', 'said': 'say', 'made': 'make',
        'came': 'come', 'took': 'take', 'ran': 'run', 'saw': 'see',
        'had': 'have', 'did': 'do', 'was': 'be', 'were': 'be',
        'ate': 'eat', 'gave': 'give', 'told': 'tell', 'found': 'find',
        'thought': 'think', 'knew': 'know', 'felt': 'feel', 'left': 'leave',
        'heard': 'hear', 'sat': 'sit', 'stood': 'stand', 'lost': 'lose',
        'brought': 'bring', 'kept': 'keep', 'began': 'begin', 'fell': 'fall',
        'children': 'child', 'mice': 'mouse', 'feet': 'foot', 'teeth': 'tooth',
    }
    if w in irregulars:
        return irregulars[w]
    # 规则词尾: -ing, -ed, -s, -er, -est, -ly
    if w.endswith('ing'):
        base = w[:-3]
        if base + 'e' in L1_WORD_LIST: return base + 'e'  # making -> make
        if len(base) > 2 and base[-1] == base[-2]: return base[:-1]  # running -> run
        if base in L1_WORD_LIST: return base
    if w.endswith('ed'):
        base = w[:-2]
        if base in L1_WORD_LIST: return base
        if base + 'e' in L1_WORD_LIST: return base + 'e'  # liked -> like
        if w[:-1] in L1_WORD_LIST: return w[:-1]  # cleaned -> clean (clean+ed)
        if len(base) > 2 and base[-1] == base[-2]: return base[:-1]  # stopped -> stop
    if w.endswith('ies'):
        base = w[:-3] + 'y'
        if base in L1_WORD_LIST: return base  # carries -> carry
    if w.endswith('es'):
        base = w[:-2]
        if base in L1_WORD_LIST: return base
    if w.endswith('s') and not w.endswith('ss'):
        base = w[:-1]
        if base in L1_WORD_LIST: return base
    if w.endswith('er'):
        base = w[:-2]
        if base in L1_WORD_LIST: return base
        if base + 'e' in L1_WORD_LIST: return base + 'e'
    if w.endswith('est'):
        base = w[:-3]
        if base in L1_WORD_LIST: return base
        if base + 'e' in L1_WORD_LIST: return base + 'e'
    if w.endswith('ly'):
        base = w[:-2]
        if base in L1_WORD_LIST: return base
    return w

def check_out_of_vocab(script_content, stage):
    """检查超纲词汇，规则从配置文件读取"""
    if not VALIDATION_CONFIG['vocab_validation']['enable_OOV_remind']:
        return []
    # 只提取纯台词行，排除元数据
    script_lines = extract_script_lines(script_content)
    script_text = ' '.join(script_lines)
    # 提取英文单词
    words = re.findall(r"[a-zA-Z']+", script_text)
    words = [word.lower().strip("'") for word in words]
    # 过滤停用词
    stop_words = set(VALIDATION_CONFIG['vocab_validation'].get('stop_words', []))
    words = [word for word in words if word not in stop_words and len(word) > 1]
    # 词形还原后再查
    out_of_vocab = []
    seen = set()
    for word in words:
        lemma = simple_lemmatize(word)
        if lemma not in L1_WORD_LIST and word not in L1_WORD_LIST and word not in seen:
            out_of_vocab.append(word)
            seen.add(word)
    # 根据配置的允许比例判断是否超标
    total_words = len([w for w in words if w not in stop_words and len(w) > 1])
    if total_words == 0:
        return []
    oov_ratio = len(out_of_vocab) / total_words
    max_ratio_key = f"{stage}_allow_OOV_ratio"
    max_ratio = VALIDATION_CONFIG['vocab_validation'].get(max_ratio_key, 0.15)
    if oov_ratio > max_ratio:
        out_of_vocab.append(f"(超纲率 {oov_ratio:.1%} 超过允许上限 {max_ratio:.0%})")
    return out_of_vocab

def validate_output(content, stage):
    """代码侧校验: 句长、词汇、AR等级等"""
    warnings = []
    script_lines = extract_script_lines(content)
    if not script_lines:
        return {'warnings': ['未检测到有效台词行']}

    # 句长校验
    stage_word_limits = {
        'S1': (4, 7), 'S2': (7, 10), 'S3': (10, 15), 'S4': (15, 20)
    }
    min_len, max_len = stage_word_limits.get(stage, (4, 20))
    for line in script_lines:
        words = re.findall(r"[a-zA-Z']+", line)
        word_count = len(words)
        if word_count > max_len + 3:  # 容差3个词
            warnings.append(f"句长超标: \"{line[:40]}...\" 共{word_count}词，{stage}要求{min_len}-{max_len}词")

    # 禁用连接词校验
    adult_connectors = ['actually', 'in fact', 'however', 'therefore', 'furthermore', 'moreover', 'nevertheless']
    all_text = ' '.join(script_lines).lower()
    for conn in adult_connectors:
        if conn in all_text:
            warnings.append(f"成人化连接词: 检测到 '{conn}'，{stage} 禁止使用")

    # S1 禁止复合从句
    if stage == 'S1':
        s1_banned = ['because', 'when', 'before', 'after', 'if', 'although', 'while']
        for word in s1_banned:
            if re.search(r'\b' + word + r'\b', all_text):
                warnings.append(f"S1禁止从句连词: 检测到 '{word}'")

    # S1/S2 禁止完成时/虚拟语气
    if stage in ['S1', 'S2']:
        complex_patterns = [r"\bhave\s+\w+ed\b", r"\bhas\s+\w+ed\b", r"\bwould\b", r"\bcould\b", r"\bshould\b"]
        for pat in complex_patterns:
            match = re.search(pat, all_text)
            if match:
                warnings.append(f"复杂时态/语气: 检测到 '{match.group()}'，{stage} 禁止使用")

    # 重复台词检测
    seen_lines = {}
    for line in script_lines:
        normalized = line.strip().lower()
        if normalized in seen_lines:
            warnings.append(f"重复台词: \"{line[:40]}\" 连续出现{seen_lines[normalized] + 1}次")
        seen_lines[normalized] = seen_lines.get(normalized, 0) + 1

    return {'warnings': warnings}

def save_result(output_dir, filename, content):
    """保存结果到文件"""
    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    output_file = output_dir / f"result_{filename}"
    with open(output_file, "w", encoding="utf-8") as f:
        f.write(content)
    return output_file

def main():
    parser = argparse.ArgumentParser(description="4-8岁儿童英文台词标准化生产工具（带全流程监控）")
    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument("--input", type=str, help="直接输入待处理的剧本文本")
    group.add_argument("--path", type=str, help="待处理的单个剧本文件路径或包含多个剧本的目录路径")
    group.add_argument("--feishu-url", type=str, help="飞书文档链接，从飞书表格读取台词并写回结果")
    parser.add_argument("--stage", type=str, choices=["S1", "S2", "S3", "S4"], required=True, help="目标难度等级 S1/S2/S3/S4")
    parser.add_argument("--output", type=str, help="结果输出目录，不指定则直接打印到控制台")
    parser.add_argument("--verbose", action="store_true", default=True, help="显示每个节点的详细过程日志（默认开启）")
    parser.add_argument("--quiet", action="store_true", help="静默模式，只输出最终结果")
    parser.add_argument("--trace", action="store_true", help="保存中间产物到output目录（需配合--output使用）")
    # 飞书模式参数
    parser.add_argument("--source-col", type=str, default="中文台词", help="飞书表格源列名称（默认: 中文台词）")
    parser.add_argument("--target-col", type=str, default="英文台词", help="飞书表格目标列名称（默认: 英文台词）")
    parser.add_argument("--table-index", type=int, default=0, help="文档中第几个表格（从0开始）")
    args = parser.parse_args()

    verbose = not args.quiet

    # 飞书模式：读取飞书表格 → pipeline处理 → 写回飞书
    if args.feishu_url:
        run_feishu_mode(args)
        return

    # 处理输入
    if args.input:
        input_list = [("direct_input", args.input)]
    else:
        input_list = load_input(args.path)

    # 批量生成
    results = []
    for filename, text in input_list:
        print(f"\n🚀 正在处理: {filename}")
        result = generate_single_script(text, args.stage, verbose=verbose)
        results.append((filename, str(result)))
        if not verbose:
            print(result)
        # 保存结果
        if args.output:
            save_path = save_result(args.output, filename, str(result))
            print(f"💾 结果已保存到: {save_path}")
            # --trace: 保存中间产物
            if args.trace and hasattr(result, 'logger'):
                trace_dir = Path(args.output) / "trace"
                trace_dir.mkdir(parents=True, exist_ok=True)
                intermediates = result.logger.get_intermediate_outputs()
                for node_name, data in intermediates.items():
                    trace_file = trace_dir / f"{filename}_{node_name}.txt"
                    with open(trace_file, "w", encoding="utf-8") as f:
                        f.write(f"节点: {node_name}\n")
                        f.write(f"状态: {data['status']}\n")
                        f.write(f"耗时: {data['duration']:.1f}s\n")
                        if data['warnings']:
                            f.write(f"警告: {'; '.join(data['warnings'])}\n")
                        f.write(f"\n--- 输出内容 ---\n{data['content']}\n")
                print(f"📝 中间产物已保存到: {trace_dir}/")

    print(f"\n✅ 全部处理完成，共处理{len(results)}个剧本")

def _preprocess_feishu_cell(text: str) -> str:
    """
    预处理飞书单元格文本，解决台词粘连问题。
    1. 清洗HTML换行标签（<br>/<br/>等）→ 真实换行
    2. 智能拆分粘连台词行（如 "Tom: Hi!Dan: Hello!" → 两行）
    3. 清理多余空行
    """
    # Step 1: HTML标签清洗
    text = re.sub(r'<br\s*/?\s*>', '\n', text, flags=re.IGNORECASE)
    text = re.sub(r'</?p\s*>', '\n', text, flags=re.IGNORECASE)

    # Step 2: 如果已有正常换行则直接返回
    lines = [l.strip() for l in text.split('\n') if l.strip()]
    if len(lines) > 1:
        return '\n'.join(lines)

    # Step 3: 单行文本 → 智能拆分粘连台词
    # 匹配模式：句末标点紧跟角色名冒号（如 "okay?Dan:" 或 "fine.Grace:"）
    # 在角色名前插入换行，角色名支持中英文
    text = re.sub(
        r'([.!?。！？"\)）])\s*(?=[A-Z][a-zA-Z]*\s*[:：]|[\u4e00-\u9fff]+\s*[:：]|User\s*[:：]|Mum\s*[:：]|Dad\s*[:：])',
        r'\1\n',
        text
    )
    lines = [l.strip() for l in text.split('\n') if l.strip()]
    return '\n'.join(lines)


def run_feishu_mode(args):
    """飞书模式：从飞书表格读取台词 → pipeline处理 → 写回飞书"""
    # 导入飞书 I/O 能力
    feishu_io_path = Path(__file__).parent.parent.parent / "feishu-table-translate-fill" / "scripts"
    sys.path.insert(0, str(feishu_io_path))
    try:
        from feishu_io import FeishuDocIO
    except ImportError:
        print("❌ 无法加载飞书 I/O 模块")
        print(f"   请确认 feishu-table-translate-fill 技能已部署在: {feishu_io_path}")
        print("   或设置环境变量 FEISHU_IO_PATH 指向 feishu_io.py 所在目录")
        sys.exit(1)

    verbose = not args.quiet
    io = FeishuDocIO()

    # Step 1: 读取飞书表格
    print(f"\n📄 正在读取飞书文档...")
    print(f"   源列: {args.source_col} | 目标列: {args.target_col}")
    try:
        rows = io.read_table_column(
            url=args.feishu_url,
            column_name=args.source_col,
            table_index=args.table_index
        )
    except Exception as e:
        print(f"❌ 读取飞书文档失败: {e}")
        sys.exit(1)

    print(f"✅ 读取到 {len(rows)} 行台词")

    # Step 2: 通过 pipeline 逐行处理
    results = []
    for i, row in enumerate(rows):
        text = row["content"]
        if not text.strip():
            print(f"  行{row['row']}: 跳过（空内容）")
            continue
        # 输入预处理：清洗HTML标签 + 智能拆分粘连台词
        text = _preprocess_feishu_cell(text)
        print(f"\n🚀 处理行 {row['row']} ({i+1}/{len(rows)})")
        result = generate_single_script(text, args.stage, verbose=verbose)

        # 提取纯台词部分（去掉追踪报告）
        result_text = str(result)
        # 截取到追踪报告之前的内容
        trace_marker = "\n" + "=" * 60 + "\n📋"
        if trace_marker in result_text:
            clean_text = result_text[:result_text.index(trace_marker)].strip()
        else:
            clean_text = result_text.strip()
        # 去掉开头的 Stage 标题行
        lines = clean_text.split("\n")
        if lines and lines[0].startswith("【"):
            clean_text = "\n".join(lines[1:]).strip()

        results.append({"row": row["row"], "content": clean_text})

    if not results:
        print("\n⚠️ 没有生成任何结果")
        return

    # Step 3: 写回飞书表格
    print(f"\n✍️  正在写回飞书文档 ({len(results)} 行)...")
    try:
        stats = io.write_table_column(
            url=args.feishu_url,
            column_name=args.target_col,
            data=results,
            table_index=args.table_index
        )
        print(f"\n✅ 写回完成: 成功{stats['success']} 跳过{stats.get('skipped',0)} 失败{stats['failed']}")
    except Exception as e:
        print(f"❌ 写回飞书失败: {e}")
        # 保存到本地作为备份
        backup_file = Path(args.output or ".") / "feishu_backup_results.json"
        import json as _json
        with open(backup_file, "w", encoding="utf-8") as f:
            _json.dump(results, f, ensure_ascii=False, indent=2)
        print(f"   结果已备份到: {backup_file}")

    # 同时保存到本地
    if args.output:
        output_dir = Path(args.output)
        output_dir.mkdir(parents=True, exist_ok=True)
        for item in results:
            save_path = output_dir / f"result_row{item['row']}.txt"
            with open(save_path, "w", encoding="utf-8") as f:
                f.write(item["content"])
        print(f"💾 本地副本已保存到: {output_dir}/")

    print(f"\n✅ 飞书模式处理完成，共处理 {len(results)} 行台词")


if __name__ == "__main__":
    main()
