"""用本地大 Qwen-VL 模型（vLLM 离线批推）为 listwise reranker 训练集构造 CoT 数据。

输入: 训练集 jsonl（格式同 train/train_reranker_qwen3vl.py 所用数据）
输出: jsonl，每行 = 原样本 + 实际使用的 docs 顺序 + cot（思考过程）+ cot_answer（标签数字串）

本地用法见 cot_gen/run_gen_cot.sh
Nebula 用法见 scripts/submit/submit_cot_gen.sh：
    nebula 按每卡一进程拉起本脚本，WORLD_SIZE > 1 时自动按 rank 分片，
    每个进程用 tp=1 独占一张卡，各自写 <output>.shard<rank>-<world>.jsonl，跑完再合并。
"""
import os

os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")

import argparse
import glob
import json
import random
import re
import sys

current_file_path = os.path.dirname(os.path.abspath(__file__))
sys.path.append(os.path.join(current_file_path, "../"))


def get_dist_env():
    rank = int(os.environ.get("RANK", 0))
    world_size = int(os.environ.get("WORLD_SIZE", 1))
    local_rank = int(os.environ.get("LOCAL_RANK", rank))
    return rank, world_size, local_rank


def bind_single_gpu(local_rank: int):
    """分片模式下每个进程只用一张卡；调用点必须早于 torch/vllm 的 import。"""
    visible = os.environ.get("CUDA_VISIBLE_DEVICES")
    if visible:
        gpu_ids = [x for x in visible.split(",") if x]
    else:
        gpu_ids = [str(i) for i in sorted(int(re.findall(r"\d+$", p)[0]) for p in glob.glob("/dev/nvidia[0-9]*"))]
    if len(gpu_ids) > 1:
        os.environ["CUDA_VISIBLE_DEVICES"] = gpu_ids[local_rank % len(gpu_ids)]
    print(f"[rank {local_rank}] CUDA_VISIBLE_DEVICES={os.environ.get('CUDA_VISIBLE_DEVICES')}")


def extract_prompt(full_prompt: str, pattern: str):
    match = re.search(pattern, full_prompt, re.DOTALL)
    return match.group(1).strip() if match else None


def load_data(data_path: str):
    if data_path.endswith(".json"):
        return json.load(open(data_path))
    return [json.loads(line) for line in open(data_path)]


def select_docs(item, num_docs: int, shuffle: bool, rng: random.Random):
    docs = list(item["docs"])
    if shuffle:
        rng.shuffle(docs)
    if num_docs > 0:
        docs = docs[:num_docs]
    return docs


def build_cot_messages(sys_prompt, user_prompt, item, docs, max_pixels):
    """将 user_prompt 中的 <image>/<instruction>/<doc_list>/<gt_labels> 占位符展开成 qwen-vl messages。"""
    labels = [doc["label"] for doc in docs]
    gt_str = "".join(str(l) for l in labels)

    content = []
    for segment in re.split(r"(<image>|<instruction>|<doc_list>|<gt_labels>)", user_prompt):
        if segment == "<image>":
            content.append({"type": "image", "image": item["query_image"], "max_pixels": max_pixels})
        elif segment == "<instruction>":
            content.append({"type": "text", "text": item.get("instruction", "")})
        elif segment == "<doc_list>":
            for i, doc in enumerate(docs):
                content.append({"type": "text", "text": f"\n{i}: "})
                content.append({"type": "image", "image": doc["crop_image"], "max_pixels": max_pixels})
                content.append({"type": "text", "text": f" {doc['doc_title']}"})
            content.append({"type": "text", "text": "\n"})
        elif segment == "<gt_labels>":
            lines = [f"{i}: {l}" for i, l in enumerate(labels)]
            content.append({"type": "text", "text": "\n".join(lines) + f"\n（拼接后应为：{gt_str}）"})
        else:
            content.append({"type": "text", "text": segment})

    return [
        {"role": "system", "content": sys_prompt},
        {"role": "user", "content": content},
    ]


ANSWER_RE = re.compile(r"最终答案[：:]\s*([0-3]+)")


def parse_output(text: str):
    """返回 (cot, answer)，解析失败返回 (None, None)。"""
    matches = list(ANSWER_RE.finditer(text))
    if not matches:
        return None, None
    last = matches[-1]
    answer = last.group(1)

    # 关闭 thinking 后模型直接输出正文，chat template 仍会留下一个空 think 块
    cot = text[: last.start()].replace("<think>", "").replace("</think>", "").strip()
    if "【逐项分析】" not in cot:
        return None, answer
    return cot, answer


LEAK_RE = re.compile(
    r"参考标签|参考答案|标注员|给定的?标签|给定的?答案|已知标签|正确答案|所给的?标签|提供的标签|ground\s*truth",
    re.IGNORECASE,
)


def is_valid(cot, answer, gt_str):
    if cot is None or answer is None:
        return False
    if answer != gt_str:
        return False
    if LEAK_RE.search(cot):
        return False
    return True


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--model-path", type=str, required=True, help="本地大 qwen-vl 模型权重路径")
    parser.add_argument("--train-data-path", type=str, required=True)
    parser.add_argument("--output-path", type=str, required=True)
    parser.add_argument("--failed-path", type=str, default=None, help="多次重试仍失败的样本记录，默认 output_path + .failed")
    parser.add_argument("--prompt-path", type=str, default=os.path.join(current_file_path, "cot_prompt.jinja"))
    parser.add_argument("--num-docs", type=int, default=10, help="每个样本使用的候选数，<=0 表示全部")
    parser.add_argument("--shuffle-docs", action="store_true", help="是否先 shuffle 再截取（与训练随机采样对齐）")
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--max-crop-tokens", type=int, default=80, help="单图最大视觉 token 数，与训练配置对齐")
    parser.add_argument("--patch-size", type=int, default=16)
    parser.add_argument("--tensor-parallel-size", type=int, default=8)
    parser.add_argument("--max-model-len", type=int, default=16384)
    parser.add_argument("--max-new-tokens", type=int, default=4096)
    parser.add_argument("--gpu-memory-utilization", type=float, default=0.9)
    parser.add_argument("--max-retries", type=int, default=3, help="答案与 GT 不一致时的重试轮数")
    parser.add_argument("--retry-temperature", type=float, default=0.7)
    parser.add_argument("--batch-size", type=int, default=64, help="每轮送入 vLLM 的样本数")
    parser.add_argument("--shard-id", type=int, default=None, help="不传则取环境变量 RANK")
    parser.add_argument("--num-shards", type=int, default=None, help="不传则取环境变量 WORLD_SIZE")
    parser.add_argument("--test", type=int, nargs="?", const=20, default=None,
                        help="试跑模式：只处理数据集前 N 条（默认 20），结果写到 <output>.test.jsonl")
    args = parser.parse_args()

    rank, world_size, local_rank = get_dist_env()
    if args.num_shards is None:
        args.num_shards = world_size
    if args.shard_id is None:
        args.shard_id = rank
    if args.num_shards > 1:
        # 每进程独占一张卡，必须早于 vllm import
        bind_single_gpu(local_rank)
        args.tensor_parallel_size = 1

    from vllm import LLM, SamplingParams
    from transformers import AutoProcessor
    from collators.vision_process import process_vision_info

    if args.test:
        root, ext = os.path.splitext(args.output_path)
        args.output_path = f"{root}.test{ext}"
    if args.num_shards > 1:
        root, ext = os.path.splitext(args.output_path)
        args.output_path = f"{root}.shard{args.shard_id}-{args.num_shards}{ext}"
    failed_path = args.failed_path or args.output_path + ".failed"
    os.makedirs(os.path.dirname(os.path.abspath(args.output_path)), exist_ok=True)

    with open(args.prompt_path, encoding="utf-8") as f:
        full_prompt = f.read()
    sys_prompt = extract_prompt(full_prompt, r"<sys_start>(.*?)<sys_end>")
    user_prompt = extract_prompt(full_prompt, r"<user_start>(.*?)<user_end>")
    assert sys_prompt and user_prompt, "prompt 模板缺少 <sys_start>/<user_start> 块"

    max_pixels = args.max_crop_tokens * (args.patch_size * 2) ** 2

    data = load_data(args.train_data_path)
    print(f"Load {len(data)} instances from {args.train_data_path}")
    if args.test:
        data = data[: args.test]
        print(f"[TEST MODE] only generate for first {len(data)} samples -> {args.output_path}")

    done_ids = set()
    if os.path.exists(args.output_path):
        for line in open(args.output_path):
            try:
                done_ids.add(json.loads(line)["sample_idx"])
            except Exception:
                pass
        print(f"Resume: {len(done_ids)} samples already done")

    rng = random.Random(args.seed)
    pending = []  # (sample_idx, item, docs, gt_str, messages)
    for idx, item in enumerate(data):
        docs = select_docs(item, args.num_docs, args.shuffle_docs, rng)  # 先采样保证各分片 rng 状态一致
        if idx % args.num_shards != args.shard_id or idx in done_ids:
            continue
        gt_str = "".join(str(doc["label"]) for doc in docs)
        messages = build_cot_messages(sys_prompt, user_prompt, item, docs, max_pixels)
        pending.append((idx, item, docs, gt_str, messages))
    print(f"Shard {args.shard_id}/{args.num_shards}: {len(pending)} samples to generate")
    if not pending:
        return

    max_images = 1 + max(len(d) for _, _, d, _, _ in pending)
    llm = LLM(
        model=args.model_path,
        tensor_parallel_size=args.tensor_parallel_size,
        max_model_len=args.max_model_len,
        gpu_memory_utilization=args.gpu_memory_utilization,
        limit_mm_per_prompt={"image": max_images},
        trust_remote_code=True,
    )
    processor = AutoProcessor.from_pretrained(args.model_path, trust_remote_code=True)

    def to_vllm_input(messages):
        text = processor.apply_chat_template(
            messages, tokenize=False, add_generation_prompt=True, enable_thinking=False
        )
        image_inputs, _ = process_vision_info([messages])
        return {"prompt": text, "multi_modal_data": {"image": image_inputs}}

    out_f = open(args.output_path, "a", encoding="utf-8")
    failed_f = open(failed_path, "a", encoding="utf-8")

    n_ok, n_fail = 0, 0
    for batch_start in range(0, len(pending), args.batch_size):
        batch = pending[batch_start: batch_start + args.batch_size]
        results = {}  # sample_idx -> (cot, answer)
        last_text = {}  # sample_idx -> 最后一次尝试的模型原始输出

        remaining = list(batch)
        for attempt in range(args.max_retries + 1):
            if not remaining:
                break
            temperature = 0.0 if attempt == 0 else args.retry_temperature
            sampling_params = SamplingParams(
                temperature=temperature,
                top_p=0.95 if temperature > 0 else 1.0,
                max_tokens=args.max_new_tokens,
                seed=args.seed + attempt,
            )
            vllm_inputs = [to_vllm_input(m) for _, _, _, _, m in remaining]
            outputs = llm.generate(vllm_inputs, sampling_params)

            still_failed = []
            for (idx, item, docs, gt_str, messages), out in zip(remaining, outputs):
                text = out.outputs[0].text
                last_text[idx] = text
                cot, answer = parse_output(text)
                if is_valid(cot, answer, gt_str):
                    results[idx] = (cot, answer)
                else:
                    still_failed.append((idx, item, docs, gt_str, messages))
            remaining = still_failed
            if remaining:
                print(f"[batch {batch_start}] attempt {attempt}: {len(remaining)} samples need retry")

        for idx, item, docs, gt_str, messages in batch:
            record = dict(item)
            record["docs"] = docs  # 保存实际使用的 docs 顺序，训练时 CoT 与顺序绑定
            record["sample_idx"] = idx
            if idx in results:
                cot, answer = results[idx]
                record["cot"] = cot
                record["cot_answer"] = answer
                out_f.write(json.dumps(record, ensure_ascii=False) + "\n")
                n_ok += 1
            else:
                record["gt_answer"] = gt_str
                record["raw_output"] = last_text.get(idx)
                failed_f.write(json.dumps(record, ensure_ascii=False) + "\n")
                n_fail += 1
        out_f.flush()
        failed_f.flush()
        print(f"Progress: {batch_start + len(batch)}/{len(pending)}, ok={n_ok}, fail={n_fail}")

    out_f.close()
    failed_f.close()
    print(f"Done. success={n_ok}, failed={n_fail}, output={args.output_path}")


if __name__ == "__main__":
    main()