from __future__ import annotations

from ...commands import CommandSpec, append_flag, conda_python
from ...config import settings
from ..dataset.tasks import _append_label_rebuild_flags


VISUAL_DIR = settings.source_root / "Tool-可视化"
VISUAL_LABEL_DIR = VISUAL_DIR / "0_图片Labels生成"


def build_visual_task(job_type: str, params: dict, conda_env: str) -> CommandSpec | None:
    if job_type == "visual.train":
        return CommandSpec(conda_python(conda_env, VISUAL_DIR / "train.py"), VISUAL_DIR, "run standalone YOLO visual training script")

    if job_type == "visual.inference":
        return CommandSpec(conda_python(conda_env, VISUAL_DIR / "inference.py"), VISUAL_DIR, "run standalone YOLO visual inference script")

    if job_type == "visual.fps":
        args = conda_python(conda_env, VISUAL_DIR / "get_FPS.py")
        append_flag(args, "--weights", params.get("weights", "yolov8n.pt"))
        append_flag(args, "--batch", params.get("batch", 1))
        imgs = params.get("imgs", [640, 640])
        if isinstance(imgs, str):
            imgs = [part for part in imgs.replace(",", " ").split() if part]
        if imgs:
            args.append("--imgs")
            args.extend(str(item) for item in imgs)
        append_flag(args, "--device", params.get("device", ""))
        append_flag(args, "--warmup", params.get("warmup"))
        append_flag(args, "--testtime", params.get("testtime"))
        if params.get("half"):
            args.append("--half")
        return CommandSpec(args, VISUAL_DIR, "measure YOLO visual model FPS")

    if job_type == "visual.yolo11_heatmap_v1":
        return CommandSpec(conda_python(conda_env, VISUAL_DIR / "yolov11_heatmap_V1.py"), VISUAL_DIR, "run standalone YOLOv11 heatmap V1")

    if job_type == "visual.yolo11_heatmap_v2":
        return CommandSpec(conda_python(conda_env, VISUAL_DIR / "yolov11_heatmap_V2.py"), VISUAL_DIR, "run standalone YOLOv11 heatmap V2")

    if job_type == "visual.label_ori":
        return CommandSpec(conda_python(conda_env, VISUAL_DIR / "Tool_Check_and_Gen_Txt_Label_ori_label.py"), VISUAL_DIR, "generate YOLO txt labels from original label ids")

    if job_type == "visual.label_sort":
        return CommandSpec(conda_python(conda_env, VISUAL_DIR / "Tool_Check_and_Gen_Txt_Label_sort_label.py"), VISUAL_DIR, "generate YOLO txt labels from sorted label ids")

    if job_type == "visual.gen_8bit_png":
        return CommandSpec(conda_python(conda_env, VISUAL_DIR / "Tool_Gen_8_Bit_PNG[没用，不认].py"), VISUAL_DIR, "generate 8-bit PNG labels with legacy visual tool")

    if job_type == "visual.deal_labels":
        args = conda_python(conda_env, VISUAL_LABEL_DIR / "4_deal_labels.py")
        _append_label_rebuild_flags(args, params)
        return CommandSpec(args, VISUAL_LABEL_DIR, "run visual label rebuild/conversion script")

    if job_type == "visual.tool_deal_labels_demo":
        return CommandSpec(conda_python(conda_env, VISUAL_LABEL_DIR / "Tool_deal_labels.py"), VISUAL_LABEL_DIR, "run visual label helper demo script")

    return None