from typing import Dict, List, Callable, Tuple
from slurming.Slurm.shellUtils import make_if_statement
from slurming.SlurmJob.job import SlurmJob
from pathlib import Path


def multicluster_submit(job_name, clusters):
    """
    生成一段嵌入式的 Shell 脚本逻辑。
    逻辑：
    1. 识别当前集群。
    2. 检查其他集群Peer是否有正在运行的同名任务。
    3. 检查本集群是否有超过1个 即除了当前进程外 正在运行的同名任务。
    4. 若满足上述任一退出条件，则在退出前清理所有集群上的 PENDING 任务。
    """
    clusters_str = ",".join(clusters)

    shell_code = f"""
exclusive_check() {{
# ==========================================================
# 自动生成的集群互斥与排队清理脚本
# Job Name: {job_name}
# ==========================================================

JOB_NAME="{job_name}"
ALL_CLUSTERS="{clusters_str}"

# 1. 获取当前集群名称
CURRENT_CLUSTER=$(scontrol show config | grep ClusterName | awk '{{print $3}}')

echo "CURRENT_CLUSTER->$CURRENT_CLUSTER"

# 2. 检查 Peer Clusters (其他集群) 状态
# 统计在其他集群上处于 RUNNING 状态的同名任务
PEER_RUNNING=$(squeue --clusters=$ALL_CLUSTERS --user=$(whoami) --name=$JOB_NAME --states=RUNNING --noheader | grep -v "$CURRENT_CLUSTER" | wc -l || true)

# 3. 检查本集群状态
# 统计在本集群上处于 RUNNING 状态的同名任务数量
LOCAL_RUNNING=$(squeue --clusters=$CURRENT_CLUSTER --user=$(whoami) --name=$JOB_NAME --states=RUNNING --noheader | wc -l || true)

EXIT_REQUIRED=0

if [ "$PEER_RUNNING" -gt 0 ]; then
    echo "[$(date)] Found $PEER_RUNNING instance(s) running on peer clusters. Preparing to exit..."
    EXIT_REQUIRED=1
elif [ "$LOCAL_RUNNING" -gt 1 ]; then
    echo "[$(date)] Found $LOCAL_RUNNING instances running on local cluster ($CURRENT_CLUSTER). Preparing to exit..."
    EXIT_REQUIRED=1
fi

# 4. 如果需要退出，或者为了保持队列整洁，清理所有集群的 PENDING 任务
if [ "$EXIT_REQUIRED" -eq 1 ] || [ "$LOCAL_RUNNING" -eq 1 ]; then
    echo "[$(date)] Cleaning up all PENDING jobs named $JOB_NAME across $ALL_CLUSTERS..."

    # 循环清理每个集群，防止 scancel 在多集群模式下的兼容性问题
    for CLSTR in ${{ALL_CLUSTERS//,/ }}; do
        PIDS=$(squeue --clusters=$CLSTR --user=$(whoami) --name=$JOB_NAME --states=PENDING --format="%i" --noheader || true)
        if [ ! -z "$PIDS" ]; then
            echo "Cancelling PENDING jobs on $CLSTR: $PIDS"
            scancel --clusters=$CLSTR $PIDS
        fi
    done
fi

# 5. 执行退出
if [ "$EXIT_REQUIRED" -eq 1 ]; then
    echo "[$(date)] Conflict detected. Exiting script to avoid duplicate work."
    exit 0
fi

echo "[$(date)] Check passed. No conflicts found. Proceeding with execution on $CURRENT_CLUSTER."
# ==========================================================
}}
exclusive_check
"""
    return shell_code


def relative_to(path: Path, base: Path, abbr: str) -> str:
    """
    Return the relative path of `path` with respect to `base`, using `abbr` as abbreviation.
    """
    rel = path.relative_to(base)
    return f"{abbr}/{rel}"


def init_case(
    shell_make_func: Callable[[Dict, int, int, str | Path], SlurmJob],
    config: Dict,
    commit_home: Path,
    shell_name: str,
    cwd: Path,
    cluster: str,
    list_of_clusters: List[str] = [],
):
    """
    Initialize a case with the given configuration and shell path.

    Args:
        config (Dict): The configuration dictionary for the case.
        shell_path (Path): The path where the shell script will be created.
    """

    src_folder = commit_home / "src"
    shell_folder = commit_home / "shell"
    sink_folder = commit_home / "storage"
    case_dir = commit_home / "cases"
    log_dir = commit_home / "log"

    script_path = shell_folder / f"{shell_name}_{cluster}.sh"
    script_abbr = f"${{SHELL_SCRIPT}}"

    num_cpus = config.pop("num_cpus")
    num_gpus = config.pop("num_gpus")
    time = config.pop("time")
    requeue: bool = config.pop("requeue")
    # requeue = False
    case_dir_path = case_dir / shell_name
    case_storage_dir = sink_folder / shell_name

    slurm_job: SlurmJob = shell_make_func(
        config,
        num_cpus,
        num_gpus,
        r"${SHELL_SCRIPT}",
    )
    slurm_job.env_vars.update({
        "SLURM_CPU_PER_TASK": str(num_cpus),
        "SLURM_GPU_PER_TASK": str(num_gpus),
        "RUNTIME_DIR": str(case_dir_path),
        "STORAGE_DIR": str(case_storage_dir),
        "SHELL_SCRIPT": str(script_path),
    })

    slurm_job.cpus_per_task = num_cpus
    slurm_job.gpus = num_gpus
    slurm_job.hours, slurm_job.minutes, slurm_job.seconds = map(int, time.split(":"))

    slurm_job.job_name = shell_name

    if (requeue):
        # Add function brackets
        slurm_job.content.insert(0, "main() {")
        slurm_job.content.append("}")
        # Call main function
        slurm_job.content.append("main &")
        # Wait for all background processes to finish
        slurm_job.content.extend([
            "job_pid=$!",
            "wait $job_pid",
            *make_if_statement(if_st=(
                [f"$? -ne 0"],
                [
                    "echo 'One of the commands failed.'",
                    "#NEXT_JOB_PLACEHOLDER#",
                    "exit 1",
                ],
            ), ),
        ])
    # Convert last user command to background command
    store_symlink = []
    store_output = []
    contain_workdir = any([len(x) == 0 for x in slurm_job.output_storage])

    dir_paths = [case_storage_dir, case_dir_path]
    runtime_abbr = r"${RUNTIME_DIR}"
    storage_abbr = r"${STORAGE_DIR}"
    runtime_abbr_safe = r"${RUNTIME_DIR:?}"
    storage_abbr_safe = r"${STORAGE_DIR:?}"
    if contain_workdir:
        # Case A: 整个工作目录同步

        store_symlink.append(f"ln -sf {runtime_abbr} {storage_abbr}")
        store_output.extend([])

    else:

        for output_file in slurm_job.output_storage:
            store_output_path = case_dir_path / output_file
            store_source_path = case_storage_dir / output_file

            relative_output_path = relative_to(store_output_path, case_dir_path, runtime_abbr)
            relative_source_path = relative_to(store_source_path, case_storage_dir, storage_abbr)

            parent_relative_output_path = relative_to(store_output_path.parent, case_dir_path, runtime_abbr)
            parent_relative_source_path = relative_to(store_source_path.parent, case_storage_dir, storage_abbr)

            if ("/" == output_file[-1]):

                mkdir: List[str] = []
                if (store_output_path.parent not in dir_paths):
                    mkdir.append(parent_relative_output_path)
                if (store_source_path not in dir_paths):
                    mkdir.append(relative_source_path)
                store_symlink.extend([
                    f"mkdir -p {' '.join(mkdir)}" if len(mkdir) > 0 else "",
                    f"ln -sf {relative_source_path} {relative_output_path}",
                ])
                store_output.extend([])
            else:
                mkdir: List[str] = []
                if (store_output_path.parent not in dir_paths):
                    mkdir.append(parent_relative_output_path)
                if (store_source_path.parent not in dir_paths):
                    mkdir.append(parent_relative_source_path)
                store_symlink.extend([
                    f"mkdir -p {' '.join(mkdir)}" if len(mkdir) > 0 else "",
                    f"touch {relative_source_path} && ln -sf {relative_source_path} {relative_output_path}",
                ])
                store_output.extend([])

    slurm_job.append(store_output)

    timeout_trap = []
    if requeue:
        timeout_trap.extend([
            "timeoutTrap() {",
            "  echo 'Caught SIGTERM, requeuing...';",
            # "  scontrol requeue $SLURM_JOB_ID;",
            f"  sbatch --begin=now+30 {script_abbr}",
            "  #NEXT_JOB_PLACEHOLDER#",
            "  exit 0",
            "}",
            "trap timeoutTrap SIGTERM",
        ])
    slurm_job.prepend([
        multicluster_submit(shell_name, list_of_clusters),
        *timeout_trap,
        f"rm -rf {runtime_abbr_safe} || true",
        f"mkdir -p {runtime_abbr}",
        f"cp -rf {src_folder}/* {runtime_abbr}",
        f"cd {runtime_abbr}",
        f"mkdir -p {storage_abbr}",
        f"find {storage_abbr} -maxdepth 1 -mindepth 1 -type l -delete || true",
        *store_symlink,
    ])

    slurm_job.append([
        # 如果 case_dir_path 是一个符号链接，使用 rm 删除它，否则使用 rm -rf 删除它
        f"if [ -L \"{runtime_abbr}\" ]; then rm {runtime_abbr_safe}; else rm -rf {runtime_abbr_safe}; fi",
        "#NEXT_JOB_PLACEHOLDER#",
        f"echo 'DONE' > {storage_abbr}/DONE",
        f"rm {storage_abbr_safe}/roll.json || true",
    ])
    slurm_job.prepend([
        f"rm -rf {storage_abbr_safe}/DONE || true",
    ])

    slurm_job.python_home = Path("$(poetry env info --path)")
    slurm_job.working_dir = cwd

    # Requeue timeout jobs
    if requeue:
        slurm_job.sbatch_args["requeue"] = True
    slurm_job.sbatch_args["signal"] = "B:SIGTERM@60"
    shell_folder.mkdir(exist_ok=True, parents=True)

    slurm_job.write(
        job_name=shell_name,
        script_path=script_path,
        log_file=log_dir / f"{shell_name}.log",
        delete_log_on_completion=True,
        delete_script_on_completion=True,
    )
