from __future__ import annotations

from ...commands import CommandSpec, append_flag, conda_python, required
from ...config import settings


MMSEG_DIR = settings.source_root / "Seg_All_In_One_MMSeg"
MY_DIR = MMSEG_DIR / "My_All_In_One"


def _stdin_for_generate_alg(params: dict) -> str:
    lines = [
        str(params.get("dataset_choice", 1)),
        str(params.get("gpu_count", 1)),
    ]
    gpu_ids = params.get("gpu_ids", [0])
    if isinstance(gpu_ids, str):
        gpu_ids = [part.strip() for part in gpu_ids.split(",") if part.strip()]
    for index in range(int(params.get("gpu_count", len(gpu_ids) or 1))):
        lines.append(str(gpu_ids[index] if index < len(gpu_ids) else 0))

    mode = str(params.get("schedule_mode", 2))
    lines.append(mode)
    if mode == "1":
        lines.extend(
            [
                str(params.get("train_k", 40)),
                str(params.get("check_count", 10)),
                str(params.get("logger_interval", 50)),
            ]
        )
    else:
        lines.extend(
            [
                str(params.get("max_epochs", 300)),
                str(params.get("val_interval", 1)),
                str(params.get("checkpoint_interval", 10)),
                str(params.get("logger_interval", "")),
            ]
        )
    lines.append(str(params.get("algorithm_choice", 1)))
    return "\n".join(lines) + "\n"


def build_mmseg_task(job_type: str, params: dict, conda_env: str) -> CommandSpec | None:
    if job_type == "mmseg.init_weights":
        return CommandSpec(conda_python(conda_env, MY_DIR / "0_Initial_Save_All_Model_locally.py"), MMSEG_DIR, "download/save MMSeg pretrained weights locally")

    if job_type == "mmseg.generate_data":
        return CommandSpec(conda_python(conda_env, MY_DIR / "1_Initial_Data_All_data_from_1_Data_Parameter-V2.py"), MMSEG_DIR, "generate MMSeg dataset configs from JSON parameters")

    if job_type == "mmseg.generate_alg":
        script = MY_DIR / "2_Initial_Alg_All_data_from_2_Alg_Program-V2.py"
        return CommandSpec(
            conda_python(conda_env, script),
            MMSEG_DIR,
            "generate MMSeg algorithm config and training command",
            stdin_text=_stdin_for_generate_alg(params),
        )

    if job_type == "mmseg.train":
        config_path = required(params, "config")
        args = conda_python(conda_env, MMSEG_DIR / "tools" / "train.py", config_path)
        append_flag(args, "--work-dir", params.get("work_dir"))
        return CommandSpec(args, MMSEG_DIR, "train MMSeg model")

    if job_type == "mmseg.metrics":
        args = conda_python(conda_env, MY_DIR / "4_2_predict_matrics_from_log_V2.py")
        append_flag(args, "--input_dir", params.get("input_dir", "../Hardisk"))
        append_flag(args, "--output_dir", params.get("output_dir", "../BestMode_Predict_Results_DataSet_Public"))
        stdin = f"{params.get('dataset_choice', 1)}\n{params.get('algorithm_choice', 0)}\n"
        return CommandSpec(args, MMSEG_DIR, "extract best MMSeg metrics from logs", stdin_text=stdin)

    if job_type == "mmseg.flops_fps":
        args = conda_python(conda_env, MY_DIR / "4_1_predict_params_FLOPs_FPS_V2.py")
        append_flag(args, "--input_dir", params.get("input_dir", "../Hardisk"))
        append_flag(args, "--output_dir", params.get("output_dir", "../BestMode_Predict_Results_DataSet_Public"))
        append_flag(args, "--repeat-times", params.get("repeat_times", 3))
        stdin = f"{params.get('dataset_choice', 1)}\n{params.get('algorithm_choice', 0)}\n"
        if "shape_h" in params and "shape_w" in params:
            stdin += f"{params['shape_h']}\n{params['shape_w']}\n"
        return CommandSpec(args, MMSEG_DIR, "calculate MMSeg FLOPs/Params/FPS", stdin_text=stdin)

    if job_type == "mmseg.draw":
        return CommandSpec(conda_python(conda_env, MY_DIR / "4_3_predict_draw_pictures_and_tabels.py"), MMSEG_DIR, "generate MMSeg prediction pictures and tables")

    if job_type == "mmseg.extract_loss_miou":
        return CommandSpec(conda_python(conda_env, MY_DIR / "4_4_extract_loss_and_best_miou.py"), MMSEG_DIR, "extract MMSeg loss and best mIoU curves")

    if job_type == "mmseg.delete_epoch":
        return CommandSpec(conda_python(conda_env, MY_DIR / "3_Find_And_Delete_Special_Epoch.py"), MMSEG_DIR, "find and delete selected epoch checkpoints")

    return None