#!/usr/bin/env python3
"""
Identify hallucinated captions (LLaVA mentions target object where absent)
and rewrite them using Qwen3-8B via vLLM.

Pipeline:
  1. Regex pre-filter  (compound mode, wide net)
  2. LLM judge         (Qwen3-8B YES/NO on each regex match)
  3. LLM rewrite       (Qwen3-8B removes object mentions)
  4. Strict check      (negation-aware keyword detector on rewrites)

Usage:
    python EFUF/scripts/clean_captions.py --relation bathroom_toilet
    python EFUF/scripts/clean_captions.py --relation bathroom_toilet --dry_run
    CUDA_VISIBLE_DEVICES=0,1 python EFUF/scripts/clean_captions.py --relation kitchen_oven --n_gpus 2
"""

from __future__ import annotations

import argparse
import json
import os
import sys
import time

sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "../.."))

from experiment.config.relation_config import get_relation_config
from experiment.evaluation.metrics import KeywordMentionDetector


def load_data(path: str) -> list[dict]:
    with open(path) as f:
        return json.load(f)


def save_data(path: str, data: list[dict]) -> None:
    os.makedirs(os.path.dirname(os.path.abspath(path)), exist_ok=True)
    tmp = path + ".tmp"
    with open(tmp, "w") as f:
        json.dump(data, f, indent=2, ensure_ascii=False)
    os.replace(tmp, path)


def load_checkpoint(path: str) -> dict:
    ckpt_path = path + ".ckpt.json"
    if not os.path.exists(ckpt_path):
        return {}
    with open(ckpt_path) as f:
        return json.load(f)


def save_checkpoint(path: str, ckpt: dict) -> None:
    ckpt_path = path + ".ckpt.json"
    tmp = ckpt_path + ".tmp"
    with open(tmp, "w") as f:
        json.dump(ckpt, f, indent=2, ensure_ascii=False)
    os.replace(tmp, ckpt_path)


def create_llm(model: str, n_gpus: int):
    from vllm import LLM
    from transformers import AutoTokenizer
    tokenizer = AutoTokenizer.from_pretrained(model)
    llm = LLM(
        model=model,
        dtype="bfloat16",
        tensor_parallel_size=n_gpus,
        gpu_memory_utilization=0.30,
        max_model_len=1024,
        enforce_eager=True,
    )
    return llm, tokenizer


def step1_regex_filter(
    data: list[dict],
    rc,
    ckpt: dict,
) -> dict:
    if ckpt.get("step1_done"):
        print("[Step 1] Regex filter already done -- skipping.")
        return ckpt

    object_key = rc.object_key
    detector = KeywordMentionDetector(keywords=rc.mention_keywords, compound_mode=True)

    regex_matched = {}

    for entry in data:
        iid = entry["image_id"]
        if entry.get(object_key, 0) == 1:
            continue
        caption = entry.get("llava_caption", "")
        if detector.mentions_object(caption):
            regex_matched[iid] = caption

    existing_matched = ckpt.get("step1_regex_matched", {})
    existing_matched.update(regex_matched)
    ckpt["step1_regex_matched"] = existing_matched
    ckpt["step1_done"] = True

    total_obj0 = sum(1 for e in data if e.get(object_key, 0) == 0)
    print(f"[Step 1] Regex filter: {len(regex_matched)}/{total_obj0} object=0 entries matched")

    return ckpt


def step2_llm_judge(
    data: list[dict],
    rc,
    ckpt: dict,
    llm,
    tokenizer,
    batch_size: int,
) -> dict:
    if ckpt.get("step2_done"):
        print("[Step 2] LLM judge already done -- skipping.")
        return ckpt

    from vllm import SamplingParams

    regex_matched = ckpt.get("step1_regex_matched", {})
    already_judged = ckpt.get("step2_judge_done", {})

    to_judge = {
        iid: cap for iid, cap in regex_matched.items()
        if iid not in already_judged
    }

    if not to_judge:
        ckpt["step2_done"] = True
        print("[Step 2] No new entries to judge.")
        return ckpt

    print(f"[Step 2] Judging {len(to_judge)} entries ...")

    judge_sys = (
        "You are determining whether an image description mentions a specific physical object. "
        "Only answer YES if the description explicitly refers to the {judge_object_name} "
        "as a physical object in the scene — not as part of a compound word like "
        "\"toilet paper\", \"toilet seat\", etc. unless it actually refers to the object itself."
    ).format(judge_object_name=rc.judge_object_name)

    prompts = []
    iids_ordered = []
    for iid, cap in to_judge.items():
        user_msg = (
            f"Does the following image description mention a {rc.judge_object_name}?\n\n"
            f'Answer ONLY "YES" or "NO".\n\n'
            f"Description: {cap}"
        )
        messages = [
            {"role": "system", "content": judge_sys},
            {"role": "user", "content": user_msg},
        ]
        prompt_text = tokenizer.apply_chat_template(
            messages, tokenize=False, add_generation_prompt=True, enable_thinking=False
        )
        prompts.append(prompt_text)
        iids_ordered.append(iid)

    sampling = SamplingParams(temperature=0.0, max_tokens=5)

    results = {}
    for i in range(0, len(prompts), batch_size):
        batch_prompts = prompts[i : i + batch_size]
        batch_iids = iids_ordered[i : i + batch_size]
        outputs = llm.generate(batch_prompts, sampling, use_tqdm=True)
        for iid, out in zip(batch_iids, outputs):
            text = out.outputs[0].text.strip()
            results[iid] = "YES" if "YES" in text.upper() else "NO"

    already_judged.update(results)
    ckpt["step2_judge_done"] = already_judged
    ckpt["step2_done"] = True

    yes_count = sum(1 for v in results.values() if v == "YES")
    no_count = sum(1 for v in results.values() if v == "NO")
    print(f"[Step 2] Judge results: {yes_count} YES, {no_count} NO (out of {len(results)} judged)")

    return ckpt


def step3_llm_rewrite(
    rc,
    ckpt: dict,
    llm,
    tokenizer,
    batch_size: int,
) -> dict:
    if ckpt.get("step3_done"):
        print("[Step 3] LLM rewrite already done -- skipping.")
        return ckpt

    from vllm import SamplingParams

    regex_matched = ckpt.get("step1_regex_matched", {})
    judge_done = ckpt.get("step2_judge_done", {})
    already_rewritten = ckpt.get("step3_rewrite_done", {})

    yes_iids = [iid for iid, verdict in judge_done.items() if verdict == "YES"]
    to_rewrite = {
        iid: regex_matched[iid] for iid in yes_iids
        if iid not in already_rewritten
    }

    if not to_rewrite:
        ckpt["step3_done"] = True
        print("[Step 3] No new entries to rewrite.")
        return ckpt

    print(f"[Step 3] Rewriting {len(to_rewrite)} entries ...")

    rewrite_sys = (
        "You are a meticulous writing assistant. "
        "Rewrite the given image description, removing ALL mentions of {object} "
        "as if the object was never present in the scene. "
        "The rewritten description must be fluent, coherent, and self-consistent -- "
        "do NOT leave gaps, placeholders, or awkward phrasing. "
        "Do not add any new content about the scene. "
        "Output only the rewritten description with no preamble or explanation."
    ).format(object=rc.judge_object_name)

    prompts = []
    iids_ordered = []
    for iid, cap in to_rewrite.items():
        messages = [
            {"role": "system", "content": rewrite_sys},
            {"role": "user", "content": f"Original description:\n{cap}"},
        ]
        prompt_text = tokenizer.apply_chat_template(
            messages, tokenize=False, add_generation_prompt=True, enable_thinking=False
        )
        prompts.append(prompt_text)
        iids_ordered.append(iid)

    sampling = SamplingParams(temperature=0.0, max_tokens=200, repetition_penalty=1.05)

    rewrite_results = {}
    for i in range(0, len(prompts), batch_size):
        batch_prompts = prompts[i : i + batch_size]
        batch_iids = iids_ordered[i : i + batch_size]
        outputs = llm.generate(batch_prompts, sampling, use_tqdm=True)
        for iid, out in zip(batch_iids, outputs):
            rewrite_results[iid] = out.outputs[0].text.strip()

    already_rewritten.update(rewrite_results)
    ckpt["step3_rewrite_done"] = already_rewritten
    ckpt["step3_done"] = True

    print(f"[Step 3] Rewrote {len(rewrite_results)} captions")

    return ckpt


def step4_strict_check(
    rc,
    ckpt: dict,
) -> dict:
    if ckpt.get("step4_done"):
        print("[Step 4] Strict quality check already done -- skipping.")
        return ckpt

    detector = KeywordMentionDetector(keywords=rc.mention_keywords, compound_mode=False)
    rewrite_done = ckpt.get("step3_rewrite_done", {})
    already_checked = ckpt.get("step4_strict_done", {})

    strict_results = {}
    for iid, rewritten in rewrite_done.items():
        if iid not in already_checked:
            strict_results[iid] = not detector.mentions_object(rewritten)

    already_checked.update(strict_results)
    ckpt["step4_strict_done"] = already_checked
    ckpt["step4_done"] = True

    passed = sum(1 for v in strict_results.values() if v)
    failed = sum(1 for v in strict_results.values() if not v)
    print(f"[Step 4] Strict check: {passed} passed, {failed} failed (out of {len(strict_results)} rewrites)")

    return ckpt


def assemble_output(
    data: list[dict],
    rc,
    ckpt: dict,
) -> list[dict]:
    object_key = rc.object_key
    regex_matched = ckpt.get("step1_regex_matched", {})
    judge_done = ckpt.get("step2_judge_done", {})
    rewrite_done = ckpt.get("step3_rewrite_done", {})
    strict_done = ckpt.get("step4_strict_done", {})

    output = []
    for entry in data:
        out = dict(entry)
        iid = out["image_id"]

        if out.get(object_key, 0) == 1:
            out["hallucinating"] = False
            out["edited_caption"] = None
            out["mention_confidence"] = None
            out["rewrite_passed_strict_check"] = None
        elif iid in regex_matched:
            verdict = judge_done.get(iid, "NO")
            if verdict == "YES":
                out["hallucinating"] = True
                out["edited_caption"] = rewrite_done.get(iid, entry.get("llava_caption", ""))
                out["mention_confidence"] = "regex+llm"
                out["rewrite_passed_strict_check"] = strict_done.get(iid, None)
            else:
                out["hallucinating"] = False
                out["edited_caption"] = None
                out["mention_confidence"] = "regex"
                out["rewrite_passed_strict_check"] = None
        else:
            out["hallucinating"] = False
            out["edited_caption"] = None
            out["mention_confidence"] = None
            out["rewrite_passed_strict_check"] = None

        output.append(out)

    return output


def main() -> None:
    ap = argparse.ArgumentParser(
        description=(
            "Identify hallucinated captions and rewrite them with Qwen3-8B via vLLM. "
            "Pipeline: regex filter -> LLM judge -> LLM rewrite -> strict quality check."
        )
    )
    ap.add_argument("--relation", default="bathroom_toilet",
                    help="Relation key (default: bathroom_toilet)")
    ap.add_argument("--gpus", default="0",
                    help="Comma-separated GPU IDs for CUDA_VISIBLE_DEVICES (default: 0)")
    ap.add_argument("--n_gpus", type=int, default=1,
                    help="Number of GPUs for vLLM tensor parallel (default: 1)")
    ap.add_argument("--model", default="Qwen/Qwen3-8B",
                    help="Qwen model path (default: Qwen/Qwen3-8B)")
    ap.add_argument("--input", default=None,
                    help="Override input path")
    ap.add_argument("--output", default=None,
                    help="Override output path (default: same as input)")
    ap.add_argument("--batch_size", type=int, default=64,
                    help="vLLM batch size per step (default: 64)")
    ap.add_argument("--dry_run", action="store_true",
                    help="Run regex filter only, skip all LLM inference")
    args = ap.parse_args()

    rc = get_relation_config(args.relation)

    if args.input:
        input_path = args.input
    else:
        input_path = os.path.join(
            os.path.dirname(os.path.abspath(__file__)),
            "..", "data", args.relation, "all_captions.json"
        )

    output_path = args.output if args.output else input_path

    print(f"[Config] relation={args.relation} object_key={rc.object_key}")
    print(f"[Config] input={input_path}")
    print(f"[Config] output={output_path}")

    data = load_data(input_path)
    print(f"[Config] Loaded {len(data)} entries")

    total_obj0 = sum(1 for e in data if e.get(rc.object_key, 0) == 0)
    total_obj1 = sum(1 for e in data if e.get(rc.object_key, 0) == 1)
    print(f"[Config] object=0: {total_obj0}, object=1: {total_obj1}")

    ckpt = load_checkpoint(output_path)

    print("\n[Step 1] Regex pre-filter (compound mode) ...")
    ckpt = step1_regex_filter(data, rc, ckpt)
    save_checkpoint(output_path, ckpt)

    if args.dry_run:
        regex_matched = ckpt.get("step1_regex_matched", {})
        print(f"\n[Dry run] Regex matched {len(regex_matched)}/{total_obj0} object=0 entries")
        print("[Dry run] Saving regex-only output ...")

        output = assemble_output(data, rc, ckpt)
        save_data(output_path, output)
        print(f"[Dry run] Saved {len(output)} entries to {output_path}")
        return

    os.environ["CUDA_VISIBLE_DEVICES"] = args.gpus
    print(f"\n[Config] CUDA_VISIBLE_DEVICES={args.gpus}, tensor_parallel_size={args.n_gpus}")

    llm, tokenizer = create_llm(args.model, args.n_gpus)

    print("\n[Step 2] LLM judge ...")
    t0 = time.time()
    ckpt = step2_llm_judge(data, rc, ckpt, llm, tokenizer, args.batch_size)
    save_checkpoint(output_path, ckpt)
    print(f"[Step 2] Completed in {time.time() - t0:.1f}s")

    print("\n[Step 3] LLM rewrite ...")
    t0 = time.time()
    ckpt = step3_llm_rewrite(rc, ckpt, llm, tokenizer, args.batch_size)
    save_checkpoint(output_path, ckpt)
    print(f"[Step 3] Completed in {time.time() - t0:.1f}s")

    del llm
    import gc
    gc.collect()
    import torch
    torch.cuda.empty_cache()

    print("\n[Step 4] Strict quality check ...")
    ckpt = step4_strict_check(rc, ckpt)
    save_checkpoint(output_path, ckpt)

    print("\n[Output] Assembling final output ...")
    output = assemble_output(data, rc, ckpt)
    save_data(output_path, output)

    hallucinating = sum(1 for e in output if e.get("hallucinating"))
    regex_only = sum(1 for e in output if e.get("mention_confidence") == "regex")
    regex_llm = sum(1 for e in output if e.get("mention_confidence") == "regex+llm")
    rewritten = sum(1 for e in output if e.get("edited_caption") is not None)
    strict_pass = sum(1 for e in output if e.get("rewrite_passed_strict_check") is True)
    strict_fail = sum(1 for e in output if e.get("rewrite_passed_strict_check") is False)

    print(f"\n[Summary]")
    print(f"  Total entries:       {len(output)}")
    print(f"  object=0 entries:    {total_obj0}")
    print(f"  Regex matches:       {regex_only + regex_llm}")
    print(f"    LLM judge NO:     {regex_only}")
    print(f"    LLM judge YES:    {regex_llm}")
    print(f"  Hallucinating:       {hallucinating}")
    print(f"  Rewritten:           {rewritten}")
    print(f"  Strict check pass:   {strict_pass}")
    print(f"  Strict check fail:   {strict_fail}")
    print(f"\nSaved {len(output)} entries to {output_path}")


if __name__ == "__main__":
    main()