"""
CoT 蒸馏数据采集脚本（实验 8-9 配套代码）

方法（对应书中实验 8-9 的三步流程之第一步"采集轨迹"）：
  1. 从 problems.jsonl 读取带标准答案的数学题（规则可验证的任务分布）；
  2. 通过 OpenRouter 调用前沿教师模型（默认 Claude），开启 reasoning 获取
     完整"思考 + 答案"轨迹（Claude 4 系列返回的是 summarized thinking——由另一个
     模型对原始思维链做的高保真摘要，原始思维链只存在于加密的 signature 字段中）；
  3. 用规则验证器核对最终答案，只把答对的轨迹写成 SFT 训练数据
     （"问题 → <think>思考</think> + 最终答案" 的 messages 格式）。

注意：本实验只使用各厂商官方 API 的 reasoning/thinking 能力获取思维链，
不涉及任何绕过厂商安全机制的手段。原始轨迹（含未通过验证的）保存在
raw_trajectories.jsonl，便于分析教师的错误模式。
"""

import argparse
import asyncio
import json
import os
import re
from pathlib import Path
from typing import Optional

from openai import AsyncOpenAI

ANSWER_SUFFIX = "\n\n请一步步推理，并在最后一行用「Final Answer: 数值」的格式给出最终答案（只写数值，不带单位）。"

def load_problems(path: str) -> list[dict]:
    problems = []
    with open(path, "r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if line:
                problems.append(json.loads(line))
    return problems


def load_jsonl(path: str | Path) -> list[dict]:
    path = Path(path)
    if not path.is_file():
        return []
    with path.open("r", encoding="utf-8") as f:
        return [json.loads(line) for line in f if line.strip()]


def write_jsonl_atomic(path: str | Path, rows: list[dict]) -> None:
    """Replace a JSONL file without exposing a partially rewritten dataset."""
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    temporary = path.with_name(f".{path.name}.{os.getpid()}.tmp")
    try:
        with temporary.open("w", encoding="utf-8") as f:
            for row in rows:
                f.write(json.dumps(row, ensure_ascii=False) + "\n")
            f.flush()
            os.fsync(f.fileno())
        os.replace(temporary, path)
    finally:
        temporary.unlink(missing_ok=True)


def records_in_problem_order(problems: list[dict], records_by_id: dict[str, dict]) -> list[dict]:
    """Return canonical problem order while retaining any legacy extra records."""
    known_ids = [str(problem["id"]) for problem in problems]
    rows = [records_by_id[problem_id] for problem_id in known_ids if problem_id in records_by_id]
    rows.extend(record for problem_id, record in records_by_id.items() if problem_id not in known_ids)
    return rows


def extract_predicted_number(text: str) -> Optional[float]:
    """从模型输出中解析最终答案数值。优先匹配 Final Answer 标记，否则取最后一个数字。"""
    m = re.findall(r"Final Answer[:：]\s*(-?[\d,]+(?:\.\d+)?)", text, re.IGNORECASE)
    if not m:
        m = re.findall(r"-?[\d,]+(?:\.\d+)?", text)
    if not m:
        return None
    try:
        return float(m[-1].replace(",", ""))
    except ValueError:
        return None


def verify(text: str, gold: float, tol: float = 1e-6) -> bool:
    """规则验证器：核对最终答案是否与标准答案一致。"""
    pred = extract_predicted_number(text)
    if pred is None:
        return False
    return abs(pred - float(gold)) <= tol * max(1.0, abs(float(gold)))


def get_reasoning(message) -> str:
    """从返回的 message 中提取思维链。

    依次尝试：OpenRouter 的 reasoning / reasoning_details 字段，
    以及 Moonshot、DeepSeek 等原生 API 的 reasoning_content 字段。
    """
    reasoning = getattr(message, "reasoning", None)
    if reasoning:
        return reasoning
    reasoning_content = getattr(message, "reasoning_content", None)
    if reasoning_content:
        return reasoning_content
    details = getattr(message, "reasoning_details", None) or []
    parts = []
    for d in details:
        if isinstance(d, dict):
            parts.append(d.get("text") or d.get("summary") or "")
        else:
            parts.append(getattr(d, "text", None) or getattr(d, "summary", None) or "")
    return "\n".join(p for p in parts if p)


def reasoning_extra_body(base_url: str, effort: str, max_tokens: int) -> dict:
    """Build the provider-specific reasoning control without silently ignoring it."""
    if effort:
        if "api.moonshot.cn" in base_url:
            # Moonshot's native OpenAI-compatible endpoint accepts the same
            # top-level control used by the Experiment 8-8 Kimi campaign.
            return {"reasoning_effort": effort}
        return {"reasoning": {"effort": effort}}
    if max_tokens:
        return {"reasoning": {"max_tokens": max_tokens}}
    return {}


async def distill_one(client: AsyncOpenAI, problem: dict, args, semaphore) -> dict:
    """对单道题调用教师模型，返回完整轨迹记录。"""
    record = {
        "id": problem["id"],
        "question": problem["question"],
        "gold_answer": problem["answer"],
        "model": args.model,
        "content": None,
        "reasoning": None,
        "verified": False,
        "usage": None,
        "error": None,
        "attempts": [],
    }
    async with semaphore:
        for attempt in range(args.max_retries + 1):
            try:
                kwargs = {}
                reasoning_body = reasoning_extra_body(
                    args.base_url, args.reasoning_effort, args.reasoning_max_tokens
                )
                if reasoning_body:
                    kwargs["extra_body"] = reasoning_body
                resp = await asyncio.wait_for(
                    client.chat.completions.create(
                        model=args.model,
                        messages=[{"role": "user", "content": problem["question"] + args.answer_suffix}],
                        max_tokens=args.max_tokens,
                        # 重试时升温换取不同轨迹；Kimi K3 等锁定 temperature=1 的模型除外
                        temperature=args.temperature + (0.2 * attempt if args.temperature < 1.0 else 0),
                        **kwargs,
                    ),
                    timeout=args.request_timeout,  # 硬超时：防止半开连接挂死
                )
                msg = resp.choices[0].message
                record["content"] = msg.content or ""
                record["reasoning"] = get_reasoning(msg)
                record["usage"] = resp.usage.model_dump() if resp.usage else None
                record["verified"] = verify(record["content"], problem["answer"])
                record["error"] = None
                record["attempts"].append({
                    "attempt": attempt,
                    "content": record["content"],
                    "reasoning": record["reasoning"],
                    "usage": record["usage"],
                    "verified": record["verified"],
                    "error": None,
                })
                if record["verified"]:
                    break
            except Exception as e:
                record["error"] = f"attempt {attempt}: {type(e).__name__}: {e}"
                record["attempts"].append({
                    "attempt": attempt,
                    "content": None,
                    "reasoning": None,
                    "usage": None,
                    "verified": False,
                    "error": record["error"],
                })
        status = "OK" if record["verified"] else ("ERR" if record["error"] else "WRONG")
        print(f"  [{status}] {record['id']}", flush=True)
        return record


def to_sft_sample(record: dict) -> dict:
    """把验证通过的轨迹转成 SFT 训练样本（messages 格式，思考包在 <think> 标签里）。"""
    if record["reasoning"]:
        assistant = f"<think>\n{record['reasoning'].strip()}\n</think>\n\n{record['content'].strip()}"
    else:
        assistant = record["content"].strip()
    return {
        "messages": [
            {"role": "user", "content": record["question"]},
            {"role": "assistant", "content": assistant},
        ]
    }


async def main():
    parser = argparse.ArgumentParser(
        description="用前沿云模型（经 OpenRouter）蒸馏 CoT 轨迹，生成 SFT 数据",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    parser.add_argument("--input", default="./problems.jsonl", help="题目文件（JSONL，含 question/answer）")
    parser.add_argument("--sft_output", default="./data/sft_cot_distill.jsonl", help="SFT 训练数据输出路径")
    parser.add_argument("--raw_output", default="./data/raw_trajectories.jsonl", help="原始轨迹（含失败样本）输出路径")
    parser.add_argument("--model", default="anthropic/claude-opus-4.8", help="教师模型 ID")
    parser.add_argument("--base_url", default="https://openrouter.ai/api/v1", help="OpenAI 兼容 API 端点")
    parser.add_argument("--api_key_env", default="OPENROUTER_API_KEY", help="存放 API Key 的环境变量名")
    parser.add_argument("--reasoning_effort", default="",
                        help="OpenRouter 风格 reasoning effort（如 high/medium/low；设置后优先于 --reasoning_max_tokens，"
                             "用于 Claude Opus 4.8 等只支持自适应思考的模型）")
    parser.add_argument("--reasoning_max_tokens", type=int, default=4096,
                        help="思维链最大 token 数（OpenRouter 风格 reasoning 参数；0 = 不传该参数，"
                             "用于 Moonshot/DeepSeek 等默认返回 reasoning_content 的原生 API）")
    parser.add_argument("--max_problems", type=int, default=0, help="最多处理多少题（0 = 全部，调试用）")
    parser.add_argument(
        "--problem-id",
        action="append",
        default=[],
        help="只运行指定题目 ID；可重复传入。用于定点重试而不重跑整套题",
    )
    parser.add_argument(
        "--resume",
        action="store_true",
        help="保留 raw_output 中已验证记录，只重试缺失或未验证题目，并原子更新数据集",
    )
    parser.add_argument("--concurrency", type=int, default=8, help="并发请求数")
    parser.add_argument("--temperature", type=float, default=0.3, help="采样温度")
    parser.add_argument("--max_tokens", type=int, default=8192, help="单条回复最大 token 数（须大于 reasoning tokens）")
    parser.add_argument("--max_retries", type=int, default=1, help="失败/出错后的最大重试次数")
    parser.add_argument("--request_timeout", type=float, default=600, help="单次请求超时（秒），超时后按失败重试")
    parser.add_argument("--answer_suffix", default=ANSWER_SUFFIX, help="附加在题目后的作答格式要求")
    args = parser.parse_args()

    api_key = os.environ.get(args.api_key_env)
    if not api_key:
        raise SystemExit(f"请先设置环境变量 {args.api_key_env}")

    all_problems = load_problems(args.input)
    problem_ids = {str(problem["id"]) for problem in all_problems}
    requested_ids = set(args.problem_id)
    unknown_ids = sorted(requested_ids - problem_ids)
    if unknown_ids:
        raise SystemExit(f"未知题目 ID: {', '.join(unknown_ids)}")
    problems = [
        problem for problem in all_problems
        if not requested_ids or str(problem["id"]) in requested_ids
    ]
    if args.max_problems:
        problems = problems[: args.max_problems]

    existing_rows = load_jsonl(args.raw_output) if args.resume else []
    records_by_id = {
        str(record["id"]): record for record in existing_rows if record.get("id") is not None
    }
    pending = [
        problem for problem in problems
        if not records_by_id.get(str(problem["id"]), {}).get("verified", False)
    ]
    print(
        f"选中 {len(problems)} 道题，待运行 {len(pending)} 道，"
        f"教师模型：{args.model} @ {args.base_url}"
    )

    client = AsyncOpenAI(base_url=args.base_url, api_key=api_key, timeout=args.request_timeout)
    semaphore = asyncio.Semaphore(args.concurrency)

    # 每题完成后原子替换：中断最多损失当前请求，不会破坏已有数据集。
    run_records = []
    tasks = [distill_one(client, p, args, semaphore) for p in pending]
    for coro in asyncio.as_completed(tasks):
        record = await coro
        previous = records_by_id.get(str(record["id"]))
        if previous and not previous.get("verified", False):
            prior_failures = list(previous.get("prior_failures") or [])
            prior_failures.append({
                "model": previous.get("model"),
                "verified": False,
                "error": previous.get("error"),
                "usage": previous.get("usage"),
            })
            record["prior_failures"] = prior_failures
        records_by_id[str(record["id"])] = record
        run_records.append(record)
        write_jsonl_atomic(
            args.raw_output,
            records_in_problem_order(all_problems, records_by_id),
        )

    records = records_in_problem_order(all_problems, records_by_id)
    write_jsonl_atomic(args.raw_output, records)
    passed = [record for record in records if record.get("verified", False)]
    write_jsonl_atomic(args.sft_output, [to_sft_sample(record) for record in passed])

    total_in = sum((r["usage"] or {}).get("prompt_tokens", 0) for r in run_records)
    total_out = sum((r["usage"] or {}).get("completion_tokens", 0) for r in run_records)
    n_err = sum(1 for r in run_records if r["error"])
    print(f"\n{'=' * 50}")
    # Empty problems JSONL yields zero records; avoid ZeroDivisionError on the rate.
    pass_rate = (len(passed) / len(records) * 100) if records else 0.0
    print(f"数据集验证通过：{len(passed)}/{len(records)}（{pass_rate:.1f}%）")
    print(
        f"本次请求：{len(run_records)}  API 最终出错：{n_err}  "
        f"无思维链返回：{sum(1 for r in run_records if not r['reasoning'])}"
    )
    print(f"本次 Token 消耗：输入 {total_in}，输出 {total_out}")
    print(f"SFT 数据已写入：{args.sft_output}")
    print(f"原始轨迹已写入：{args.raw_output}")


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