from tqdm import trange
import json
import subprocess
from typing import List, Dict
import pandas as pd
import numpy as np
from pathlib import Path
import re
import os

import math
from investigation import investigation_home, cwd, storage_path, entry_path

from investigation.doeargs.args2config import args2config
from investigation.doe.generate import config2jobs
from investigation.utils import get_hashname
from investigation.shell.case import init_case
from investigation.database.dataframe import format_and_Q_eq_in
from investigation.database.house import scan_cases
from investigation.database.query import filter_entry
from slurming.Slurm.slurm import get_job_dict
from investigation.command.utils import global_io_lock


def init_entry(file_name: str):
    from investigation.entry import LIST_OF_ARGS

    doe = args2config(LIST_OF_ARGS)

    job_file_path = cwd / f"{file_name}.json"

    if (job_file_path.exists()):
        print(
            f"There exists an uncommited investigation job {job_file_path}, please finish and commit it or erase it to stare a new job."
        )
        exit(1)
    try:
        job_file_path.write_text(json.dumps(doe))
    except Exception as e:
        print(f"Error writing investigation job file: {e}")
        print(f"Problem caused by {doe}")
        exit(1)

    print(f"Initialized investigation DoE job {job_file_path}.")
    print(f"Run 'control commit' when ready.")


def sample_entry(commit_hash: str, force: bool, cases: Dict[str, Dict]):
    from investigation.entry import LIST_OF_CLUSTERS, hostname, parameter2SlurmJob

    if (force):
        jobs = list(cases.keys())
    else:
        jobs = []
        for job_hash, job_config in cases.items():
            commit_hash = job_config.get("commit_hash", commit_hash)

            # commit_folder = storage_path / commit_hash
            # sink_dir = commit_folder / "storage"
            # existing_cache = list(sink_dir.glob(f"{job_hash}*"))
            # if (len(existing_cache) > 0):
            #     continue
            job_state = get_job_state(commit_hash, job_hash, job_config)
            if (job_state in ["DONE"]):
                continue

            jobs.append(job_hash)
    for job_hash in jobs:

        job = cases[job_hash]
        commit_folder = storage_path / job.pop("commit_hash", commit_hash)
        job.pop("batch_hash", None)

        init_case(
            shell_make_func=parameter2SlurmJob,
            config={
                **job,
                "job_name": job_hash,
            },
            commit_home=commit_folder,
            shell_name=f"{job_hash}",
            cwd=cwd,
            list_of_clusters=LIST_OF_CLUSTERS,
            cluster=hostname,
        )
    return len(jobs)


def commit_entry(file_name: str, commit_hash: str = ""):
    if (not storage_path.exists()):
        print(f"Making {storage_path}...")
        storage_path.mkdir(exist_ok=True, parents=True)

    print(f"Reading current project HEAD hash")
    if (commit_hash == ""):
        commit_hash = subprocess.check_output(["git", "rev-parse", "--short", "HEAD"]).decode().strip()
    job_file_path = cwd / f"{file_name}.json"

    commit_folder = storage_path / commit_hash
    if (not commit_folder.exists()):

        print(f"Making current project HEAD Home {commit_folder}")

        commit_folder.mkdir(parents=True, exist_ok=True)

    src_folder = commit_folder / "src"

    if (not src_folder.exists()):
        print(f"Cloning current project HEAD to {src_folder}")

        subprocess.check_call(cwd=commit_folder, args=["git", "clone", str(cwd), str(src_folder)])
        subprocess.check_call(cwd=src_folder, args=["git", "checkout", commit_hash])
        subprocess.check_call(cwd=src_folder, args=["rm", "-rf", ".git"])

    shell_folder = commit_folder / "shell"

    sink_folder = commit_folder / "storage"

    case_dir = commit_folder / "cases"
    log_dir = commit_folder / "log"

    for folder in [shell_folder, sink_folder, case_dir, log_dir]:
        if (not folder.exists()):
            folder.mkdir(parents=True, exist_ok=True)

    # Read pending investigation json
    if (not job_file_path.exists()):
        print(f"Pending investigation path {job_file_path} does not exist, please run 'control init' first.")
        exit(1)

    pending_investigation: Dict = json.loads(job_file_path.read_text())

    meta_json_path = commit_folder / "meta.json"

    if (not meta_json_path.exists()):
        meta_json = {
            "seed": {
                "type": str(int),
                "default": 1,
                "value": 1,
            }
        }
        with open(meta_json_path, "w") as f:
            json.dump(meta_json, f, indent=4)

    meta_json = json.loads(meta_json_path.read_text())
    pending_investigation["meta"].update(meta_json)

    print(f"Generating jobs from pending investigation job")
    segs, jobs = config2jobs(pending_investigation)
    print(f"Reading existing jobs from {commit_folder}/*.json")

    existing_jobs = scan_cases(commit_hash)

    job_df = pd.DataFrame.from_dict(existing_jobs, orient='index')
    job_df.index.name = "job_name"

    # input(job_df)
    # Exclude existing jobs from cases, where each case is a dictionary [parameter: value]
    if (len(job_df) == 0):
        print(f"No existing jobs found, all cases will be new.")
        new_jobs = jobs
    else:
        yes_to_all = False

        def find_exact_match(job: Dict, yta: bool) -> bool:
            job_df_key_diff = set(job.keys()).difference(set(job_df.columns))
            if (len(job_df_key_diff) > 0):
                if (yta):
                    return False
                resume = input(
                    f"Warning: The existing jobs have no columns {job_df_key_diff}, which are present in the new job {job}.\nThis could potentially due to any recently added parameters.\nDo you want to continue? (yes/No/all)"
                )
                match resume.strip().lower():
                    case "yes":
                        return False
                    case "all":
                        nonlocal yes_to_all
                        yes_to_all = True
                        return False
                    case _:
                        print(f"Aborting commit.")
                        exit(1)
            filter_keys = set(job.keys()) - set(["num_cpus", "num_gpus", "time", "requeue"])
            # print(f"Checking existing jobs for match on keys {filter_keys}...")
            filter_job_df = job_df[list(filter_keys)]

            # Find rows in filter_job_df that match all key-value pairs in job
            job_filter = {key: job[key] for key in filter_keys}
            query_str, local_dict = format_and_Q_eq_in(job_filter)
            flags = filter_job_df.query(query_str, local_dict=local_dict)

            return len(flags) > 0

        new_jobs = list(filter(lambda job: not find_exact_match(job, yes_to_all), jobs))

    if (len(new_jobs) == 0):
        print(f"No new cases to commit, exiting.")
        return
    else:
        print(f"Found {len(new_jobs)} new cases to commit.")
        # Assign a unique name to each job
        new_job_dict = {}
        for job in new_jobs:
            job_hash = get_hashname(8)
            while (job_hash in new_job_dict or job_hash in existing_jobs):
                job_hash = get_hashname(8)
            new_job_dict[job_hash] = job

        # Write new jobs to commit_folder with a unique name
        config_hash = get_hashname(8)
        (commit_folder / f"{config_hash}.json").write_text(json.dumps(new_job_dict, indent=4))
        (commit_folder / f"{config_hash}_config.json").write_text(json.dumps(pending_investigation, indent=4))

    print(f"Writing new jobs to {shell_folder}")

    case_df, cases = filter_entry(commit_hash=commit_hash, batch_hash=[config_hash])

    pending_shells = sample_entry(commit_hash, force=True, cases=cases)
    print(f"Generated {pending_shells} pending shell scripts.")


def batch_submit(jobs: List[Path], parellels: int = -1):

    def wipe_tail(job: Path):
        lines = job.read_text().splitlines()
        filtered_lines = []
        insert_position = -1
        for line in lines:
            if re.match(r'^sbatch ".*', line):
                continue
            filtered_lines.append(line)
            if (re.match(r'#NEXT_JOB_PLACEHOLDER#', line)):
                insert_position = len(filtered_lines)
        return filtered_lines, insert_position

    if (len(jobs) == 0):
        print(f"No un-run jobs found in {storage_path}, exiting.")
        return
    print(f"Found {len(jobs)} un-run jobs in {storage_path}, starting to run them.")

    if (parellels <= 0):
        parellels = len(jobs)
        print(f"Running all {len(jobs)} jobs at once.")

    for i in trange(int(np.ceil(len(jobs) / parellels))):

        batch = jobs[i * parellels:(i + 1) * parellels]

        n_batch = jobs[(i + 1) * parellels:(i + 2) * parellels]

        for i in range(len(n_batch)):
            c_job = batch[i]
            n_job = n_batch[i]

            # Remove all lines in c_job following the regex ^sbatch ".*
            filtered_lines, i_p = wipe_tail(c_job)
            filtered_lines = [
                *filtered_lines[:i_p],
                f"sbatch \"{n_job}\" || echo \"Failed to submit {n_job}\"",
                *filtered_lines[i_p:],
            ]

            c_job.write_text('\n'.join(filtered_lines))
        for i in range(len(n_batch), len(batch)):
            batch[i].write_text('\n'.join(wipe_tail(batch[i])[0]))

    for job in jobs[:parellels]:
        os.system(f"sbatch {job}")


def get_job_state(commit_hash: str, case_hash: str, case_config: Dict) -> str:
    commit_path: Path = storage_path / case_config.get("commit_hash", commit_hash)
    case_stor_path = commit_path / "storage"
    case_log_path = commit_path / "log" / f"{case_hash}.log"
    case_run_path = commit_path / "cases" / case_hash
    shell_path = commit_path / "shell" / f"{case_hash}.sh"
    shell_paths = list(commit_path.glob(f"shell/{case_hash}_*.sh"))

    DONE_path = commit_path / "storage" / f"{case_hash}" / "DONE"
    runtime_files = [x for x in [case_log_path, case_run_path] if x.exists()]

    if (len(runtime_files) > 0):
        return "STARTED"
    if (DONE_path.exists()):
        return "DONE"

    if (shell_path.exists() or len(shell_paths) > 0):
        return "SAMPLED"

    return "PENDING"


def run_entry(commit_hash: str, parellels: int, force: bool, cases: Dict[str, Dict], cluster: str):
    jobs = []

    for job_hash, job_config in cases.items():
        commit_hash = job_config.get("commit_hash", commit_hash)

        commit_folder = storage_path / commit_hash
        if (not force):
            # existing_cache = list((commit_folder / "storage").glob(f"{job_hash}*"))
            # existing_cache = [x for x in existing_cache if not x.is_symlink()]
            # if (len(existing_cache) > 0):
            #     continue
            job_state = get_job_state(commit_hash, job_hash, job_config)
            if (job_state in ["UNINITIALIZED", "FAILED", "DONE", "PENDING"]):
                continue

        shell_dir = commit_folder / "shell"
        jobs.append(shell_dir / f"{job_hash}_{cluster}.sh")

    batch_submit(jobs, parellels=parellels)


def git_check():

    # print(f"Detect git registration...")
    try:
        subprocess.check_output(["git", "rev-parse", "--is-inside-work-tree"], cwd=cwd)
    except subprocess.CalledProcessError:
        print("Current directory is not a valid git repository. Please initialize git first.")
        exit(1)
    gitignore_path = cwd / ".gitignore"
    if (gitignore_path.exists()):
        gitignore = gitignore_path.read_text()
    else:
        gitignore = ""

    gitignore_add: List[str] = []
    for ignore_path in [
            entry_path,
            storage_path,
    ]:
        relative_ignore_path = str(ignore_path.relative_to(cwd))
        if (relative_ignore_path in gitignore):
            continue
        print(f"Ignoring {relative_ignore_path} in .gitignore")

        gitignore_add.append(relative_ignore_path.strip())
    if (len(gitignore_add) == 0):
        return
    print(f"Adding {' '.join(gitignore_add)} to .gitignore")
    with open(gitignore_path, 'a') as fout:
        fout.write("\n".join(gitignore_add))
    os.system(f"git add {gitignore_path}")


def check_entry():

    if (not entry_path.exists()):

        response = input(
            f"{entry_path} is not found, you are expected to define list of args in it. Do you want to let investigation self-check and then initialize ./entry/__init__.py? (yes/No)"
        )
        if (response.strip().lower() != "yes"):
            print(f"{entry_path} is not found, failed to generate DoE.")
            exit(1)

        print(f"Making {entry_path.parent}")
        entry_path.parent.mkdir(exist_ok=True, parents=True)

        print(f"Making {entry_path}")
        entry_path.write_text((investigation_home / "investigation" / "entry" / "__init__.py").read_text())

        print(f"Please implement {entry_path} and re-run.")

        exit(0)


def locate_entry(commit_hash: str, case_hash: str, case_config: Dict):

    commit_path: Path = storage_path / case_config.get("commit_hash", commit_hash)
    batch_hash = case_config.get("batch_hash", "????????")
    batch_found = ""
    for config_json in commit_path.glob(f"{batch_hash}.json"):
        case_configs = json.loads(config_json.read_text())
        if (case_hash in case_configs):
            batch_found = config_json.stem
            break
    assert batch_found != "", f"Cannot locate case {case_hash} in commit {commit_hash}."
    return batch_found


def delete_entry(commit_hash: str, case_hash: str, case_config: Dict):

    commit_path: Path = storage_path / case_config.get("commit_hash", commit_hash)
    batch_hash = case_config.get("batch_hash", "????????")
    for config_json in commit_path.glob(f"{batch_hash}.json"):
        case_configs = json.loads(config_json.read_text())
        if (case_hash in case_configs):
            del case_configs[case_hash]
            # print(f"Deleting case {case_hash} from {config_json}")
            if (len(case_configs) == 0):
                os.system(f"rm -rfv {config_json}")
            else:
                config_json.write_text(json.dumps(case_configs))

    for case_json in commit_path.glob(f"*/{case_hash}*"):
        os.system(f"rm -rfv {case_json}")


def insert_to_entry(
    commit_hash: str,
    case_hash: str,
    case_config: Dict,
    new_config: Dict,
    dry_run: bool = False,
):

    commit_path: Path = storage_path / case_config.get("commit_hash", commit_hash)
    batch_hash = case_config.get("batch_hash", "????????")
    for config_json in commit_path.glob(f"{batch_hash}.json"):
        case_configs = json.loads(config_json.read_text())
        if (case_hash in case_configs):
            entry_dict = case_configs[case_hash]
            assert set(new_config.keys()).isdisjoint(set(entry_dict.keys(
            ))), f"New config keys {set(new_config.keys())} already exist in the case config {set(entry_dict.keys())}."

            if (dry_run):
                print(f"Will update case {case_hash} in {case_configs[case_hash]} with {new_config}.\n")
                continue

            case_configs[case_hash].update(new_config)
            config_json.write_text(json.dumps(case_configs))


def reset_entry(commit_hash: str, case_hash: str, case_config: Dict):

    commit_path: Path = storage_path / case_config.get("commit_hash", commit_hash)

    for case_path in commit_path.glob(f"*/{case_hash}*"):
        if (case_path.suffix == ".sh"):
            continue
        os.system(f"rm -rfv {case_path}")


def clean_runtime_cache(commit_hash: str, case_hash: str, case_config: Dict, dry_run: bool = False):

    commit_path: Path = storage_path / case_config.get("commit_hash", commit_hash)

    log_path = commit_path / "log" / f"{case_hash}.log"
    runtime_path = commit_path / "cases" / case_hash

    toBeRemoved = [x for x in [log_path, runtime_path] if x.exists()]
    if (dry_run):
        for p in toBeRemoved:
            print(f"Will remove {p}")
        return
    for p in toBeRemoved:
        os.system(f"rm -rfv {p}")


def sync_src_to_commit(base_folder: Path, commit_hash: str):

    if (not storage_path.exists()):
        print(f"Making {storage_path}...")
        storage_path.mkdir(exist_ok=True, parents=True)

    commit_folder = storage_path / commit_hash
    if (not commit_folder.exists()):

        print(f"Making current project HEAD Home {commit_folder}")

        commit_folder.mkdir(parents=True, exist_ok=True)

    src_folder = commit_folder / "src"
    case_dir = commit_folder / "cases"

    if (not src_folder.exists()):
        print(f"Cloning current project HEAD to {src_folder}")

        subprocess.check_call(cwd=commit_folder, args=["git", "clone", str(cwd), str(src_folder)])
        subprocess.check_call(cwd=src_folder, args=["git", "checkout", commit_hash])
        subprocess.check_call(cwd=src_folder, args=["rm", "-rf", ".git"])
    else:
        for folder in src_folder.iterdir():
            if (folder.name == ".git"):
                continue
            if (folder.is_dir()):
                source_path = str(base_folder / folder.name)
                # print(f"rm -rf {folder.name} && cp -rf {source_path} {folder.name} in {src_folder}")
                subprocess.check_call(cwd=src_folder, args=["rm", "-rf", folder.name])
                subprocess.check_call(cwd=src_folder, args=["cp", "-rf", source_path, folder.name])
                for case_path in case_dir.iterdir():
                    if (case_path.is_dir()):
                        subprocess.check_call(cwd=case_path, args=["rm", "-rf", folder.name])
                        subprocess.check_call(cwd=case_path, args=["cp", "-rf", source_path, folder.name])
                        # print(f"rm -rf {folder.name} && cp -rf {source_path} {folder.name} in {case_path}")


def entry_stat(
    commit_hash: str,
    batch_hash: List[str],
    cluster_name: str,
    all_cluster: bool,
    unknown_args,
):
    case_df, cases = filter_entry(
        commit_hash=commit_hash,
        batch_hash=batch_hash,
        **unknown_args,
    )
    active_jobs, running_jobs, queue_jobs = get_job_dict(cluster_name)

    # for cluster, job_dict in active_jobs.items():
    #     all_job_keys.extend(list(job_dict.keys()))

    # joint_keys = set(cases.keys()).intersection(set(all_job_keys))

    # for cluster, job_dict in active_jobs.items():
    #     for key in list(job_dict.keys()):
    #         if key not in cases:
    #             job_dict.pop(key)
    # for key in joint_keys:
    #     cases.pop(key)
    for key in list(running_jobs.keys()):
        if (key not in cases):
            running_jobs.pop(key)
    for key in list(queue_jobs.keys()):
        if (key not in cases):
            queue_jobs.pop(key)

    all_pending_jobs = {
        k: v
        for cluster, job_dict in active_jobs.items()
        for k, v in job_dict.items() if "PENDING" == v['state'] and k in cases
    }
    if (all_cluster):
        for key in list(all_pending_jobs.keys()) + list(running_jobs.keys()):
            cases.pop(key, None)
        queue_jobs = all_pending_jobs
    else:
        for key in list(queue_jobs.keys()) + list(running_jobs.keys()):
            cases.pop(key, None)
    job_states = {x: get_job_state(commit_hash, x, cases[x]) for x in cases.keys()}
    # running_jobs = {
    #     k: v
    #     for cluster, job_dict in active_jobs.items()
    #     for k, v in job_dict.items() if "RUNNING" == v['state']
    # }
    job_stat_dict = {
        "RUNNING": list(running_jobs.keys()),
        "QUEUED": list(queue_jobs.keys()),
        "PENDING": list(filter(lambda x: job_states[x] == 'PENDING', job_states.keys())),
        "DONE": list(filter(lambda x: job_states[x] == 'DONE', job_states.keys())),
        "FAILED": list(filter(lambda x: job_states[x] == 'STARTED', job_states.keys())),
        "SAMPLED": list(filter(lambda x: job_states[x] == 'SAMPLED', job_states.keys())),
    }

    return job_stat_dict


def roll_script(commit_hash: str, case_hash: str):
    from investigation.entry import LIST_OF_CLUSTERS, hostname, CLUSTER_CONFIG
    """
    能力导向的多集群重试脚本。

    逻辑：
    1. 读取全局状态：任务目前至少需要多少 RAM (GB) 和 VRAM (GB)。
    2. 检查本地日志：如果上次是在本机跑挂的，分析错误并提升全局 RAM/VRAM 需求。
    3. 能力评估：计算本地为了满足全局 RAM/VRAM 需求，需要分配多少 CPU/GPU。
    4. 适格性判断：
       - 如果本地硬件上限 < 需求：标记本机 Abandoned(无能为力）。
       - 如果本地硬件上限 >= 需求：更新 Config 并提交 Slurm(尝试接单）。
    """

    # ---------------- Step 0: 环境与路径加载 ----------------
    commit_dir = storage_path / commit_hash
    log_file = commit_dir / "log" / f"{case_hash}.log"
    roll_log_file = commit_dir / "cases" / f"{case_hash}" / "roll.json"
    batch_config_path = commit_dir / f"{locate_entry(commit_hash, case_hash, {})}.json"

    # ---------------- Step 2: 获取本地硬件规格 ----------------
    assert batch_config_path.exists(
    ), f"Batch config {batch_config_path} not found for case {case_hash} in commit {commit_hash}."
    with open(batch_config_path, 'r') as f:
        batch_config = json.load(f)
    job_config = batch_config[case_hash]

    # ---------------- Step 1: 读取/初始化共享状态 (原子操作) ----------------
    with global_io_lock(roll_log_file):

        # 确定使用 GPU 还是 CPU 分区的配置
        use_gpu = job_config.get("num_gpus", 0) > 0
        local_config_key = f"{hostname}-gpu" if use_gpu else f"{hostname}-cpu"

        cluster_limits = CLUSTER_CONFIG.get(local_config_key, {})
        local_mem_per_cpu = cluster_limits["mem_per_cpu"]  # 防止除零
        local_mem_per_gpu = cluster_limits["mem_per_gpu"]

        if roll_log_file.exists() and roll_log_file.stat().st_size > 0:
            with open(roll_log_file, 'r') as f:
                roll_log = json.load(f)
        else:
            roll_log = {
                "progress": {
                    "min_ram_gb": job_config["num_cpus"] * local_mem_per_cpu,  # 全局知识：目前已知的最小内存需求
                    "min_vram_gb": job_config["num_gpus"] * local_mem_per_gpu,  # 全局知识：目前已知的最小显存需求
                },
                "clusters": {
                    x: {
                        "attempts": 0,
                        "abandoned": False
                    }
                    for x in LIST_OF_CLUSTERS
                }
            }

        # 检查自己是否已经彻底放弃过该任务（硬件不够格）
        if roll_log["clusters"][hostname]["abandoned"]:
            print(f"Cluster {hostname} previously abandoned {case_hash}. Skipping.")
            return

        # ---------------- Step 3: 错误分析与全局需求更新 ----------------
        # 只有 Log 存在且包含错误时才更新。
        # 这里假设 Log 是最近一次运行产生的。
        policy_triggered = False
        if log_file.exists():
            log_content = log_file.read_text()

            cluster_running_this = log_content.find(f"CURRENT_CLUSTER->{hostname}") != -1
            # 策略：错误关键字 -> (资源类型, 资源增量 count)
            # 注意：我们先把 count 转换成 GB 更新到全局 progress
            POLICY = {
                "CUDA out of memory": ("num_gpus", 1),  # +1 GPU
                "oom_kill": ("num_cpus", 5),  # +2 CPU (for RAM)
                "Out of memory": ("num_cpus", 5),
            }

            # 检查这是否是上次运行刚产生的新 Log (通过对比 attempts 或者简单地检查错误)
            # 为了简化，我们假设如果 Log 里有 OOM，且当前的配置还小于 Log 暗示的需求，就需要更新

            for error_key, (res_type, increment) in POLICY.items():
                policy_triggered = policy_triggered or (error_key in log_content)
                if (not cluster_running_this):
                    continue
                if error_key in log_content:
                    # 计算如果按当前 config 跑挂了，加了 increment 之后应该是多少 GB
                    if res_type == "num_cpus":
                        # 估算挂掉时的内存：当前配置的 CPU 数 * 本地单核内存
                        # 注意：这里可能有点模糊，因为不知道挂掉的是哪个集群。
                        # 但如果 Log 里有 OOM，说明"之前试过的最大值"都不够。
                        # 稳健策略：基于当前 Global 值增加

                        # 假设每次 OOM 至少需要增加 2 * 2GB (假设基准) = 4GB 步进，或者直接基于当前 job_config
                        # 更科学的方法：如果是本机跑挂的，用本机的 mem_per_cpu 计算。
                        # 如果不知道谁跑挂的，直接在 Global 上加一个固定步长 (例如 8GB)

                        # 简单起见：每次 OOM，内存需求 + 4GB，显存 + 10GB (或根据 increment)
                        increase_amount = increment * local_mem_per_cpu  # 尽力估算

                        roll_log["progress"]["min_ram_gb"] += increase_amount
                        print(
                            f"[Roll] OOM Detected. Increasing Global RAM Req to {roll_log['progress']['min_ram_gb']}GB")

                    elif res_type == "num_gpus":
                        increase_amount = increment * local_mem_per_gpu
                        roll_log["progress"]["min_vram_gb"] += increase_amount
                        print(
                            f"[Roll] CUDA OOM Detected. Increasing Global VRAM Req to {roll_log['progress']['min_vram_gb']}GB"
                        )
            if (cluster_running_this):
                print(
                    f"[Roll] Log for {case_hash} exists but does not indicate it ran on {hostname}. Skipping log analysis."
                )
        if not policy_triggered:
            print(
                f"[Roll] No resource limitations triggered for {case_hash} on {hostname}. No update to global requirements."
            )
            return
        # else:
        #     print(f"[Roll] No log found for {case_hash} at {log_file}, skipping rolling restart.")
        #     return
        # ---------------- Step 4: 能力适格性计算 (Capability Check) ----------------

        target_ram = roll_log["progress"]["min_ram_gb"]
        target_vram = roll_log["progress"]["min_vram_gb"]
        print(f"[Roll] Global Target RAM: {target_ram}GB, VRAM: {target_vram}GB")
        # 计算本机需要多少 CPU/GPU 才能满足 target
        needed_cpus = math.ceil(target_ram / local_mem_per_cpu)
        needed_gpus = math.ceil(target_vram / local_mem_per_gpu)
        print(f"[Roll] Cluster {hostname} needs {needed_cpus} CPUs / {needed_gpus} GPUs to meet global requirements.")
        # assert False

        # ---------------- Step 5: 硬件上限检查 (Limit Check) ----------------
        limit_cpus = cluster_limits["max_num_cpus"]
        limit_gpus = cluster_limits["max_num_gpus"]
        # assert False
        if needed_cpus > limit_cpus or needed_gpus > limit_gpus:
            print(
                f"[Roll] Cluster {hostname} NOT QUALIFIED. Needs {needed_cpus}/{limit_cpus} CPUs / {needed_gpus}/{limit_gpus} GPUs. Exceeds limit."
            )
            # 标记自己为 Abandoned，以后不再尝试这个 case
            roll_log["clusters"][hostname]["abandoned"] = True
            with open(roll_log_file, 'w') as f:
                json.dump(roll_log, f, indent=4)
            return

        # ---------------- Step 6: 提交执行 ----------------
        # 只要能运行到这里，说明本机有能力（硬件满足需求）
        # 无论是否是新错误触发的，还是单纯的接力，只要 qualified 就可以尝试提交。
        # 为了防止重复提交，可以检查 attempts 或者 squeue (外部脚本处理并发)

        # 确保不低于 Job 原始定义的最小值（防止 job 定义了 4 卡，结果 progress 算是 0，被改成 0 卡）
        # needed_cpus = max(needed_cpus, job_config["num_cpus"])
        # needed_gpus = max(needed_gpus, job_config["num_gpus"])
        # 更新 Config
        updated_config = {}
        if needed_cpus != job_config.get("num_cpus"):
            updated_config["num_cpus"] = max(needed_cpus, job_config["num_cpus"])
        if needed_gpus != job_config.get("num_gpus"):
            updated_config["num_gpus"] = max(needed_gpus, job_config["num_gpus"])

        # 记录尝试
        roll_log["clusters"][hostname]["attempts"] += 1
        with open(roll_log_file, 'w') as f:
            json.dump(roll_log, f, indent=4)

        # 提交
        # 注意：这里我们假设外部有机制防止无限提交(比如 Cron 频率或 check running)
        print(f"[Roll] Cluster {hostname} is QUALIFIED. Submitting {case_hash}...")
        sample_entry(
            commit_hash,
            force=True,
            cases={case_hash: batch_config[case_hash] | updated_config},
        )
        run_entry(
            commit_hash,
            parellels=1,
            force=True,
            cases={case_hash: batch_config[case_hash] | updated_config},
            cluster=hostname,
        )
