"""
Generate CoT data for embedding model using local Qwen3.5-397B-A17B via vLLM.

Usage:
    # 本地测试（前100条）
    python cot_gen/generate_cot.py --test

    # 本地全量运行（断点续传）
    python generate_cot.py --resume

    # 指定范围（分片输出）
    python cot_gen/generate_cot.py --start_idx 0 --end_idx 1000 --shard

    # 通过 nebula 提交（自动分片，结束后运行 merge_shards.py 合并）
    sh run_nebula.sh
"""

import os
import json
import argparse
import random
from pathlib import Path

from PIL import Image
from tqdm import tqdm
from vllm import LLM, SamplingParams

from config import (
    MODEL_PATH, TENSOR_PARALLEL_SIZE, GPU_MEMORY_UTILIZATION, MAX_MODEL_LEN,
    INPUT_PATH, OUTPUT_DIR, OUTPUT_PATH, FAILED_LOG_PATH,
    IMAGE_CDN_PREFIX, IMAGE_MAX_PIXELS, PROMPT_SELECT_MODE,
    BATCH_SIZE, MAX_TOKENS, TEMPERATURE, TOP_P, SAVE_INTERVAL,
    SYSTEM_PROMPT, COT_PROMPT_TEMPLATE,
)


def load_image_pair(image_name: str, bbox: str) -> tuple[Image.Image, Image.Image]:
    """Load original image and cropped image from local filesystem."""
    image_path = os.path.join(IMAGE_CDN_PREFIX, image_name)
    original = Image.open(image_path)
    if original.mode != "RGB":
        original = original.convert("RGB")

    if bbox and bbox != "ori":
        parts = bbox.split()
        if len(parts) >= 4:
            xmin, ymin, xmax, ymax = [int(float(z)) for z in parts[:4]]
            xmin = max(0, xmin)
            xmax = min(original.width - 1, xmax)
            ymin = max(0, ymin)
            ymax = min(original.height - 1, ymax)
            if (ymax - ymin) > 2 and (xmax - xmin) > 2:
                cropped = original.crop((xmin, ymin, xmax, ymax))
            else:
                cropped = original.copy()
        else:
            cropped = original.copy()
    else:
        cropped = original.copy()

    return original, cropped


def parse_records(raw_data: list[dict]) -> list[dict]:
    """
    Parse raw records into individual samples.
    Behavior depends on PROMPT_SELECT_MODE:
      - "all": each prompt in pos_prompts becomes a separate sample (data augmentation)
      - "random_one": randomly pick one prompt per record
    """
    all_samples = []
    for record in raw_data:
        query_tfsid = record["query_tfsid"]
        query_bbox = record["query_bbox"]
        rn = record["rn"]

        pos_prompts = record["pos_prompts"].split("##")
        pos_titles_list = record["pos_item_titles"].split("##")
        pos_tfsids_list = record["pos_item_tfsids"].split("##")
        pos_bboxs_list = record["pos_item_bboxs"].split("##")

        neg_tfsids_list = [x for x in record.get("neg_item_tfsids", "").split("##") if x.strip()]
        neg_titles_list = [x for x in record.get("neg_item_titles", "").split("##") if x.strip()]
        neg_bboxs_list = [x for x in record.get("neg_item_bboxs", "").split("##") if x.strip()]

        def pick_random_neg():
            if not neg_tfsids_list:
                return "", "", ""
            idx = random.randint(0, len(neg_tfsids_list) - 1)
            tfsid = neg_tfsids_list[idx]
            title = neg_titles_list[idx] if idx < len(neg_titles_list) else ""
            bbox = neg_bboxs_list[idx] if idx < len(neg_bboxs_list) else ""
            return tfsid, title, bbox

        def build_sample(i):
            title_indices = [max(0, i - 1), i, min(len(pos_titles_list) - 1, i + 1)]
            ref_titles = list(dict.fromkeys(
                pos_titles_list[idx] for idx in title_indices if idx < len(pos_titles_list)
            ))
            neg_tfsid, neg_title, neg_bbox = pick_random_neg()
            return {
                "rn": rn,
                "query_tfsid": query_tfsid,
                "query_bbox": query_bbox,
                "instruction": pos_prompts[i].strip(),
                "pos_titles_ref": ref_titles,
                "prompt_index": i,
                "pos_item_tfsid": pos_tfsids_list[i] if i < len(pos_tfsids_list) else "",
                "pos_item_title": pos_titles_list[i] if i < len(pos_titles_list) else "",
                "pos_item_bbox": pos_bboxs_list[i] if i < len(pos_bboxs_list) else "",
                "neg_item_tfsid": neg_tfsid,
                "neg_item_title": neg_title,
                "neg_item_bbox": neg_bbox,
            }

        if PROMPT_SELECT_MODE == "random_one":
            valid_indices = [i for i, p in enumerate(pos_prompts) if p.strip()]
            if not valid_indices:
                continue
            i = random.choice(valid_indices)
            all_samples.append(build_sample(i))
        else:
            for i, prompt in enumerate(pos_prompts):
                if not prompt.strip():
                    continue
                all_samples.append(build_sample(i))

    return all_samples


def build_vllm_prompt(sample: dict) -> str:
    """Build the text prompt part for vLLM (image is passed separately)."""
    instruction = sample["instruction"]
    pos_titles = "\n".join(f"- {t}" for t in sample["pos_titles_ref"])

    user_text = COT_PROMPT_TEMPLATE.format(
        instruction=instruction,
        pos_titles=pos_titles,
    )

    return user_text


def build_messages(sample: dict, original: Image.Image, cropped: Image.Image) -> list[dict]:
    """Build chat messages for vLLM with original and cropped images."""
    user_text = build_vllm_prompt(sample)

    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {
            "role": "user",
            "content": [
                {"type": "image_pil", "image_pil": original},
                {"type": "image_pil", "image_pil": cropped},
                {"type": "text", "text": user_text},
            ],
        },
    ]
    return messages


def load_processed_keys(output_path: str) -> set:
    """Load already processed (rn, prompt_index) pairs for resume."""
    processed = set()
    if os.path.exists(output_path):
        with open(output_path, "r", encoding="utf-8") as f:
            for line in f:
                if line.strip():
                    try:
                        item = json.loads(line)
                        processed.add((item["rn"], item["prompt_index"]))
                    except (json.JSONDecodeError, KeyError):
                        continue
    return processed


def save_results(results: list[dict], output_path: str):
    """Append results to JSONL file."""
    with open(output_path, "a", encoding="utf-8") as f:
        for item in results:
            f.write(json.dumps(item, ensure_ascii=False) + "\n")


def save_failed(failed_samples: list[dict], failed_path: str):
    """Append failed samples to log."""
    with open(failed_path, "a", encoding="utf-8") as f:
        for item in failed_samples:
            f.write(json.dumps(item, ensure_ascii=False) + "\n")


def main(args):
    # Load data
    print(f"Loading data from {INPUT_PATH} ...")
    with open(INPUT_PATH, "r", encoding="utf-8") as f:
        raw_data = json.load(f)
    print(f"Total raw records: {len(raw_data)}")

    # Slice
    end_idx = len(raw_data) if args.end_idx < 0 else min(args.end_idx, len(raw_data))
    raw_data = raw_data[args.start_idx:end_idx]
    print(f"Processing records [{args.start_idx}, {end_idx})")

    # Parse into samples
    all_samples = parse_records(raw_data)
    print(f"Prompt select mode: {PROMPT_SELECT_MODE}")
    print(f"Total samples (query x instruction pairs): {len(all_samples)}")

    # Output setup
    os.makedirs(OUTPUT_DIR, exist_ok=True)

    if args.shard:
        base, ext = os.path.splitext(OUTPUT_PATH)
        output_path = f"{base}_{args.start_idx}_{end_idx}{ext}"
        failed_path = f"{base}_{args.start_idx}_{end_idx}_failed{ext}"
    else:
        output_path = OUTPUT_PATH
        failed_path = FAILED_LOG_PATH

    # Resume
    if args.resume:
        processed = load_processed_keys(output_path)
        before = len(all_samples)
        all_samples = [
            s for s in all_samples
            if (s["rn"], s["prompt_index"]) not in processed
        ]
        print(f"Resume: skipping {before - len(all_samples)} processed, {len(all_samples)} remaining")

    if not all_samples:
        print("No samples to process. Done.")
        return

    # Initialize vLLM
    print(f"Loading model from {MODEL_PATH} ...")
    print(f"  tensor_parallel_size={args.tp or TENSOR_PARALLEL_SIZE}")
    print(f"  gpu_memory_utilization={GPU_MEMORY_UTILIZATION}")
    print(f"  image_dir={IMAGE_CDN_PREFIX}")

    llm = LLM(
        model=MODEL_PATH,
        tensor_parallel_size=args.tp or TENSOR_PARALLEL_SIZE,
        gpu_memory_utilization=GPU_MEMORY_UTILIZATION,
        max_model_len=MAX_MODEL_LEN,
        trust_remote_code=True,
        dtype="bfloat16",
        limit_mm_per_prompt={"image": 2},
        mm_processor_kwargs={"max_pixels": IMAGE_MAX_PIXELS},
        enable_prefix_caching=True,
    )
    sampling_params = SamplingParams(
        max_tokens=MAX_TOKENS,
        temperature=TEMPERATURE,
        top_p=TOP_P,
    )

    print(f"Model loaded. Starting generation...")
    print(f"Output: {output_path}")

    # Process in batches
    total_success = 0
    total_failed = 0
    batch_results = []

    for batch_start in tqdm(range(0, len(all_samples), BATCH_SIZE), desc="Batches"):
        batch_samples = all_samples[batch_start:batch_start + BATCH_SIZE]

        # Prepare batch inputs: load images and build prompts
        valid_samples = []
        batch_inputs = []

        for sample in batch_samples:
            try:
                original, cropped = load_image_pair(sample["query_tfsid"], sample["query_bbox"])
                messages = build_messages(sample, original, cropped)
                valid_samples.append(sample)
                batch_inputs.append(messages)
            except Exception as e:
                total_failed += 1
                save_failed([{**sample, "error": f"image_load: {str(e)}"}], failed_path)

        if not batch_inputs:
            continue

        # vLLM batch inference using chat interface
        try:
            outputs = llm.chat(
                messages=batch_inputs,
                sampling_params=sampling_params,
                chat_template_kwargs={"enable_thinking": False},
            )

            for sample, output in zip(valid_samples, outputs):
                generated_text = output.outputs[0].text.strip()

                if generated_text and len(generated_text) >= 10:
                    result = {
                        "rn": sample["rn"],
                        "query_tfsid": sample["query_tfsid"],
                        "query_bbox": sample["query_bbox"],
                        "instruction": sample["instruction"],
                        "cot": generated_text,
                        "prompt_index": sample["prompt_index"],
                        "pos_item_tfsid": sample["pos_item_tfsid"],
                        "pos_item_title": sample["pos_item_title"],
                        "pos_item_bbox": sample["pos_item_bbox"],
                        "neg_item_tfsid": sample["neg_item_tfsid"],
                        "neg_item_title": sample["neg_item_title"],
                        "neg_item_bbox": sample["neg_item_bbox"],
                    }
                    batch_results.append(result)
                    total_success += 1
                else:
                    total_failed += 1
                    save_failed([{**sample, "error": f"short_output: {generated_text}"}], failed_path)

        except Exception as e:
            print(f"Batch inference error: {e}")
            total_failed += len(valid_samples)
            for sample in valid_samples:
                save_failed([{**sample, "error": f"inference: {str(e)}"}], failed_path)

        # Periodic save
        if len(batch_results) >= SAVE_INTERVAL:
            save_results(batch_results, output_path)
            batch_results = []

    # Save remaining
    if batch_results:
        save_results(batch_results, output_path)

    print(f"\nDone! Success: {total_success}, Failed: {total_failed}")
    print(f"Output: {output_path}")
    if total_failed > 0:
        print(f"Failed log: {failed_path}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Generate CoT data using local Qwen3.5-397B-A17B")
    parser.add_argument("--test", action="store_true", help="Test mode: process first 100 records")
    parser.add_argument("--start_idx", type=int, default=0, help="Start index of raw records")
    parser.add_argument("--end_idx", type=int, default=-1, help="End index (-1 for all)")
    parser.add_argument("--tp", type=int, default=None, help="Override tensor_parallel_size")
    parser.add_argument("--resume", action="store_true", help="Resume from previous progress")
    parser.add_argument("--shard", action="store_true", help="Shard output by index range")
    args = parser.parse_args()

    if args.test:
        args.start_idx = 0
        args.end_idx = 100
        args.resume = False
    else:
        # Auto-shard based on RANK/WORLD_SIZE env vars (set by nebula/distributed launcher)
        rank = int(os.environ.get("RANK", os.environ.get("INDEX", 0)))
        world_size = int(os.environ.get("WORLD_SIZE", os.environ.get("WORKER_NUM", 1)))
        if world_size > 1:
            with open(INPUT_PATH, "r", encoding="utf-8") as f:
                total_records = len(json.load(f))
            if args.end_idx < 0:
                args.end_idx = total_records
            total = min(args.end_idx, total_records) - args.start_idx
            shard_size = (total + world_size - 1) // world_size
            args.start_idx = args.start_idx + rank * shard_size
            args.end_idx = min(args.start_idx + shard_size, total_records)
            args.shard = True
            print(f"[Auto-shard] rank={rank}, world_size={world_size}, "
                  f"processing [{args.start_idx}, {args.end_idx})")

    main(args)