#!/usr/bin/env python3
"""Generate multiple images from a kq4-decompile visual image with different seeds."""

import random
import subprocess
import sys
import os
import urllib.request


def check_server(server_address: str = "127.0.0.1:8188", timeout: int = 5) -> bool:
    """Check if ComfyUI server is running and accessible."""
    try:
        req = urllib.request.Request(
            f"http://{server_address}/system_stats",
            method="GET",
        )
        with urllib.request.urlopen(req, timeout=timeout) as response:
            return response.status == 200
    except Exception:
        return False


def count_existing_variations(image_path: str) -> int:
    """Count existing image variations for a visual image.

    Args:
        image_path: Path to the kq4-decompile visual image

    Returns:
        Number of existing generated images matching the room name
    """
    image_dir = os.path.dirname(os.path.abspath(image_path)) or "."
    klein_outputs_dir = os.path.join(image_dir, "klein_outputs")

    if not os.path.exists(klein_outputs_dir):
        return 0

    image_basename = os.path.basename(image_path)
    room_name = (
        image_basename.replace("pic_", "").replace("_visual.png", "").replace("_", "-")
    )

    count = 0
    for item in os.listdir(klein_outputs_dir):
        if item.startswith(f"{room_name}_") and item.endswith("_generated.png"):
            count += 1

    return count


def generate_multiple_klein(
    image_path: str,
    count: int = 2,
    server: str = "127.0.0.1:8188",
    dry_run: bool = False,
    lora_disney: float = 0.8,
    lora_kq5hoyos: float = 0.7,
    lora_sylvain: float = 0.5,
    caption_prompt: str = None,
) -> tuple[list[str], int]:
    """Generate multiple images with different random seeds.

    Args:
        image_path: Path to kq4-decompile visual image
        count: Total number of variations to have (default: 2)
        server: ComfyUI server address
        dry_run: If True, validate only without generating

    Returns:
        Tuple of (list of created image paths, number of existing variations)
    """
    if not os.path.exists(image_path):
        print(f"Error: Image file not found: {image_path}")
        sys.exit(1)

    image_basename = os.path.basename(image_path)
    room_name = (
        image_basename.replace("pic_", "").replace("_visual.png", "").replace("_", "-")
    )

    if dry_run:
        print(f"  Image: {image_basename}")
        return [], count_existing_variations(image_path)

    existing = count_existing_variations(image_path)
    print(f"Found {existing} existing variations")
    print(f"Requested total: {count}")

    needed = count - existing

    if needed <= 0:
        print(f"\nAlready have {existing} variations (requested {count})")
        print("No new images needed.")
        return [], existing

    print(f"Will generate {needed} new variation(s)")

    output_paths = []
    script_dir = os.path.dirname(os.path.abspath(__file__))
    generator_script = os.path.join(script_dir, "generate_from_caption_klein.py")

    for i in range(needed):
        seed = random.randint(1, 2**32 - 1)
        print(f"\n{'=' * 60}")
        print(f"Generation {i + 1} of {needed} (seed: {seed})")
        print(f"Progress: {existing + i + 1} of {count} total")
        print(f"{'=' * 60}")

        cmd = [
            sys.executable,
            generator_script,
            image_path,
            str(seed),
            "--server",
            server,
            "--lora-disney",
            str(lora_disney),
            "--lora-kq5hoyos",
            str(lora_kq5hoyos),
            "--lora-sylvain",
            str(lora_sylvain),
        ]
        if caption_prompt:
            cmd.extend(["--caption-prompt", caption_prompt])

        result = subprocess.run(cmd, capture_output=False)

        if result.returncode != 0:
            print(
                f"Warning: Generation {i + 1} of {needed} failed with exit code {result.returncode}"
            )
        else:
            image_dir = os.path.dirname(os.path.abspath(image_path)) or "."
            klein_outputs_dir = os.path.join(image_dir, "klein_outputs")
            output_path = os.path.join(
                klein_outputs_dir, f"{room_name}_{seed}_generated.png"
            )
            output_paths.append(output_path)

    return output_paths, existing


def main():
    import argparse

    parser = argparse.ArgumentParser(
        description="Generate multiple images from kq4-decompile visual with different seeds"
    )
    parser.add_argument(
        "image_path",
        help="Path to kq4-decompile visual image (e.g., kq4-sierra-decompile/rooms/kq4-057-witch-cave/pic_057_visual.png)",
    )
    parser.add_argument(
        "--count",
        "-n",
        type=int,
        default=2,
        help="Number of variations to generate (default: 2)",
    )
    parser.add_argument(
        "--server",
        default="127.0.0.1:8188",
        help="ComfyUI server address (default: 127.0.0.1:8188)",
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Test mode: validate image file and server connection without generating",
    )
    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=None,
        help="Custom caption prompt for scene description",
    )

    args = parser.parse_args()

    if args.dry_run:
        print("\n" + "=" * 60)
        print("DRY RUN MODE - Validating without generating")
        print("=" * 60)

        print(f"\nChecking ComfyUI server at {args.server}...")
        if check_server(args.server):
            print("✓ Server is accessible")
        else:
            print(f"✗ Server is NOT accessible at {args.server}")
            print("  Please ensure ComfyUI is running before generating images.")
            sys.exit(1)

        print(f"\nValidating: {args.image_path}")
        output_paths, existing = generate_multiple_klein(
            args.image_path,
            args.count,
            args.server,
            dry_run=True,
            lora_disney=args.lora_disney,
            lora_kq5hoyos=args.lora_kq5hoyos,
            lora_sylvain=args.lora_sylvain,
            caption_prompt=args.caption_prompt,
        )
        print(f"✓ Image file is valid (found {existing} existing variations)")

        print("\n" + "=" * 60)
        print("✓ Dry run successful! All checks passed.")
        print("=" * 60)
        sys.exit(0)

    print(f"\nChecking ComfyUI server at {args.server}...")
    if not check_server(args.server):
        print(f"Error: ComfyUI server is not running at {args.server}")
        print("Please start ComfyUI first or check the server address.")
        print(f"\nTo test without generating, use: --dry-run")
        sys.exit(1)
    print("✓ Server is running\n")

    output_paths, existing = generate_multiple_klein(
        args.image_path,
        args.count,
        args.server,
        lora_disney=args.lora_disney,
        lora_kq5hoyos=args.lora_kq5hoyos,
        lora_sylvain=args.lora_sylvain,
        caption_prompt=args.caption_prompt,
    )

    print(f"\n{'=' * 60}")
    print("All generations complete!")
    print(f"{'=' * 60}")
    if output_paths:
        print(f"\nCreated {len(output_paths)} new image(s):")
        for output_path in output_paths:
            print(f"  - {output_path}")
    print(f"\nTotal variations now: {existing + len(output_paths)}")


if __name__ == "__main__":
    main()