"""
1_build_dataset.py  —  STAGE 1: build the captioned dataset CSV
================================================================
Scans a folder of images, auto-captions each with BLIP, and writes a CSV:

    image_name,caption
    img_001.jpg,gr4v_style a man seated at a desk, navy shirt, office background
    img_002.jpg,gr4v_style outdoor portrait, green park, overcast light
    ...

Then YOU open dataset.csv and edit the captions by hand before training.
This is the corrected version of the create_dataset() idea from your main.py
(which wrote image *paths* as "input_ids" — not usable). Here we write clean
image_name + caption columns, with the trigger word prepended.

Captioning principle: describe what you DON'T want fused to the trigger
(clothing, pose, background). Leave the subject's identity UNdescribed so it
binds to the trigger word.

Run:
    python 1_build_dataset.py --images ./train --out dataset.csv --trigger gr4v_style

Install:
    pip install -U transformers torch pillow pandas tqdm
"""

import argparse
import csv
from pathlib import Path

import torch
from PIL import Image
from tqdm import tqdm
from transformers import BlipProcessor, BlipForConditionalGeneration

VALID_EXT = {".jpg", ".jpeg", ".png"}


def build_dataset(image_dir: str, out_csv: str, trigger: str):
    device = "cuda" if torch.cuda.is_available() else "cpu"

    print("Loading BLIP captioner...")
    processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
    model = BlipForConditionalGeneration.from_pretrained(
        "Salesforce/blip-image-captioning-base"
    ).to(device)

    image_dir = Path(image_dir)
    rows = []

    files = sorted(p for p in image_dir.iterdir() if p.suffix.lower() in VALID_EXT)
    if not files:
        raise SystemExit(f"No .jpg/.jpeg/.png images found in {image_dir}")

    for p in tqdm(files, desc="Captioning"):
        try:
            image = Image.open(p).convert("RGB")
        except Exception as e:
            print(f"  skip {p.name}: {e}")
            continue

        inputs = processor(image, return_tensors="pt").to(device)
        out = model.generate(**inputs, max_new_tokens=40)
        caption = processor.decode(out[0], skip_special_tokens=True).strip()

        # Prepend the trigger word so it sits at the front of every caption.
        full_caption = f"{trigger} {caption}"
        rows.append({"image_name": p.name, "caption": full_caption})

    # Write CSV
    with open(out_csv, "w", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(f, fieldnames=["image_name", "caption"])
        writer.writeheader()
        writer.writerows(rows)

    print(f"\nWrote {len(rows)} rows -> {out_csv}")
    print("NEXT: open the CSV and edit captions before training.")
    print("  • Describe clothing, pose, background, lighting (things to stay variable)")
    print("  • Do NOT describe the subject's face/identity (let it bind to the trigger)")


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("--images", default="./train", help="folder of training images")
    ap.add_argument("--out", default="dataset.csv", help="output CSV path")
    ap.add_argument("--trigger", default="gr4v_style", help="unique trigger word")
    args = ap.parse_args()
    build_dataset(args.images, args.out, args.trigger)