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.benchmark.score_images import get_score_images_argparser

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


def collect_benchmark_img_dirs(
    top_level_dir: Path, exactly_n_images: int | None
) -> list[Path]:
    """
    Recursively collects benchmark image directories under top_level_dir that contain PNG files.
    If exactly_n_images is provided, only directories with that exact number of PNGs are returned.
    """
    img_dirs = []
    for img_dir in top_level_dir.glob("**/"):
        if img_dir.is_dir():
            num_imgs = len(list(img_dir.glob("*.png")))
            if num_imgs > 0:
                if exactly_n_images is None or num_imgs == exactly_n_images:
                    img_dirs.append(img_dir)
                else:
                    print(
                        f"WARNING: Directory {img_dir} has {num_imgs} images; expected {exactly_n_images}. Skipping."
                    )
    return img_dirs


def score_command(
    python_path: Path,
    benchmark_prompts: Path,
    image_dir: Path,
    output_subpath: Path,
    exactly_n_images: int | None,
    delete_after: bool,
    rescore_existing: bool,
    image_naming_mode: str,
    file_mode: str | None = None,
) -> str:
    """
    Builds the command string to run score_images.py on a given image directory.
    """
    exactly_str = (
        f"--exactly-n-images {exactly_n_images}" if exactly_n_images else ""
    )
    delete_str = "--delete-after" if delete_after else ""
    rescore_str = "--rescore-existing" if rescore_existing else ""
    file_mode_str = f"--file-mode {file_mode}" if file_mode else ""
    return (
        f"{python_path} {PROJECT_ROOT}/editing_DiTs/benchmark/score_images.py"
        f" --benchmark-prompts {benchmark_prompts}"
        f" --image-dir {image_dir}"
        f" --output-subpath {output_subpath}"
        f" --image-naming-mode {image_naming_mode}"
        f" {file_mode_str} {exactly_str} {delete_str} {rescore_str}"
    )


def generate_score_commands(
    python_path: Path,
    benchmark_prompts: Path,
    image_dir: Path,
    output_subpath: Path,
    exactly_n_images: int | None,
    delete_after: bool,
    rescore_existing: bool,
    image_naming_mode: str,
    file_mode: str,
) -> list[str]:
    """
    Generates a flat list of command strings to score benchmark image directories.
    """
    img_dirs = collect_benchmark_img_dirs(image_dir, exactly_n_images)
    print(f"Found {len(img_dirs)} benchmark image directories to score.")
    commands = []
    for dir in img_dirs:
        cmd = score_command(
            python_path,
            benchmark_prompts,
            dir,
            output_subpath,
            exactly_n_images,
            delete_after,
            rescore_existing,
            image_naming_mode,
            file_mode,
        )
        commands.append(cmd)
    return commands


def main() -> None:
    # Merge the common multi-job arguments with score_images-specific arguments.
    common_parser = get_common_argparser()
    score_parser = get_score_images_argparser()
    parser = argparse.ArgumentParser(parents=[common_parser, score_parser])

    parser.add_argument(
        "--image-dir-list",
        required=False,
        nargs="+",
        help="A space seperated list of image dir paths to score.",
    )

    args = parser.parse_args()

    # Validate and resolve common arguments.
    validate_common_args(args)

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

    args.benchmark_prompts = args.benchmark_prompts.resolve()
    if not args.benchmark_prompts.exists():
        print(
            f"WARNING: Benchmark prompts file {args.benchmark_prompts} not found."
        )

    if args.image_dir is None and args.image_dir_list is None:
        raise ValueError("No image directory or list of directories provided.")

    if args.image_dir is not None and args.image_dir_list is not None:
        raise ValueError(
            "Provide either an image directory or a list of directories, not both."
        )

    image_dir_list = []
    if args.image_dir_list is not None:
        image_dir_list = [Path(d).resolve() for d in args.image_dir_list]
    if args.image_dir is not None:
        image_dir_list = [args.image_dir]

    all_commands = []
    for img_dir in image_dir_list:
        if not img_dir.exists():
            print(f"WARNING: Image directory {img_dir} not found.")
            continue
        print(img_dir)
        commands = generate_score_commands(
            args.python_path,
            args.benchmark_prompts,
            img_dir,
            args.output_subpath,
            args.exactly_n_images,
            args.delete_after,
            args.rescore_existing,
            args.image_naming_mode,
            args.file_mode,
        )
        all_commands.extend(commands)

    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(all_commands)


if __name__ == "__main__":
    main()
