"""Caption CC3M general images for data mixing.

Uses LLaVA to generate captions for images that will be used as
"general/unrelated" samples in adversarial training.

Usage:
    python -m experiment.data.caption_cc3m_general \\
        --cc3m_dir /path/to/cc3m_images/train \\
        --output /path/to/output/cc3m_captions.json \\
        --batch_size 8
"""
from __future__ import annotations

import argparse
import json
import os
import gc
from pathlib import Path

import torch
from tqdm import tqdm
from PIL import Image


def _clear_gpu():
    gc.collect()
    torch.cuda.empty_cache()
    torch.cuda.synchronize()


def load_model(model_name: str, device: str):
    """Load LLaVA model for captioning."""
    from transformers import LlavaForConditionalGeneration, AutoProcessor

    processor = AutoProcessor.from_pretrained(model_name)
    model = LlavaForConditionalGeneration.from_pretrained(
        model_name,
        torch_dtype=torch.bfloat16,
        device_map={"": device},
        attn_implementation="sdpa",
    ).eval()

    return model, processor


def caption_images(
    image_paths: list[str],
    model,
    processor,
    device: str,
    batch_size: int,
    prompt: str = "Describe this image.",
):
    """Caption a list of images."""
    results = []

    for i in tqdm(range(0, len(image_paths), batch_size), desc="Captioning"):
        batch_paths = image_paths[i:i + batch_size]
        batch_images = []

        for path in batch_paths:
            try:
                img = Image.open(path).convert("RGB")
                batch_images.append(img)
            except Exception as e:
                print(f"Warning: Could not load {path}: {e}")

        if not batch_images:
            continue

        # Prepare inputs
        texts = [f"USER: <image>\n{prompt} ASSISTANT:" for _ in batch_images]
        inputs = processor(text=texts, images=batch_images, return_tensors="pt", padding=True)
        inputs = {k: v.to(device) if hasattr(v, "to") else v for k, v in inputs.items()}

        # Generate
        with torch.inference_mode():
            outputs = model.generate(
                **inputs,
                max_new_tokens=100,
                do_sample=False,
                use_cache=True,
            )

        # Decode
        input_len = inputs["input_ids"].shape[1]
        for path, seq in zip(batch_paths, outputs):
            caption = processor.decode(seq[input_len:], skip_special_tokens=True).strip()
            results.append({
                "image_path": path,
                "image_id": Path(path).stem,
                "caption": caption,
            })

        _clear_gpu()

    return results


def main():
    parser = argparse.ArgumentParser(description="Caption CC3M general images")
    parser.add_argument("--cc3m_dir", required=True, help="Path to CC3M images directory")
    parser.add_argument("--output", required=True, help="Output JSON file path")
    parser.add_argument("--model", default="llava-hf/llava-1.5-7b-hf", help="LLaVA model")
    parser.add_argument("--batch_size", type=int, default=8, help="Batch size")
    parser.add_argument("--prompt", default="Describe this image.", help="Captioning prompt")
    parser.add_argument("--max_samples", type=int, default=None, help="Max samples to caption")
    args = parser.parse_args()

    device = "cuda" if torch.cuda.is_available() else "cpu"

    print(f"Loading CC3M images from {args.cc3m_dir}...")
    image_paths = list(Path(args.cc3m_dir).glob("*.jpg"))

    if args.max_samples:
        import random
        random.seed(42)
        random.shuffle(image_paths)
        image_paths = image_paths[:args.max_samples]

    print(f"Found {len(image_paths)} images")

    print(f"Loading LLaVA model: {args.model}")
    model, processor = load_model(args.model, device)

    print("Generating captions...")
    results = caption_images(
        [str(p) for p in image_paths],
        model,
        processor,
        device,
        args.batch_size,
        args.prompt,
    )

    # Save results
    os.makedirs(os.path.dirname(args.output) or ".", exist_ok=True)
    with open(args.output, "w") as f:
        json.dump(results, f, indent=2)

    print(f"Saved {len(results)} captions to {args.output}")

    # Cleanup
    del model, processor
    _clear_gpu()


if __name__ == "__main__":
    main()