"""
multi_job.py

This module provides common functionality for generating and submitting slurm job scripts
for multi-job arrays. It defines shared constants, a MultiJobRunner class, and common
argument parsing/validation functions so that job-specific scripts can simply extend the
common argparser and then call the validation routine before continuing.

Usage example in a job-specific script:
    from editing_DiTs.benchmark.multi_job import MultiJobRunner, get_common_argparser, validate_common_args

    parser = get_common_argparser()
    parser.add_argument("--my-param", type=int, required=True, help="Job-specific parameter")
    args = parser.parse_args()
    validate_common_args(args)

    # Build a flat list of commands.
    commands = my_command_generator(...)
    
    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,  # set via validation routine if --gpus is provided
        dry_run=args.dry_run,
    )
    runner.submit_jobs(commands)
"""

import argparse
import subprocess
import json
from pathlib import Path
from typing import List

# Shared constants.
DEFAULT_PYTHON_PATH = Path(
    "~/micromamba/envs/ecad/bin/python"
)
PROJECT_ROOT = Path("~/dev/editing-DiTs")
DEFAULT_ENV_SETUP = f"""
. /usr/share/Modules/init/bash
module add cuda/12.4.1

if [ -f "$(pwd)/.env" ]; then
    export $(grep -v '^#' "$(pwd)/.env" | xargs)
elif [[ -f "{PROJECT_ROOT}/.env" ]]; then
    export $(grep -v '^#' "{PROJECT_ROOT}/.env" | xargs)
fi

export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
"""
MAX_ARRAY_SIZE = 500

RERUN_SCRIPT = """
import subprocess, time
interval = 90

def main(failed_jobs):
    while failed_jobs:
        to_del = []
        for i, (cmd, script) in enumerate(failed_jobs):
            run_res = subprocess.run([cmd, script], text=True)
            if run_res.returncode == 0:
                print(f"Successfully reran {script}")
                to_del.append(i)
        for i in reversed(to_del):
            del failed_jobs[i]
        print(f"\\n{len(failed_jobs)} jobs left to rerun.")
        time.sleep(interval)
    print("\\n\\nDone.")
"""


class MultiJobRunner:
    """
    Encapsulates functionality for creating and submitting slurm job scripts for multi-job arrays.
    """

    def __init__(
        self,
        scripts_dir: Path,
        time_per_job: str,
        qos: str,
        gpu: str,
        cpu: int,
        mem: int,
        nodelist: List[str] | None = None,
        env_setup: str = DEFAULT_ENV_SETUP,
        submit_cmd: str = "sbatch",
        dry_run: bool = False,
    ):
        self.scripts_dir = scripts_dir
        self.time_per_job = time_per_job
        self.qos = qos
        self.gpu = gpu
        self.cpu = cpu
        self.mem = mem
        self.nodelist = nodelist
        self.env_setup = env_setup
        self.submit_cmd = submit_cmd
        self.dry_run = dry_run

    def script_header(
        self, num_jobs: int, slurm_output_file: Path, slurm_error_file: Path
    ) -> str:
        if "REDACTED" in self.qos:
            partition = self.qos
            if self.qos == "REDACTED":
                account = "REDACTED"
            else:
                account = self.qos
        elif "REDACTED" in self.qos:
            partition = "REDACTED"
            account = "REDACTED"
        else:
            partition = "REDACTED"
            account = "REDACTED"

        gpu_nodelist = (
            f"--nodelist={self.nodelist}"
            if self.nodelist is not None
            else f"--gres=gpu:{self.gpu}"
        )
        return (
            "#!/bin/bash\n"
            f"#SBATCH --array=1-{num_jobs}\n"
            "#SBATCH --requeue\n"
            "#SBATCH --job-name=multi_job\n"
            f"#SBATCH --output={slurm_output_file}\n"
            f"#SBATCH --error={slurm_error_file}\n"
            f"#SBATCH --time={self.time_per_job}\n"
            f"#SBATCH --qos={self.qos}\n"
            f"#SBATCH --partition={partition}\n"
            f"#SBATCH --account={account}\n"
            f"#SBATCH --cpus-per-task={self.cpu}\n"
            f"#SBATCH --mem={self.mem}G\n"
            f"#SBATCH {gpu_nodelist}\n"
        )

    @staticmethod
    def script_srun(
        job_out_path_file: Path,
        job_err_path_file: Path,
        job_command_file: Path,
    ) -> str:
        return (
            f"srun --output=$(head -n $SLURM_ARRAY_TASK_ID {job_out_path_file} | tail -n 1)"
            f" --error=$(head -n $SLURM_ARRAY_TASK_ID {job_err_path_file} | tail -n 1)"
            f" $(head -n $SLURM_ARRAY_TASK_ID {job_command_file} | tail -n 1)"
        )

    @staticmethod
    def _batch_commands(commands: List[str]) -> List[List[str]]:
        """Automatically batch a flat list of commands into chunks of size MAX_ARRAY_SIZE."""
        return [
            commands[i : i + MAX_ARRAY_SIZE]
            for i in range(0, len(commands), MAX_ARRAY_SIZE)
        ]

    def submit_jobs(self, commands: List[str]) -> list[tuple[str, str]]:
        """
        Accepts a flat list of command strings, automatically batches them,
        creates job scripts, and submits the jobs.

        Returns a list of tuples of failed jobs (sbatch command, script path).
        """
        failed_jobs = []

        batches = self._batch_commands(commands)
        self.scripts_dir.mkdir(parents=True, exist_ok=True)
        for i, batch in enumerate(batches):
            script_dir = self.scripts_dir / f"batch_{i:04d}"
            script_dir.mkdir(parents=True, exist_ok=True)
            slurm_output_file = script_dir / "slurm-out.txt"
            slurm_error_file = script_dir / "slurm-err.txt"
            job_outs_dir = script_dir / "job-outs"
            job_errs_dir = script_dir / "job-errs"
            job_out_path_file = script_dir / "job-out-paths.txt"
            job_err_path_file = script_dir / "job-err-paths.txt"
            job_command_file = script_dir / "job-commands.txt"
            job_map_file = script_dir / "job-map.json"
            job_run_file = script_dir / "run.sh"

            job_outs_dir.mkdir(parents=True, exist_ok=True)
            job_errs_dir.mkdir(parents=True, exist_ok=True)

            job_outs_fnames = []
            job_errs_fnames = []
            job_map = {}

            for j, command in enumerate(batch):
                job_out = job_outs_dir / f"job-{j:04d}.out"
                job_err = job_errs_dir / f"job-{j:04d}.err"
                job_outs_fnames.append(str(job_out))
                job_errs_fnames.append(str(job_err))
                job_map[f"job-{j:04d}"] = {
                    "out": str(job_out),
                    "err": str(job_err),
                    "cmd": command,
                }

            with open(job_out_path_file, "w") as f:
                f.write("\n".join(job_outs_fnames + [""]))
            with open(job_err_path_file, "w") as f:
                f.write("\n".join(job_errs_fnames + [""]))
            with open(job_command_file, "w") as f:
                f.write("\n".join(batch + [""]))
            with open(job_map_file, "w") as f:
                json.dump(job_map, f, indent=4)

            header = self.script_header(
                len(batch), slurm_output_file, slurm_error_file
            )
            srun_str = self.script_srun(
                job_out_path_file, job_err_path_file, job_command_file
            )
            run_script_content = "\n".join(
                [header, self.env_setup, srun_str, ""]
            )
            with open(job_run_file, "w") as f:
                f.write(run_script_content)

            print(f"Created script {job_run_file}.")
            if not self.dry_run:
                print("\n" + "-" * 80 + "\n")
                print(f"Running {self.submit_cmd} {job_run_file}")
                run_res = subprocess.run(
                    [self.submit_cmd, str(job_run_file)],
                    stdout=subprocess.PIPE,
                    stderr=subprocess.PIPE,
                    text=True,
                )
                print(f"STDOUT: {run_res.stdout.strip()}")
                print(f"STDERR: {run_res.stderr.strip()}")
                if run_res.returncode != 0:
                    print(f"ERROR: failed to submit {job_run_file}")
                    failed_jobs.append((self.submit_cmd, str(job_run_file)))

        self.rerun_failed_jobs_script(failed_jobs)
        return failed_jobs

    def rerun_failed_jobs_script(
        self, failed_jobs: list[tuple[str, str]]
    ) -> None:
        rerun_script_path = self.scripts_dir / "rerun_failed_jobs.py"

        script_content = RERUN_SCRIPT + f"\n\nmain({failed_jobs})\n"
        with open(rerun_script_path, "w") as f:
            f.write(script_content)
        print(
            "\n\n"
            + "=" * 80
            + "\n"
            + f"Created auto-rerun script: {rerun_script_path}"
        )


def get_common_argparser() -> argparse.ArgumentParser:
    """
    Returns an argparse.ArgumentParser pre-loaded with common multi-job arguments.
    """
    parser = argparse.ArgumentParser(
        add_help=False, formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )
    parser.add_argument(
        "--python-path",
        type=Path,
        default=DEFAULT_PYTHON_PATH,
        help="Path to the python executable to use for the job.",
    )
    parser.add_argument(
        "--time-per-job",
        type=str,
        default="00:07:00",
        help="Time per job in the format D-HH:MM:SS",
    )
    parser.add_argument(
        "--qos",
        type=str,
        default="REDACTED",
        choices=["REDACTED"],
        help="Quality of service to use for the job array.",
    )
    parser.add_argument(
        "--gpu",
        type=str,
        default="rtxa5000:1",
        help="The gpu type and count per job in format NAME:COUNT.",
    )
    parser.add_argument(
        "--gpus",
        type=str,
        nargs="+",
        help="List of acceptable GPUs to use for the job.",
    )
    parser.add_argument(
        "--cpu", type=int, default=4, help="The number of cpu cores per job."
    )
    parser.add_argument(
        "--mem",
        type=int,
        default=32,
        help="The amount of memory per job in G.",
    )
    parser.add_argument(
        "--scripts-dir",
        type=Path,
        required=True,
        help="Directory to save generated job scripts and logs.",
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Create the job script but do not submit it.",
    )
    return parser


def validate_common_args(args: argparse.Namespace) -> None:
    """
    Resolves paths and performs error checking on the common multi-job arguments.
    Also sets a new attribute 'nodes' if --gpus is provided.
    """
    args.python_path = args.python_path.resolve()
    args.scripts_dir = args.scripts_dir.resolve()
    if not args.python_path.exists():
        print(f"WARNING: Python path {args.python_path} not found.")
    while args.scripts_dir.exists():
        print(f"WARNING: Scripts directory {args.scripts_dir} already exists.")
        args.scripts_dir = (
            args.scripts_dir.parent / f"{args.scripts_dir.name}-new"
        )
        print(f"Using new scripts directory: {args.scripts_dir}")
    args.scripts_dir.mkdir(parents=True)
    if args.gpus is not None:
        if args.gpu is not None:
            print("WARN: Both --gpu and --gpus provided. Using --gpus.")
        nodes = _get_nodes_with_gpus(args.gpus)
        if not nodes:
            print(f"ERROR: No nodes found with GPUs {args.gpus}.")
            exit(1)
        args.nodes = nodes
    else:
        args.nodes = None


def _get_nodes_with_gpus(target_gpus: List[str]) -> List[str]:
    """
    Searches for nodes that have any of the specified GPUs.
    (Currently a stub that calls 'show_nodes' and parses its output.)
    """
    print("ERROR: get_nodes_with_gpus not yet functional!")
    result = subprocess.run(["show_nodes"], capture_output=True, text=True)
    nodes = []
    for line in result.stdout.splitlines():
        parts = line.split()
        if len(parts) < 6:
            continue
        node_name = parts[0]
        gpus = parts[-2]
        if any(gpu in gpus for gpu in target_gpus):
            nodes.append(node_name)
    return nodes
