"""
SAM 3 local inference server for annotation tools.

Exposes a small FastAPI service that takes an image + bbox and returns a
segmentation mask. Designed for the classic annotation workflow:
    user draws a box around ONE object -> server returns the mask for that object.

Two endpoints are provided:

  POST /segment/box
      Uses Sam3Tracker (the SAM2-style Promptable Visual Segmentation head).
      Given ONE bbox, returns ONE mask for the object inside that box.
      This is what you want for 99% of annotation-tool use cases.

  POST /segment/exemplar
      Uses the full SAM 3 image predictor with a box as a visual "exemplar".
      Given one bbox, returns masks for ALL similar objects in the image.
      Useful when you want to label one instance and auto-propagate.

Run:
    python sam3_server.py --host 0.0.0.0 --port 8000

Request format (JSON):
    {
        "image":  "<base64-encoded image>"   # OR
        "image_url": "http://..."            # OR
        "image_path": "/absolute/path.jpg",
        "bbox": [x1, y1, x2, y2],            # absolute pixel coords (xyxy)
        "multimask_output": false             # optional, /segment/box only
    }

Response:
    {
        "mask":  "<base64 PNG, single-channel 0/255>",
        "rle":   {"size":[H,W], "counts":"..."},   # COCO RLE, handy for COCO-format tools
        "polygon": [[x,y], [x,y], ...],             # largest contour, optional
        "score": 0.97,
        "bbox_out": [x1,y1,x2,y2]
    }
"""

from __future__ import annotations

import argparse
import base64
import io
import logging
from typing import List, Optional

import numpy as np
import torch
import uvicorn
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from PIL import Image
from pydantic import BaseModel, Field

# ----- SAM 3 imports ---------------------------------------------------------
# The SAM 3 image predictor (concept / exemplar segmentation):
from sam3.model_builder import build_sam3_image_model
from sam3.model.sam3_image_processor import Sam3Processor

# The SAM 3 tracker (SAM2-style single-instance PVS on images):
# This is exposed in the HF `transformers` integration as Sam3TrackerModel,
# and in the reference repo via the image predictor for single-box prompts.
# We use HuggingFace transformers for the tracker because its API is clean
# and documented for bbox -> single mask. Fall back to the repo's predictor
# if transformers isn't available.
try:
    from transformers import Sam3TrackerModel, Sam3TrackerProcessor  # type: ignore
    HAS_TRACKER = True
except Exception:
    HAS_TRACKER = False

log = logging.getLogger("sam3-server")
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")


# ============================================================================
# Request / response schemas
# ============================================================================

class SegmentRequest(BaseModel):
    # Exactly one of these three must be provided.
    image: Optional[str] = Field(None, description="Base64-encoded image bytes.")
    image_url: Optional[str] = None
    image_path: Optional[str] = None

    bbox: List[float] = Field(..., description="Box in absolute pixels, [x1,y1,x2,y2].")
    multimask_output: bool = Field(
        False,
        description="If true (tracker only), return the 3 candidate masks instead of picking the highest-scoring one.",
    )
    return_polygon: bool = Field(True, description="Also return the largest contour as [[x,y], ...].")
    return_rle: bool = Field(True, description="Also return COCO RLE encoding of the mask.")


class SegmentResponse(BaseModel):
    mask: str  # base64 PNG
    score: float
    bbox_out: List[float]
    rle: Optional[dict] = None
    polygon: Optional[List[List[float]]] = None
    extra_masks: Optional[List[str]] = None  # only populated when multimask_output=True


# ============================================================================
# Model holder (loaded once at startup)
# ============================================================================

class SAM3Models:
    def __init__(self, device: str, checkpoint_path: Optional[str] = None):
        self.device = device

        log.info("Loading SAM 3 image model (concept/exemplar predictor)...")
        if checkpoint_path:
            self.image_model = build_sam3_image_model(checkpoint_path=checkpoint_path)
        else:
            self.image_model = build_sam3_image_model()
        self.image_model.to(device).eval()
        self.image_processor = Sam3Processor(self.image_model)

        if HAS_TRACKER:
            log.info("Loading SAM 3 tracker (single-instance PVS) from HuggingFace...")
            # facebook/sam3 exposes the tracker weights under this name
            self.tracker_processor = Sam3TrackerProcessor.from_pretrained("facebook/sam3")
            self.tracker_model = Sam3TrackerModel.from_pretrained("facebook/sam3").to(device).eval()
        else:
            log.warning(
                "transformers not available or missing Sam3Tracker classes. "
                "The /segment/box endpoint will fall back to the image predictor "
                "(still works, but treats the box as an exemplar)."
            )
            self.tracker_processor = None
            self.tracker_model = None


MODELS: Optional[SAM3Models] = None


# ============================================================================
# Helpers
# ============================================================================

def _load_image(req: SegmentRequest) -> Image.Image:
    if req.image:
        raw = base64.b64decode(req.image)
        return Image.open(io.BytesIO(raw)).convert("RGB")
    if req.image_path:
        return Image.open(req.image_path).convert("RGB")
    if req.image_url:
        import requests
        r = requests.get(req.image_url, timeout=20)
        r.raise_for_status()
        return Image.open(io.BytesIO(r.content)).convert("RGB")
    raise HTTPException(400, "Provide one of `image` (base64), `image_path`, or `image_url`.")


def _mask_to_png_b64(mask: np.ndarray) -> str:
    """mask: HxW bool/uint8 -> base64-encoded PNG string."""
    if mask.dtype != np.uint8:
        mask = (mask.astype(np.uint8) * 255)
    img = Image.fromarray(mask, mode="L")
    buf = io.BytesIO()
    img.save(buf, format="PNG", optimize=True)
    return base64.b64encode(buf.getvalue()).decode("ascii")


def _mask_to_rle(mask: np.ndarray) -> dict:
    """COCO-style RLE. Uses pycocotools if available, else a pure-python fallback."""
    mask = (mask > 0).astype(np.uint8)
    try:
        from pycocotools import mask as mutil  # type: ignore
        rle = mutil.encode(np.asfortranarray(mask))
        rle["counts"] = rle["counts"].decode("ascii")
        return rle
    except Exception:
        # Pure python fallback (uncompressed)
        h, w = mask.shape
        flat = mask.T.flatten()  # column-major, like COCO
        counts: List[int] = []
        last = 0
        run = 0
        for v in flat:
            if v == last:
                run += 1
            else:
                counts.append(run)
                run = 1
                last = v
        counts.append(run)
        return {"size": [h, w], "counts": counts}


def _mask_to_polygon(mask: np.ndarray) -> Optional[List[List[float]]]:
    """Largest-contour polygon, approximated. Returns None if mask is empty."""
    try:
        import cv2  # type: ignore
    except Exception:
        return None
    m = (mask > 0).astype(np.uint8) * 255
    contours, _ = cv2.findContours(m, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_TC89_L1)
    if not contours:
        return None
    c = max(contours, key=cv2.contourArea)
    # Simplify a bit — saves bytes and is friendlier for annotation tools
    epsilon = 0.0015 * cv2.arcLength(c, True)
    c = cv2.approxPolyDP(c, epsilon, True)
    return [[float(pt[0][0]), float(pt[0][1])] for pt in c]


def _mask_bbox(mask: np.ndarray) -> List[float]:
    ys, xs = np.where(mask > 0)
    if len(xs) == 0:
        return [0.0, 0.0, 0.0, 0.0]
    return [float(xs.min()), float(ys.min()), float(xs.max()), float(ys.max())]


def _build_response(mask: np.ndarray, score: float, req: SegmentRequest,
                    extra_masks: Optional[List[np.ndarray]] = None) -> SegmentResponse:
    return SegmentResponse(
        mask=_mask_to_png_b64(mask),
        score=float(score),
        bbox_out=_mask_bbox(mask),
        rle=_mask_to_rle(mask) if req.return_rle else None,
        polygon=_mask_to_polygon(mask) if req.return_polygon else None,
        extra_masks=[_mask_to_png_b64(m) for m in extra_masks] if extra_masks else None,
    )


# ============================================================================
# FastAPI app
# ============================================================================

app = FastAPI(title="SAM 3 local server", version="1.0")

# Allow local annotation tool UIs to call us from the browser
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)


@app.get("/health")
def health():
    return {
        "ok": True,
        "device": MODELS.device if MODELS else "unloaded",
        "tracker_available": HAS_TRACKER,
    }


@app.post("/segment/box", response_model=SegmentResponse)
@torch.inference_mode()
def segment_box(req: SegmentRequest):
    """
    Annotation-tool workflow:
        user draws a bbox around ONE object -> return ONE mask for that object.

    Uses the SAM3 tracker (SAM2-style PVS) for a clean single-instance result.
    """
    if MODELS is None:
        raise HTTPException(500, "Models not loaded.")
    image = _load_image(req)
    x1, y1, x2, y2 = req.bbox

    # ---- Preferred path: HuggingFace Sam3Tracker -------------------------
    if HAS_TRACKER and MODELS.tracker_model is not None:
        inputs = MODELS.tracker_processor(
            images=image,
            input_boxes=[[[x1, y1, x2, y2]]],  # batch[1] x obj[1] x 4
            return_tensors="pt",
        ).to(MODELS.device)

        outputs = MODELS.tracker_model(**inputs, multimask_output=req.multimask_output)

        # Post-process to image resolution, binary masks
        masks = MODELS.tracker_processor.post_process_masks(
            outputs.pred_masks.cpu(),
            original_sizes=inputs["original_sizes"].cpu(),
            reshaped_input_sizes=inputs["reshaped_input_sizes"].cpu(),
        )[0]  # shape: (num_masks, H, W) tensor
        scores = outputs.iou_scores[0].cpu().numpy().ravel()

        masks_np = masks.numpy().astype(bool)
        if req.multimask_output and masks_np.shape[0] > 1:
            best = int(np.argmax(scores))
            extras = [m for i, m in enumerate(masks_np) if i != best]
            return _build_response(masks_np[best], scores[best], req, extra_masks=extras)

        best = int(np.argmax(scores)) if masks_np.shape[0] > 1 else 0
        return _build_response(masks_np[best], float(scores[best]), req)

    # ---- Fallback: use image predictor with box as exemplar, pick the
    #      detection whose predicted box has highest IoU with the user bbox.
    state = MODELS.image_processor.set_image(image)
    out = MODELS.image_processor.set_box_prompt(
        state=state,
        boxes=[[x1, y1, x2 - x1, y2 - y1]],  # xywh as documented
        box_labels=[1],
    )
    masks = out["masks"]
    boxes = out["boxes"]
    scores = out["scores"]
    if len(masks) == 0:
        raise HTTPException(404, "No segmentation produced for the given bbox.")

    # Pick mask whose predicted bbox most overlaps user bbox
    def _iou(a, b):
        ax1, ay1, ax2, ay2 = a
        bx1, by1, bx2, by2 = b
        ix1, iy1 = max(ax1, bx1), max(ay1, by1)
        ix2, iy2 = min(ax2, bx2), min(ay2, by2)
        iw, ih = max(0, ix2 - ix1), max(0, iy2 - iy1)
        inter = iw * ih
        ua = (ax2 - ax1) * (ay2 - ay1) + (bx2 - bx1) * (by2 - by1) - inter
        return inter / ua if ua > 0 else 0.0

    user_box = [x1, y1, x2, y2]
    ious = [_iou(user_box, list(map(float, b))) for b in boxes]
    idx = int(np.argmax(ious))
    mask_np = np.asarray(masks[idx]).astype(bool)
    return _build_response(mask_np, float(scores[idx]), req)


@app.post("/segment/exemplar", response_model=SegmentResponse)
@torch.inference_mode()
def segment_exemplar(req: SegmentRequest):
    """
    Auto-propagation workflow:
        user draws a box around ONE instance -> return a merged mask of ALL
        similar instances found in the image. `extra_masks` contains each
        individual instance separately.
    """
    if MODELS is None:
        raise HTTPException(500, "Models not loaded.")
    image = _load_image(req)
    x1, y1, x2, y2 = req.bbox

    state = MODELS.image_processor.set_image(image)
    out = MODELS.image_processor.set_box_prompt(
        state=state,
        boxes=[[x1, y1, x2 - x1, y2 - y1]],
        box_labels=[1],
    )
    masks = out["masks"]
    scores = out["scores"]
    if len(masks) == 0:
        raise HTTPException(404, "No instances found.")

    masks_np = [np.asarray(m).astype(bool) for m in masks]
    merged = np.zeros_like(masks_np[0], dtype=bool)
    for m in masks_np:
        merged |= m
    best_score = float(np.max(scores))
    return _build_response(merged, best_score, req, extra_masks=masks_np)


# ============================================================================
# Entrypoint
# ============================================================================

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--host", default="0.0.0.0")
    parser.add_argument("--port", type=int, default=8000)
    parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
    parser.add_argument("--checkpoint", default=None,
                        help="Optional path to sam3 checkpoint (otherwise uses HF default).")
    args = parser.parse_args()

    global MODELS
    MODELS = SAM3Models(device=args.device, checkpoint_path=args.checkpoint)
    log.info("SAM 3 server ready on %s:%d (device=%s)", args.host, args.port, args.device)
    uvicorn.run(app, host=args.host, port=args.port, log_level="info")


if __name__ == "__main__":
    main()