"""FikoYT SAM2 Mask — point-based mask via ComfyUI-segment-anything-2."""
import os
import json

_SAM2_LOADER = None
_SAM2_RUN = None
try:
    from custom_nodes.ComfyUI_segment_anything_2 import nodes as _sam2_mod  # type: ignore
    _SAM2_LOADER = getattr(_sam2_mod, "DownloadAndLoadSAM2Model", None)
    _SAM2_RUN = getattr(_sam2_mod, "Sam2Segmentation", None)
except Exception:
    try:
        from ComfyUI_segment_anything_2 import nodes as _sam2_mod  # type: ignore
        _SAM2_LOADER = getattr(_sam2_mod, "DownloadAndLoadSAM2Model", None)
        _SAM2_RUN = getattr(_sam2_mod, "Sam2Segmentation", None)
    except Exception:
        pass

from ..core import ensure_model, get_models_dir


SAM2_PRESETS = {
    "sam2_hiera_large":  {"repo": "facebook/sam2-hiera-large",  "file": "sam2_hiera_large.pt",  "subdir": "sam2"},
    "sam2_hiera_base+":  {"repo": "facebook/sam2-hiera-base-plus", "file": "sam2_hiera_base_plus.pt", "subdir": "sam2"},
    "sam2_hiera_small":  {"repo": "facebook/sam2-hiera-small",  "file": "sam2_hiera_small.pt",  "subdir": "sam2"},
}


class FikoYTSAM2Mask:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "image": ("IMAGE",),
                "points_json": ("STRING", {
                    "multiline": True,
                    "default": "[[256, 256, 1]]",
                    "tooltip": "JSON: [[x, y, label], ...] — label 1=positive, 0=negative",
                }),
                "preset": (list(SAM2_PRESETS.keys()), {"default": "sam2_hiera_large"}),
                "auto_download": ("BOOLEAN", {"default": True}),
            }
        }

    RETURN_TYPES = ("MASK", "IMAGE")
    RETURN_NAMES = ("mask", "masked_image")
    FUNCTION = "segment"
    CATEGORY = "FikoYT/image"

    def segment(self, image, points_json, preset, auto_download):
        meta = SAM2_PRESETS[preset]
        sam_dir = get_models_dir(meta["subdir"])
        sam_path = os.path.join(sam_dir, meta["file"])

        if not os.path.isfile(sam_path) and auto_download:
            sam_path = ensure_model(meta["repo"], meta["file"], meta["subdir"])

        if _SAM2_LOADER is None or _SAM2_RUN is None:
            raise RuntimeError(
                "ComfyUI-segment-anything-2 node bulunamadı. "
                "https://github.com/kijai/ComfyUI-segment-anything-2 yükleyin."
            )

        try:
            points = json.loads(points_json)
            if not isinstance(points, list) or not points:
                raise ValueError("points_json boş veya geçersiz")
        except Exception as e:
            raise ValueError(f"points_json parse hatası: {e}")

        coords = [[p[0], p[1]] for p in points]
        labels = [int(p[2]) if len(p) > 2 else 1 for p in points]

        try:
            loader = _SAM2_LOADER()
            sam_model = loader.loadmodel(model=meta["file"], segmentor="single_image", device="cuda", precision="fp16")[0]

            runner = _SAM2_RUN()
            if hasattr(runner, "segment"):
                out = runner.segment(
                    sam2_model=sam_model, image=image,
                    coordinates_positive=json.dumps([c for c, l in zip(coords, labels) if l == 1]),
                    coordinates_negative=json.dumps([c for c, l in zip(coords, labels) if l == 0]),
                    individual_objects=False,
                    mask_threshold=0.0,
                )
            else:
                raise AttributeError("Sam2Segmentation segment API'si tanınmıyor.")
        except Exception as e:
            raise RuntimeError(
                f"SAM2 segmentasyonu başarısız: {e}\n"
                "ComfyUI-segment-anything-2 node'unu doğrudan kullanın."
            )

        if isinstance(out, tuple):
            mask = out[0]
            masked_img = out[1] if len(out) > 1 else image
        else:
            mask = out
            masked_img = image

        return (mask, masked_img)


NODE_CLASS_MAPPINGS = {"FikoYTSAM2Mask": FikoYTSAM2Mask}
NODE_DISPLAY_NAME_MAPPINGS = {"FikoYTSAM2Mask": "FikoYT • SAM2 Mask"}