import argparse
from pathlib import Path

from editing_DiTs.benchmark.multi_job import (
    MultiJobRunner,
    get_common_argparser,
    validate_common_args,
    PROJECT_ROOT,
)
from editing_DiTs.image_generators.load_image_generator import (
    ImageGeneratorRegistry,
)

# Constants specific to image generation.
DEFAULT_BENCHMARK_DIR = PROJECT_ROOT / "results" / "benchmark"
DEFAULT_EMBEDDING_DIR = DEFAULT_BENCHMARK_DIR / "embeddings"
DEFAULT_IMAGE_DIR = DEFAULT_BENCHMARK_DIR / "full-benchmark" / "images"


def collect_schedules(schedule_dir: Path) -> list[Path]:
    """Recursively collect all JSON schedule files."""
    return list(schedule_dir.glob("**/*.json"))


def generate_benchmark_commands(
    image_generator: str,
    schedule_jsons: list[Path],
    name_rel_to: Path,
    python_path: Path,
    seed: int,
    seed_step: int,
    embedding_dir: Path,
    image_dir: Path,
    num_images_per_prompt: int,
    batch_size: int,
    regen_not_n_imgs: int | None,
) -> list[str]:
    """
    Generates a flat list of command strings for each schedule file.
    """
    commands = []
    for schedule_json in schedule_jsons:
        subdir = schedule_json.parent.resolve().relative_to(name_rel_to)
        regen_str = (
            f"--regen-if-not-n-images {regen_not_n_imgs}"
            if regen_not_n_imgs
            else ""
        )
        command = (
            f"{python_path} {PROJECT_ROOT}/editing_DiTs/benchmark/generate_images.py"
            f" --image-generator {image_generator}"
            f" --seed {seed}"
            f" --seed-step {seed_step}"
            f" --input-dir {embedding_dir}"
            f" --output-dir {image_dir / subdir}"
            f" --schedule-dir {schedule_json}"
            f" --num-images-per-prompt {num_images_per_prompt}"
            f" --batch-size {batch_size}"
            f" {regen_str}"
        )
        commands.append(command)
    return commands


def main() -> None:
    # Create the common argparser and add image-generation-specific arguments.
    common_parser = get_common_argparser()
    parser = argparse.ArgumentParser(parents=[common_parser])
    parser.add_argument(
        "--image-generator",
        type=str,
        required=True,
        help="The name of the image generator to use.",
        choices=list(ImageGeneratorRegistry.registry.keys()),
    )
    parser.add_argument(
        "--embedding-dir",
        type=Path,
        default=DEFAULT_EMBEDDING_DIR,
        help="Directory containing the embeddings for the benchmark.",
    )
    parser.add_argument(
        "--image-dir",
        type=Path,
        default=DEFAULT_IMAGE_DIR,
        help="Directory where benchmark images will be saved.",
    )
    parser.add_argument(
        "--schedule-dir",
        type=Path,
        required=False,
        help="Directory containing schedules to run (recursively searched).",
    )
    parser.add_argument(
        "--schedule-list",
        required=False,
        nargs="+",
        help="A comma seperated list of paths to <schedule-name>.json files.",
    )
    parser.add_argument(
        "--name-schedule-dir-rel-to",
        type=Path,
        help="Anchor path for subdirectories found under --schedule-dir.",
    )
    parser.add_argument(
        "--num-images-per-prompt",
        type=int,
        default=10,
        help="Number of images to generate per prompt.",
    )
    parser.add_argument(
        "--batch-size",
        type=int,
        default=100,
        help="Batch size while generating images.",
    )
    parser.add_argument(
        "--seed",
        type=int,
        default=0,
        help="Random seed to start image generation.",
    )
    parser.add_argument(
        "--seed-step",
        type=int,
        default=1,
        help="Amount to increment the seed for each subsequent generation.",
    )
    parser.add_argument(
        "--regen-if-not-n-images",
        type=int,
        help="Force re-generation if output dir does not contain exactly n images.",
    )
    args = parser.parse_args()

    # Validate and resolve the common arguments.
    validate_common_args(args)
    if args.name_schedule_dir_rel_to is None:
        if args.schedule_dir is not None:
            args.name_schedule_dir_rel_to = args.schedule_dir.resolve()
            if not args.schedule_dir.exists():
                print(
                    f"WARNING: Schedule directory {args.schedule_dir} not found."
                )
        else:
            raise ValueError("Must provide --name-schedule-dir-rel-to.")

    args.embedding_dir = args.embedding_dir.resolve()
    args.image_dir = args.image_dir.resolve()
    args.name_schedule_dir_rel_to = args.name_schedule_dir_rel_to.resolve()

    if args.schedule_dir is not None:
        args.schedule_dir = args.schedule_dir.resolve()

    if not args.embedding_dir.exists():
        print(f"WARNING: Embedding directory {args.embedding_dir} not found.")

    # Generate a flat list of commands.
    if args.schedule_dir is not None:
        schedule_jsons = collect_schedules(args.schedule_dir)
    elif args.schedule_list is not None:
        schedule_jsons = [Path(s) for s in args.schedule_list]
    else:
        raise ValueError(
            "Must provide either --schedule-dir or --schedule-list."
        )

    print(f"Found {len(schedule_jsons)} schedules.")
    commands = generate_benchmark_commands(
        args.image_generator,
        schedule_jsons,
        args.name_schedule_dir_rel_to,
        args.python_path,
        args.seed,
        args.seed_step,
        args.embedding_dir,
        args.image_dir,
        args.num_images_per_prompt,
        args.batch_size,
        args.regen_if_not_n_images,
    )

    runner = MultiJobRunner(
        scripts_dir=args.scripts_dir,
        time_per_job=args.time_per_job,
        qos=args.qos,
        gpu=args.gpu,
        cpu=args.cpu,
        mem=args.mem,
        nodelist=args.nodes,
        dry_run=args.dry_run,
    )
    runner.submit_jobs(commands)  # automatically batched


if __name__ == "__main__":
    main()
