# from functools import partial
from functools import partial
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,
    pjsub_large_spec,
    pjsub_small_spec,
)


_COMMAND_PREFIX = "#PJM"


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

# --- Core Job Specifications ---"""
    ]
    # 2. Name
    lines.append(f"{_COMMAND_PREFIX} -N {spec.job_name}")
    # 3. Resources (node exclusive)
    lines.append(f'{_COMMAND_PREFIX} -L "elapse={spec.walltime}"')
    lines.append(f'{_COMMAND_PREFIX} -L "node={spec.resource.nodes}"')
    if spec.resource.memory_gb is not None:
        lines.append(f'{_COMMAND_PREFIX} -L "node-mem={spec.resource.memory_gb}"')
    if spec.resource.gpus_per_node is not None:
        lines.append(f'{_COMMAND_PREFIX} -L "gpu={spec.resource.gpus_per_node}"')

    if spec.queue is not None:
        lines.append(f'{_COMMAND_PREFIX} -L "rscqrp={spec.queue}"')
    # 4. User settings
    lines.append("\n# --- User Details ---")
    if spec.account is not None:
        lines.append(f"{_COMMAND_PREFIX} -g {spec.account}")
    if spec.user is not None:
        if spec.user.mail is not None:
            lines.append(f"{_COMMAND_PREFIX} -m e")  # end only
            lines.append(f"{_COMMAND_PREFIX} --mail-list {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} -e {spec.error_path}")
    else:
        lines.append(f"{_COMMAND_PREFIX} -j")
    if spec.output_path is None:
        spec.output_path = Path(f"./{spec.job_name}.o")
    lines.append(f"{_COMMAND_PREFIX} -o {spec.output_path}")

    # 6. MPI
    if spec.mpi is not None:
        lines.append("\n# --- MPI ---")
        if spec.mpi.proc is not None:
            lines.append(f'{_COMMAND_PREFIX} --mpi "proc={spec.mpi.proc}"')
        if spec.mpi.max_proc_per_node is not None:
            lines.append(
                f'{_COMMAND_PREFIX} --mpi "max-proc-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 ---")
    for key, value in spec.environment.items():
        lines.append(f'{_COMMAND_PREFIX} -x "{key}"="{value}"')

    # 9. Container logic

    # 10. User Command, (prologue), command, (epilogue)
    lines.append("\n# --- User Command ---")
    lines.append(spec.command)

    with open("./script", "w+") as fh:
        fh.write("\n".join(lines))

    return "\n".join(lines)


class PJSUBExecutor:
    def __init__(
        self,
        registry_path: Path | None,
        logs_path: Path,
        spec: JobSpec,
        command: str = "pjsub",
    ) -> 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_pjsub_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)


PJSUB_EXECUTOR_SMALL = partial(PJSUBExecutor, spec=pjsub_small_spec())
PJSUB_EXECUTOR_LARGE = partial(PJSUBExecutor, spec=pjsub_large_spec())
