from __future__ import annotations

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


YOLO_DIR = settings.source_root / "Seg_All_In_One_YoloModel"
VIDEO_YOLO_DIR = settings.source_root / "Seg_Predict_YoloModel"
CUSTOM_TRAIN = settings.project_root / "backend" / "app" / "modules" / "yolo" / "custom_train.py"


def build_yolo_task(job_type: str, params: dict, conda_env: str) -> CommandSpec | None:
    env = {"SEG_CONDA_ENV": conda_env}

    if job_type == "yolo.train":
        args = conda_python(conda_env, YOLO_DIR / "yolo_train.py")
        append_flag(args, "--model", required(params, "model"))
        return CommandSpec(args, YOLO_DIR, "train one Ultralytics YOLO segmentation model")

    if job_type == "yolo.train_custom":
        args = conda_python(conda_env, CUSTOM_TRAIN)
        append_flag(args, "--data", required(params, "data"))
        append_flag(args, "--model", params.get("model", "YOLO11n-seg"))
        append_flag(args, "--epochs", params.get("epochs", 10))
        append_flag(args, "--imgsz", params.get("imgsz", 640))
        append_flag(args, "--batch", params.get("batch", 1))
        append_flag(args, "--workers", params.get("workers", 0))
        append_flag(args, "--device", params.get("device", "cpu"))
        append_flag(args, "--project", params.get("project", settings.project_root / "var" / "custom_yolo_runs"))
        append_flag(args, "--name", params.get("name", "custom_upload"))
        if params.get("exist_ok", True):
            args.append("--exist-ok")
        return CommandSpec(args, YOLO_DIR, "train YOLO segmentation on a supplied dataset.yaml")

    if job_type == "yolo.batch_train":
        return CommandSpec(bash(YOLO_DIR / "yolo_train.sh"), YOLO_DIR, "run legacy YOLO batch training", env=env)

    if job_type == "yolo.predict":
        args = conda_python(conda_env, YOLO_DIR / "yolo_predict_V2.py")
        append_flag(args, "--model", required(params, "model"))
        append_flag(args, "--source", params.get("source"))
        append_flag(args, "--pt_name", params.get("pt_name", "best.pt"))
        append_flag(args, "--conf", params.get("conf", 0.2))
        choice = str(params.get("run_choice", 1))
        return CommandSpec(args, YOLO_DIR, "predict with one YOLO model", stdin_text=f"{choice}\n")

    if job_type == "yolo.predict_v1":
        args = conda_python(conda_env, YOLO_DIR / "yolo_predict_V1.py")
        append_flag(args, "--model", required(params, "model"))
        append_flag(args, "--source", params.get("source"))
        append_flag(args, "--pt_name", params.get("pt_name", "best.pt"))
        append_flag(args, "--conf", params.get("conf", 0.2))
        choice = str(params.get("run_choice", 1))
        return CommandSpec(args, YOLO_DIR, "predict with legacy YOLO V1 script", stdin_text=f"{choice}\n")

    if job_type == "yolo.batch_predict":
        args = bash(YOLO_DIR / "yolo_predict.sh")
        append_flag(args, "--pt_name", params.get("pt_name", "best.pt"))
        append_flag(args, "--conf", params.get("conf", 0.2))
        append_flag(args, "--heatmap_method", params.get("heatmap_method"))
        return CommandSpec(args, YOLO_DIR, "run legacy YOLO batch prediction", env=env)

    if job_type == "yolo.heatmap":
        args = conda_python(conda_env, YOLO_DIR / "yolo_predict_visualize_nn.py")
        append_flag(args, "--model", required(params, "model"))
        append_flag(args, "--target_layers", params.get("target_layers", "default"))
        append_flag(args, "--cam_method", params.get("cam_method", "All"))
        append_flag(args, "--pt_name", params.get("pt_name", "best.pt"))
        choice = str(params.get("run_choice", 1))
        return CommandSpec(args, YOLO_DIR, "generate YOLO heatmaps", stdin_text=f"{choice}\n")

    if job_type == "yolo.compare":
        args = conda_python(conda_env, YOLO_DIR / "yolo_predict_V2_compare_all.py")
        append_flag(args, "--pt_name", params.get("pt_name", "all"))
        return CommandSpec(args, YOLO_DIR, "compare all YOLO prediction outputs")

    if job_type == "yolo.raw_mask_check":
        args = conda_python(conda_env, YOLO_DIR / "yolo_predict_raw_masks_check.py")
        append_flag(args, "--pt_name", params.get("pt_name", "best.pt"))
        return CommandSpec(args, YOLO_DIR, "check YOLO raw mask completeness")

    if job_type == "yolo.copy_best":
        args = bash(YOLO_DIR / "Tool_Yolo_Copy_Best_Model.sh")
        append_flag(args, "--pt_name", params.get("pt_name", "best.pt"))
        return CommandSpec(args, YOLO_DIR, "copy YOLO best weights into prediction area")

    if job_type == "yolo.video_visible":
        return CommandSpec(conda_python(conda_env, VIDEO_YOLO_DIR / "yolo_Seg_Video-V1-Visible.py"), VIDEO_YOLO_DIR, "render visible YOLO video prediction")

    if job_type == "yolo.video_unvisible":
        return CommandSpec(conda_python(conda_env, VIDEO_YOLO_DIR / "yolo_Seg_Video-V2-UnVisible.py"), VIDEO_YOLO_DIR, "render invisible/headless YOLO video prediction")

    if job_type == "yolo.layer_tester":
        return CommandSpec(conda_python(conda_env, YOLO_DIR / "Yolo可视化测试" / "yolo_layer_tester.py"), YOLO_DIR, "test YOLO heatmap target layers")

    return None