"""
Retry failed records using local vLLM inference.

Usage:
    python retry_failed.py --failed_path /path/to/failed_records.jsonl
"""

import os
import json
import argparse

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,
    BATCH_SIZE, MAX_TOKENS, TEMPERATURE, TOP_P,
    SYSTEM_PROMPT, COT_PROMPT_TEMPLATE,
)
from generate_cot import (
    load_image_pair, build_messages, save_results,
)


def main(args):
    # Load failed records
    failed_samples = []
    with open(args.failed_path, "r", encoding="utf-8") as f:
        for line in f:
            if line.strip():
                try:
                    item = json.loads(line)
                    item.pop("error", None)
                    failed_samples.append(item)
                except json.JSONDecodeError:
                    continue

    print(f"Loaded {len(failed_samples)} failed records to retry")
    if not failed_samples:
        return

    base, ext = os.path.splitext(args.failed_path)
    output_path = base.replace("_failed", "") + "_retried" + ext
    still_failed_path = base + "_still_failed" + ext

    # Init
    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},
    )
    sampling_params = SamplingParams(
        max_tokens=MAX_TOKENS,
        temperature=TEMPERATURE,
        top_p=TOP_P,
    )

    total_success = 0
    total_failed = 0
    results = []
    still_failed = []

    for batch_start in tqdm(range(0, len(failed_samples), BATCH_SIZE), desc="Retrying"):
        batch = failed_samples[batch_start:batch_start + BATCH_SIZE]

        valid_samples = []
        batch_inputs = []
        for sample in batch:
            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:
                still_failed.append(sample)
                total_failed += 1

        if not batch_inputs:
            continue

        try:
            outputs = llm.chat(messages=batch_inputs, sampling_params=sampling_params)
            for sample, output in zip(valid_samples, outputs):
                text = output.outputs[0].text.strip()
                if text and len(text) >= 10:
                    results.append({
                        "rn": sample["rn"],
                        "query_tfsid": sample["query_tfsid"],
                        "query_bbox": sample["query_bbox"],
                        "instruction": sample["instruction"],
                        "cot": text,
                        "prompt_index": sample["prompt_index"],
                    })
                    total_success += 1
                else:
                    still_failed.append(sample)
                    total_failed += 1
        except Exception:
            still_failed.extend(valid_samples)
            total_failed += len(valid_samples)

    if results:
        save_results(results, output_path)
    if still_failed:
        with open(still_failed_path, "w", encoding="utf-8") as f:
            for item in still_failed:
                f.write(json.dumps(item, ensure_ascii=False) + "\n")

    print(f"\nRetry done! Success: {total_success}, Still failed: {total_failed}")
    if total_success > 0:
        print(f"Results: {output_path}")
    if total_failed > 0:
        print(f"Still failed: {still_failed_path}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--failed_path", type=str, required=True)
    parser.add_argument("--tp", type=int, default=None)
    args = parser.parse_args()
    main(args)