#!/usr/bin/env python3
"""Generate combo images from all kq4-decompile visual images that have captions."""

import subprocess
import sys
import os
import re
import argparse
from pathlib import Path

DEFAULT_CAPTION_TEXT = "oil painting style"


def find_captions_for_image(image_path: Path) -> list[tuple[Path, str, int]]:
    """Find all caption files for a visual image.

    Looks for caption files in asset-work/kq4_XXX_room_name/ directory.
    Returns list of (caption_path, caption_text, caption_number) tuples.
    Uses a set to avoid duplicates when multiple room folders match.

    Args:
        image_path: Path to the visual image

    Returns:
        List of tuples with (caption_file_path, caption_text, caption_number)
    """
    image_basename = os.path.basename(image_path)
    room_number = image_basename.replace("pic_", "").split("_")[0]

    room_folder_pattern = f"kq4_{room_number}_*"
    script_dir = Path(__file__).parent.resolve()

    captions_set = set()
    captions = []
    for room_dir in script_dir.glob(room_folder_pattern):
        if room_dir.is_dir():
            for caption_file in sorted(room_dir.glob("caption_*.txt")):
                match = re.match(r"caption_(\d+)\.txt", caption_file.name)
                if match:
                    caption_num = int(match.group(1))
                    caption_key = (caption_num, caption_file.name)
                    if caption_key not in captions_set:
                        captions_set.add(caption_key)
                        with open(caption_file, "r") as f:
                            caption_text = f.read().strip()
                        captions.append((caption_file, caption_text, caption_num))

    return sorted(captions, key=lambda x: x[2]) if captions else [(None, DEFAULT_CAPTION_TEXT, 0)]


def find_caption_for_image(image_path: Path) -> tuple[Path | None, str]:
    """Find first caption file for a visual image.

    Args:
        image_path: Path to the visual image

    Returns:
        Tuple of (caption_file_path or None, caption_text or default)
    """
    captions = find_captions_for_image(image_path)
    if captions and captions[0][0]:
        return captions[0][0], captions[0][1]
    return None, DEFAULT_CAPTION_TEXT


def main():
    parser = argparse.ArgumentParser(
        description="Generate combo images from all kq4-decompile visual images with captions"
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Dry run mode - validate without generating",
    )
    args = parser.parse_args()

    script_dir = Path(__file__).parent.resolve()
    generator = script_dir / "generate_multiple_combo.py"

    dry_run = args.dry_run

    if not generator.exists():
        print(f"Error: generate_multiple_combo.py not found in {script_dir}")
        sys.exit(1)

    print("Finding all kq4-decompile visual images...")
    visual_images = list(Path("kq4-sierra-decompile/rooms").rglob("pic_*_visual.png"))

    total = len(visual_images)

    if total == 0:
        print("No kq4-decompile visual images found!")
        sys.exit(0)

    print(f"Found {total} visual image(s)")

    all_images_to_process = []
    for image_path in visual_images:
        captions = find_captions_for_image(image_path)
        for caption_path, caption_text, caption_num in captions:
            all_images_to_process.append((image_path, caption_text, caption_path, caption_num))

    total_captions = len(all_images_to_process)
    print(f"Found {total_captions} captions to process ({len(visual_images)} rooms)")

    if dry_run:
        print("\nDRY RUN MODE - Validating all visual images")

    counter = 0
    failed = 0
    total_existing = 0
    total_created = 0
    total_needed = 0

    variations_per_image = 1

    for image_path, caption_text, caption_path, caption_num in all_images_to_process:
        counter += 1

        caption_info = f"[caption_{caption_num}]" if caption_path else "[default]"

        cmd = [
            "python3",
            str(generator),
            str(image_path),
            caption_text,
            "--count",
            str(variations_per_image),
            "--caption-num",
            str(caption_num),
        ]
        if dry_run:
            cmd.append("--dry-run")

        try:
            result = subprocess.run(cmd, capture_output=True, text=True, check=True)
            print(result.stdout)

            existing_match = re.search(
                r"[Ff]ound (\d+) existing variations?", result.stdout
            )
            created_match = re.search(r"Created (\d+) new image", result.stdout)

            if existing_match:
                existing = int(existing_match.group(1))
                total_existing += existing
                needed = max(0, variations_per_image - existing)
                total_needed += needed
            if created_match:
                total_created += int(created_match.group(1))

            status = "✓" if result.returncode == 0 else "✗"
            print(f"[{counter}/{total_captions}] {status} {image_path.name} {caption_info}")
        except subprocess.CalledProcessError as e:
            failed += 1
            print(f"[{counter}/{total_captions}] ✗ {image_path.name} {caption_info} - FAILED")

    print()
    print("=" * 50)

    if dry_run:
        print("DRY RUN COMPLETE")
        print("=" * 50)
        print(f"Rooms found: {total}")
        print(f"Captions to process: {total_captions}")
        print(f"Images already existing: {total_existing}")
        print(f"Images would be created: {total_needed}")
        print(f"Total when complete:     {total_existing + total_needed}")
        if failed == 0:
            print(f"✓ All captions validated successfully!")
        else:
            print(f"✗ {failed} of {total_captions} captions failed validation")
            sys.exit(1)
    else:
        print("GENERATION COMPLETE")
        print("=" * 50)
        print(f"Rooms processed: {total}")
        print(f"Captions processed: {total_captions}")
        print(f"Images already existing: {total_existing}")
        print(f"Images newly created:     {total_created}")
        print(f"Total images now:         {total_existing + total_created}")
        if failed > 0:
            print(f"Failed:                   {failed}")


if __name__ == "__main__":
    main()