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.compute_macs import get_compute_macs_argparser


def collect_schedule_files(input_dir: Path) -> list[Path]:
    """
    Recursively collects all JSON schedule files from the given input directory.
    """
    return list(input_dir.glob("**/*.json"))


def generate_compute_macs_commands(
    image_generator: str,
    python_path: Path,
    input_dir: Path,
    recompute_existing: bool,
) -> list[str]:
    """
    Generates a flat list of command strings for computing MACs/FLOPs on each schedule file.

    Each command calls the standalone compute_macs.py with --input-dir set to a single JSON file.
    If 'recompute_existing' is True, the flag is added.
    """
    schedule_files = collect_schedule_files(input_dir)
    print(f"Found {len(schedule_files)} schedule files for MACs computation.")
    commands = []
    for schedule in schedule_files:
        cmd = (
            f"{python_path} {PROJECT_ROOT}/editing_DiTs/benchmark/compute_macs.py"
            f" --image-generator {image_generator}"
            f" --input-dir {schedule}"
        )
        if recompute_existing:
            cmd += " --recompute-existing"
        commands.append(cmd)
    return commands


def main() -> None:
    # Merge the common multi-job arguments with compute_macs-specific ones.
    common_parser = get_common_argparser()
    compute_macs_parser = get_compute_macs_argparser()
    parser = argparse.ArgumentParser(
        parents=[common_parser, compute_macs_parser],
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    args = parser.parse_args()

    # Validate and resolve common arguments.
    validate_common_args(args)
    args.input_dir = args.input_dir.resolve()
    if not args.input_dir.exists():
        print(f"ERROR: Input directory/file {args.input_dir} not found.")
        exit(1)

    commands = generate_compute_macs_commands(
        args.image_generator,
        args.python_path,
        args.input_dir,
        args.recompute_existing,
    )

    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)


if __name__ == "__main__":
    main()
