#!/usr/bin/env python3
"""Generate multiple images from all kq4-decompile visual images recursively."""

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

DEFAULT_CAPTION_PROMPT = 'describe the mood and atmosphere of the scene without describing colors or style. Focus on the framing, composition, lighting mood, and elements visible. Respond with just the result, preceded by "convert to a painting in the Disney Renaissance Style and kq5hoyos style and sylvain style. ". Keep it to around 20-25 words. If the image is framed in a particular way, make sure to call that out. never describe colors.'


def main():
    parser = argparse.ArgumentParser(
        description="Generate multiple images from all kq4-decompile visual images recursively"
    )
    parser.add_argument(
        "--lora-disney",
        type=float,
        default=0.8,
        help="Disney LoRA strength (default: 0.8)",
    )
    parser.add_argument(
        "--lora-kq5hoyos",
        type=float,
        default=0.7,
        help="KQ5Hoyos LoRA strength (default: 0.7)",
    )
    parser.add_argument(
        "--lora-sylvain",
        type=float,
        default=0.5,
        help="Sylvain LoRA strength (default: 0.5)",
    )
    parser.add_argument(
        "--caption-prompt",
        default=DEFAULT_CAPTION_PROMPT,
        help="Custom caption prompt for scene description",
    )
    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_klein.py"

    dry_run = args.dry_run

    if not generator.exists():
        print(f"Error: generate_multiple_klein.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)")

    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 in visual_images:
        counter += 1

        cmd = [
            "python3",
            str(generator),
            str(image_path),
            "--count",
            str(variations_per_image),
            "--lora-disney",
            str(args.lora_disney),
            "--lora-kq5hoyos",
            str(args.lora_kq5hoyos),
            "--lora-sylvain",
            str(args.lora_sylvain),
        ]
        if args.caption_prompt != DEFAULT_CAPTION_PROMPT:
            cmd.extend(["--caption-prompt", args.caption_prompt])
        if dry_run:
            cmd.append("--dry-run")

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

            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}] {status} {image_path}")
        except subprocess.CalledProcessError as e:
            failed += 1
            print(f"[{counter}/{total}] ✗ {image_path} - FAILED")

    print()
    print("=" * 50)

    if dry_run:
        print("DRY RUN COMPLETE")
        print("=" * 50)
        print(f"Visual images validated: {total}")
        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 visual images validated successfully!")
        else:
            print(f"✗ {failed} of {total} visual image(s) failed validation")
            sys.exit(1)
    else:
        print("GENERATION COMPLETE")
        print("=" * 50)
        print(f"Visual images processed: {total}")
        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()
