from __future__ import annotations

from typing import Any, Dict

import numpy as np
from fastapi import APIRouter
from PIL import Image

from ..deps import OUTPUT_DIR
from ..schemas import SamPromptSegmentRequest, SegmentRequest
from ..services.image_io import b64_to_pil_image
from ..services.predictors import get_seg_predictor
from model.Seg.seg_loader import expand_mask, mask_to_contour_xy, run_sam_prompt

router = APIRouter(tags=["segment"])


@router.post("/segment")
def segment(req: SegmentRequest) -> Dict[str, Any]:
    try:
        model_name = req.model_name or "sam"
        pil = b64_to_pil_image(req.image_b64).convert("RGB")
        rgb = np.array(pil, dtype=np.uint8)

        predictor = get_seg_predictor(model_name)
        label_map = predictor(rgb).astype(np.int32)

        out_dir = OUTPUT_DIR / "segment"
        out_dir.mkdir(parents=True, exist_ok=True)
        out_path = out_dir / f"{model_name}_label.png"
        Image.fromarray(np.clip(label_map, 0, 255).astype(np.uint8), mode="L").save(out_path)

        return {"success": True, "label_path": str(out_path)}
    except Exception as e:
        return {"success": False, "error": str(e)}


@router.post("/segment/sam_prompt")
def segment_sam_prompt(req: SamPromptSegmentRequest) -> Dict[str, Any]:
    """
    交互式 SAM：裁剪图 + 点/框提示，返回掩膜外轮廓点列（裁剪像素坐标）。
    """
    try:
        pil = b64_to_pil_image(req.image_b64).convert("RGB")
        rgb = np.array(pil, dtype=np.uint8)
        h, w = rgb.shape[0], rgb.shape[1]

        if req.overlay_b64:
            ov = b64_to_pil_image(req.overlay_b64)
            if ov.size != (w, h):
                return {
                    "success": False,
                    "error": f"overlay 尺寸 {ov.size} 与 image {w}x{h} 不一致",
                    "contour": [],
                }

        if len(req.point_coords) != len(req.point_labels):
            return {
                "success": False,
                "error": "point_coords 与 point_labels 长度不一致",
                "contour": [],
            }
        if len(req.point_coords) < 1:
            return {"success": False, "error": "至少需要一个提示点", "contour": []}

        pc = np.array(req.point_coords, dtype=np.float32)
        if pc.ndim != 2 or pc.shape[1] != 2:
            return {"success": False, "error": "point_coords 每项须为 [x,y]", "contour": []}

        pl = np.array(req.point_labels, dtype=np.int64)
        box = np.array(req.box_xyxy, dtype=np.float32)

        mask = run_sam_prompt(rgb, pc, pl, box_xyxy=box)
        if not np.any(mask):
            return {"success": False, "error": "SAM 未产生有效掩膜", "contour": []}

        mask = expand_mask(mask, int(req.expand_px))
        contour = mask_to_contour_xy(mask, epsilon_px=2.0)
        if len(contour) < 3:
            return {"success": False, "error": "轮廓点数不足", "contour": []}

        return {"success": True, "contour": contour, "error": None}
    except Exception as e:
        return {"success": False, "error": str(e), "contour": []}