"""
CoT 数据清洗脚本

功能:
1. 自动识别 think 模式新数据 / 旧数据两种格式
   - 新数据 (think 模式): <think>[草稿]</think> Ref-Delta <answer>
     → 丢弃草稿，将 Ref-Delta 包入 <think>
   - 旧数据: <think>[Ref-Delta]</think> <answer>  (保持原逻辑)
2. 处理重复 <think></think> 标签
3. 验证必要字段完整性
4. 输出清洗报告

用法:
    python clean_cot_data.py input.jsonl -o output.jsonl
    python clean_cot_data.py input.jsonl                    # 覆盖原文件
"""

import re
import sys
import json
import argparse
from pathlib import Path


def clean_response(response: str, gt_answer: str = "") -> tuple[str, list[str]]:
    """清洗单条 response，返回 (清洗后文本, 修复记录列表)"""
    fixes = []

    # ---- 1. 保护 <answer> 和 <key_frame> 块 ----
    answer_match = re.search(r"<answer>(.*?)</answer>", response, re.DOTALL)
    answer_block = answer_match.group(0).strip() if answer_match else ""

    kf_match = re.search(r"<key_frame>(.*?)</key_frame>", response, re.DOTALL)
    kf_block = kf_match.group(0).strip() if kf_match else ""

    # ---- 2. 识别格式：新(think模式) vs 旧 ----
    # 新格式特征：Ref-Delta 在 </think> 之后（草稿在 <think> 内部）
    # 旧格式特征：Ref-Delta 在 <think> 内部
    think_match = re.search(r"<think>(.*?)</think>(.*)", response, re.DOTALL)
    is_new_format = False
    if think_match:
        after_think = think_match.group(2)
        if "Reference frame at" in after_think or "<time_range>" in after_think:
            is_new_format = True

    if is_new_format:
        # 新格式：丢弃 <think> 内的草稿，取 </think> 之后的 Ref-Delta 作为内容
        after_think = think_match.group(2)
        content = re.sub(r"<answer>.*?</answer>", "", after_think, flags=re.DOTALL)
        content = re.sub(r"<key_frame>.*?</key_frame>", "", content, flags=re.DOTALL)
        content = content.strip()
        fixes.append("think模式新格式: 去除内部草稿，保留Ref-Delta")
        n_think_open = n_think_close = 1
    else:
        # 旧格式：去掉所有 <think></think> 标签，内容原样保留
        body = response
        body = re.sub(r"<answer>.*?</answer>", "", body, flags=re.DOTALL)
        body = re.sub(r"<key_frame>.*?</key_frame>", "", body, flags=re.DOTALL)
        n_think_open = len(re.findall(r"<think>", body))
        n_think_close = len(re.findall(r"</think>", body))
        content = re.sub(r"</?think>", "", body).strip()

    # ---- 3. 清理 think 内容 ----
    # 3a. 把 bullet list 转为散文 (去掉行首 - • * 符号)
    if re.search(r"^\s*[-•*]\s", content, re.MULTILINE):
        lines = content.split("\n")
        cleaned_lines = []
        for line in lines:
            stripped = re.sub(r"^\s*[-•*]\s+", "", line)
            if stripped != line:
                fixes.append(f"去除bullet: {line.strip()[:60]}")
            cleaned_lines.append(stripped)
        content = "\n".join(cleaned_lines)

    # 3b. 把 "At <time>..." 句式统一为 Reference frame 句式
    def _fix_at_time(m):
        full = m.group(0)
        t = m.group(1)
        rest = m.group(2).strip().rstrip(".")
        fixes.append(f"At <time> → Reference frame: {full[:60]}")
        return f"Reference frame at <time>{t}</time>: {rest}."
    content = re.sub(
        r"(?<!Reference frame )At <time>(.*?)</time>[,:]\s*(.*?)(?=\.\s|\n\n|$)",
        _fix_at_time, content, flags=re.DOTALL,
    )

    # ---- 4. 记录修复 ----
    if not is_new_format and (n_think_open > 1 or n_think_close > 1):
        fixes.append(f"多余think标签: {n_think_open}个<think>, {n_think_close}个</think>")

    if not is_new_format:
        think_matches = re.findall(r"<think>(.*?)</think>", response, re.DOTALL)
        all_empty = all(m.strip() == "" for m in think_matches) if think_matches else True
        if all_empty and content:
            fixes.append("think为空,内容在外部,已合并")

    # ---- 5. 清理 answer ----
    if gt_answer and answer_block:
        ans_inner = re.search(r"<answer>(.*?)</answer>", answer_block, re.DOTALL)
        if ans_inner:
            ans_text = ans_inner.group(1).strip()
            if ans_text != gt_answer and len(ans_text) > len(gt_answer):
                fixes.append(f"answer冗余: \"{ans_text[:50]}\" → \"{gt_answer}\"")
                answer_block = f"<answer>{gt_answer}</answer>"

    # ---- 6. 重新组装 ----
    parts = []
    if content:
        parts.append(f"<think>\n{content}\n</think>")
    else:
        parts.append("<think>\n</think>")

    if answer_block:
        parts.append(answer_block)
    if kf_block:
        parts.append(kf_block)

    cleaned = "\n\n".join(parts)
    return cleaned, fixes


def _extract_option_letter(s: str) -> str:
    """从答案文本中提取选项字母 A-E，失败返回空串"""
    m = re.match(r"^\(?([A-E])\)?[\.\s]?", s.strip())
    return m.group(1) if m else ""


_PLACEHOLDERS = {
    "concise answer", "the answer", "answer", "...", "……",
    "yes", "no", "n/a", "none", "unknown",
}


def answers_match(model_ans: str, gt: str) -> bool:
    """答案匹配，分两套策略：
    - 选择题（双方均可提取 A-E 字母）：严格比较字母
    - 自由文本：只过滤占位符/空答/极短答案，其余保留
      （因为 GT 已作为引导传给模型，有实质内容即视为正确）
    """
    if not model_ans or not gt:
        return False
    if model_ans == gt:
        return True

    # ---- 选择题：提取字母比较 ----
    ml = _extract_option_letter(model_ans)
    gl = _extract_option_letter(gt)
    if ml and gl:
        return ml == gl

    # ---- 自由文本：过滤明显失败 ----
    mn = model_ans.strip().lower()
    # 占位符
    if mn in _PLACEHOLDERS:
        return False
    # 极短（不足 5 字符）
    if len(mn) < 5:
        return False
    # 纯省略号 / 重复符号
    if re.match(r"^[.\s…]+$", mn):
        return False

    # 有实质内容 → 保留（GT 已引导模型，不做严格文本比对）
    return True


def validate_item(item: dict) -> list[str]:
    """验证单条数据的必要字段"""
    warnings = []
    resp = item.get("response", "")

    if not resp:
        warnings.append("response为空")
        return warnings

    # 检查 think 内容
    think_match = re.search(r"<think>(.*?)</think>", resp, re.DOTALL)
    think_text = think_match.group(1).strip() if think_match else ""
    if not think_text:
        warnings.append("think内容为空")

    # 检查 answer
    if not re.search(r"<answer>", resp):
        warnings.append("缺少answer")

    # 检查 Reference frame
    if not re.search(r"Reference frame at", resp):
        warnings.append("无Reference frame")

    # 检查 delta
    if not re.search(r"<delta>", resp):
        warnings.append("无delta")

    # 检查 time 标签
    n_time = len(re.findall(r"<time>", resp))
    n_time_range = len(re.findall(r"<time_range>", resp))
    if n_time + n_time_range == 0:
        warnings.append("无时间标签")

    # 检查 answer 是否与 GT 一致（宽松匹配）
    gt = item.get("answer", "")
    if gt:
        ans_match = re.search(r"<answer>(.*?)</answer>", resp, re.DOTALL)
        model_ans = ans_match.group(1).strip() if ans_match else ""
        if model_ans and not answers_match(model_ans, gt):
            warnings.append(f"答案错误: model={model_ans[:40]} gt={gt[:40]}")

    return warnings


def main():
    parser = argparse.ArgumentParser(description="清洗 CoT 数据")
    parser.add_argument("input", help="输入 JSONL 文件路径")
    parser.add_argument("-o", "--output", help="输出路径 (不指定则覆盖原文件)")
    parser.add_argument("--filter-wrong", action="store_true",
                        help="过滤掉答案与GT不一致的条目")
    args = parser.parse_args()

    input_path = Path(args.input)
    output_path = Path(args.output) if args.output else input_path

    # 读取
    items = []
    with open(input_path, encoding="utf-8") as f:
        for line in f:
            if line.strip():
                items.append(json.loads(line))

    print(f"读取 {len(items)} 条数据\n")

    # 清洗
    total_fixes = 0
    total_warnings = 0
    cleaned_items = []

    for i, item in enumerate(items):
        pid = item.get("problem_id", i)
        resp = item.get("response", "")

        if not resp:
            cleaned_items.append(item)
            continue

        # 清洗 response
        gt_answer = item.get("answer", "")
        cleaned_resp, fixes = clean_response(resp, gt_answer)
        item["response"] = cleaned_resp

        # 验证
        warnings = validate_item(item)

        if fixes or warnings:
            print(f"  [{pid}]", end="")
            if fixes:
                total_fixes += len(fixes)
                print(f"  修复: {'; '.join(fixes)}", end="")
            if warnings:
                total_warnings += len(warnings)
                print(f"  ⚠ {'; '.join(warnings)}", end="")
            print()

        cleaned_items.append(item)

    # 过滤答案错误的条目
    n_filtered = 0
    if args.filter_wrong:
        filtered = []
        for item in cleaned_items:
            gt = item.get("answer", "")
            resp = item.get("response", "")
            ans_m = re.search(r"<answer>(.*?)</answer>", resp, re.DOTALL)
            model_ans = ans_m.group(1).strip() if ans_m else ""
            if not gt or answers_match(model_ans, gt):
                filtered.append(item)
            else:
                n_filtered += 1
        cleaned_items = filtered

    # 写入
    with open(output_path, "w", encoding="utf-8") as f:
        for item in cleaned_items:
            f.write(json.dumps(item, ensure_ascii=False) + "\n")

    # 报告
    print(f"\n{'='*50}")
    print(f"清洗完成:")
    print(f"  总条数:   {len(cleaned_items)}")
    print(f"  修复数:   {total_fixes}")
    print(f"  警告数:   {total_warnings}")
    if args.filter_wrong:
        print(f"  过滤错误: {n_filtered}")
    print(f"  输出到:   {output_path}")


if __name__ == "__main__":
    main()