"""
Pre-process dataset for Z-Image LoRA training.
Reads images from --input, resizes to [768, 1280] step 16,
writes as sequentially numbered JPEGs + paired .txt captions to --output.

Usage:
  python preprocess_dataset.py --input /path/to/img --output /path/to/img3
"""

import argparse
import os
import struct
from PIL import Image, ImageFile

# Allow truncated images
ImageFile.LOAD_TRUNCATED_IMAGES = True
from tqdm import tqdm


MIN_SIZE = 768
MAX_SIZE = 1280
STEP = 16
JPEG_QUALITY = 97
NAME_WIDTH = 7  # 0000001, 0000002, ...

IMAGE_EXTS = {'.jpg', '.jpeg', '.png', '.webp', '.bmp'}


def process_image(img: Image.Image) -> Image.Image:
    """Resize + crop to fit [MIN_SIZE, MAX_SIZE] with STEP alignment."""
    w, h = img.size

    # 1. Fit max side to MAX_SIZE (only downsample if too big)
    if max(w, h) > MAX_SIZE:
        if w >= h:
            new_w = MAX_SIZE
            new_h = int(MAX_SIZE * h / w)
        else:
            new_h = MAX_SIZE
            new_w = int(MAX_SIZE * w / h)
    else:
        new_w, new_h = w, h

    # 2. If min side is still too small, fit to MIN_SIZE
    if min(new_w, new_h) < MIN_SIZE:
        if new_w <= new_h:
            new_w = MIN_SIZE
            new_h = int(MIN_SIZE * new_h / new_w)
        else:
            new_h = MIN_SIZE
            new_w = int(MIN_SIZE * new_w / new_h)

    # 3. Crop to STEP-aligned, clamped to [MIN_SIZE, MAX_SIZE]
    crop_w = min(MAX_SIZE, (new_w // STEP) * STEP)
    crop_h = min(MAX_SIZE, (new_h // STEP) * STEP)
    crop_w = max(MIN_SIZE, crop_w)
    crop_h = max(MIN_SIZE, crop_h)

    # 4. Resize + center crop
    img = img.convert("RGB").resize((new_w, new_h), Image.LANCZOS)
    left = (new_w - crop_w) // 2
    top = (new_h - crop_h) // 2
    img = img.crop((left, top, left + crop_w, top + crop_h))

    return img


def main():
    parser = argparse.ArgumentParser(description="Resize dataset for Z-Image")
    parser.add_argument("--input", required=True)
    parser.add_argument("--output", required=True)
    parser.add_argument("--dry-run", action="store_true")
    args = parser.parse_args()

    os.makedirs(args.output, exist_ok=True)

    # Collect image files and find paired .txt captions
    all_files = sorted(os.listdir(args.input))
    image_files = [f for f in all_files if os.path.splitext(f)[1].lower() in IMAGE_EXTS]

    # Build mapping: basename -> caption text (if .txt exists with same basename)
    captions = {}
    for f in all_files:
        if f.lower().endswith(".txt"):
            stem = os.path.splitext(f)[0]
            with open(os.path.join(args.input, f)) as fh:
                captions[stem] = fh.read()

    print(f"Found {len(image_files)} images, {len(captions)} paired .txt captions")
    print(f"Min: {MIN_SIZE}, Max: {MAX_SIZE}, Step: {STEP}, Quality: {JPEG_QUALITY}")

    if args.dry_run:
        for fname in tqdm(image_files, desc="Dry run"):
            try:
                img = Image.open(os.path.join(args.input, fname))
                img.load()
                out = process_image(img)
                print(f"{fname}: {img.size} -> {out.size}")
            except (OSError, IOError, struct.error) as e:
                print(f"{fname}: SKIP ({e})")
        return

    stats = {"upscaled": 0, "downscaled": 0, "unchanged": 0, "skipped": 0}

    for idx, fname in enumerate(tqdm(image_files, desc="Processing"), start=1):
        src_path = os.path.join(args.input, fname)
        stem = os.path.splitext(fname)[0]

        try:
            img = Image.open(src_path)
            img.load()  # force load to catch truncated images early
            orig_w, orig_h = img.size
            out = process_image(img)
        except (OSError, IOError, struct.error) as e:
            tqdm.write(f"  Skip {fname}: {e}")
            stats["skipped"] += 1
            continue

        if out.size == (orig_w, orig_h):
            stats["unchanged"] += 1
        elif max(out.size) > max(orig_w, orig_h):
            stats["upscaled"] += 1
        else:
            stats["downscaled"] += 1

        # Sequential output names: 0000001.jpg, 0000001.txt
        name = f"{idx:0{NAME_WIDTH}d}"
        out_path = os.path.join(args.output, name + ".jpg")
        out.save(out_path, quality=JPEG_QUALITY)

        # Copy caption if paired .txt exists
        if stem in captions:
            txt_path = os.path.join(args.output, name + ".txt")
            with open(txt_path, "w") as f:
                f.write(captions[stem])

    print(f"\nDone: {stats['unchanged']} unchanged, "
          f"{stats['upscaled']} upscaled, "
          f"{stats['downscaled']} downscaled, "
          f"{stats['skipped']} skipped")
    print(f"Output: {args.output}")


if __name__ == "__main__":
    main()