from pathlib import Path
from typing import Callable
from tierkreis.controller.executor.hpc.hpc_executor import run_hpc_executor
from tierkreis.controller.executor.hpc.job_spec import JobSpec


_COMMAND_PREFIX = "#SBATCH"


def generate_slurm_script(spec: JobSpec) -> str:
    # 1. Shebang and file header
    lines = [
        """#!/bin/bash
#
# SLURM Job Script generated by TIERKREIS

# --- Core Job Specifications ---"""
    ]
    # 2. Name
    lines.append(f"{_COMMAND_PREFIX} --job-name={spec.job_name}")
    # 3. Resources (node exclusive)
    lines.append(f"{_COMMAND_PREFIX} --time={spec.walltime}")
    lines.append(f"{_COMMAND_PREFIX} --nodes={spec.resource.nodes}")
    if spec.resource.cores_per_node is not None:
        lines.append(
            f"{_COMMAND_PREFIX} --cpus-per-task={spec.resource.cores_per_node}"
        )
    if spec.resource.memory_gb is not None:
        lines.append(f"{_COMMAND_PREFIX} --mem={spec.resource.memory_gb}G")
    if spec.resource.gpus_per_node is not None:
        lines.append(f"{_COMMAND_PREFIX} --gpus-per-node={spec.resource.gpus_per_node}")
    if spec.queue is not None:
        lines.append(f"{_COMMAND_PREFIX} --partition={spec.queue}")

    # 4. User settings
    lines.append("\n# --- User Details ---")
    if spec.account is not None:
        lines.append(f"{_COMMAND_PREFIX} --account={spec.account}")
    if spec.user is not None:
        if spec.user.mail is not None:
            lines.append(f"{_COMMAND_PREFIX} --mail-type=END")  # end only
            lines.append(f"{_COMMAND_PREFIX} --mail-user={spec.user.mail}")

    # 5. Output and Error handling
    lines.append("\n# --- Output and Error Handling ---")
    if spec.error_path is not None:
        lines.append(f"{_COMMAND_PREFIX} --error={spec.error_path}")
    if spec.output_path is not None:
        lines.append(f"{_COMMAND_PREFIX} --output={spec.output_path}")

    # 6. MPI, #TODO check if this makes sense
    if spec.mpi is not None:
        lines.append("\n# --- MPI ---")
        if spec.mpi.proc is not None:
            lines.append(f"{_COMMAND_PREFIX} --ntasks={spec.mpi.proc}")
        if spec.mpi.max_proc_per_node is not None:
            lines.append(
                f"{_COMMAND_PREFIX} --ntasks-per-node={spec.mpi.max_proc_per_node}"
            )

    # 7. User specific
    lines.append("\n# --- User-Specific Arguments ---")
    for key, value in spec.extra_scheduler_args.items():
        lines.append(f"{_COMMAND_PREFIX} {key} {value if value is not None else ''}")

    # 8. Environment
    lines.append("\n# --- Environment Setup ---")
    if spec.environment != {}:
        env = ",".join(
            f"{key}={value if value else '""'}"
            for key, value in spec.environment.items()
        )
        lines.append(f"--export={env}")
    # 9. Container logic

    # 10. User Command, (prologue), command, (epilogue)
    lines.append("\n# --- User Command ---")
    if spec.mpi is not None:
        if spec.mpi.max_proc_per_node is None:
            spec.mpi.max_proc_per_node = 1
        lines.append(
            f"mpirun -n {spec.resource.nodes * spec.mpi.max_proc_per_node} {spec.command}"
        )
    else:
        lines.append(spec.command)

    return "\n".join(lines)


class SLURMExecutor:
    def __init__(
        self,
        registry_path: Path | None,
        logs_path: Path,
        spec: JobSpec,
        command: str = "sbatch",
    ) -> None:
        self.launchers_path = registry_path
        self.logs_path = logs_path
        self.errors_path = logs_path
        self.spec = spec
        self.script_fn: Callable[[JobSpec], str] = generate_slurm_script
        self.command = command

    def run(self, launcher_name: str, worker_call_args_path: Path) -> None:
        self.errors_path = (
            self.logs_path.parent.parent / worker_call_args_path.parent / "errors"
        )
        run_hpc_executor(self, launcher_name, worker_call_args_path)
