"""
Scenedrobe API – FastAPI server for Grounded SAM 2.

Models are loaded once at startup and kept in memory.
"""

from __future__ import annotations

import base64
import io
import os
import time
import uuid
from contextlib import asynccontextmanager
from typing import Optional

import numpy as np
import torch
from pathlib import Path

from fastapi import FastAPI, File, Form, UploadFile
from fastapi.responses import JSONResponse, Response
from fastapi.staticfiles import StaticFiles
from qdrant_client import QdrantClient
from qdrant_client.models import Distance, PointStruct, VectorParams

from detect import (
    detect,
    draw_results,
    embed_masked_crop,
    load_clip,
    load_dinov2,
    load_gdino,
    load_sam2,
    segment,
    verify_detections,
)

# ── shared state ────────────────────────────────────────────────────────
models = {}

QDRANT_URL = os.getenv("QDRANT_URL", "")  # empty = use local file storage
QDRANT_PATH = os.getenv("QDRANT_PATH", "./qdrant_data")
COLLECTION_NAME = "catalog"
EMBED_DIM = 384  # DINOv2-small output dimension

# ── detection cache (stores masks + image for re-embedding with erosion) ─
_detection_cache = {}


def get_device() -> str:
    if torch.backends.mps.is_available():
        return "mps"
    if torch.cuda.is_available():
        return "cuda"
    return "cpu"


def init_qdrant() -> QdrantClient:
    if QDRANT_URL:
        print(f"Connecting to remote Qdrant at {QDRANT_URL}")
        client = QdrantClient(url=QDRANT_URL)
    else:
        print(f"Using local Qdrant storage at {QDRANT_PATH}")
        client = QdrantClient(path=QDRANT_PATH)
    collections = [c.name for c in client.get_collections().collections]
    if COLLECTION_NAME not in collections:
        client.create_collection(
            collection_name=COLLECTION_NAME,
            vectors_config=VectorParams(size=EMBED_DIM, distance=Distance.COSINE),
        )
        print(f"Created Qdrant collection '{COLLECTION_NAME}'")
    else:
        print(f"Qdrant collection '{COLLECTION_NAME}' already exists")
    return client


@asynccontextmanager
async def lifespan(app: FastAPI):
    device = get_device()
    print(f"Device: {device}")

    print("Loading Grounding DINO...")
    gdino_processor, gdino_model = load_gdino(device)

    print("Loading SAM 2...")
    sam2_processor, sam2_model = load_sam2(device)

    print("Loading CLIP...")
    clip_processor, clip_model = load_clip(device)

    print("Loading DINOv2...")
    dinov2_processor, dinov2_model = load_dinov2(device)

    print("Connecting to Qdrant...")
    qdrant = init_qdrant()

    models["device"] = device
    models["gdino_processor"] = gdino_processor
    models["gdino_model"] = gdino_model
    models["sam2_processor"] = sam2_processor
    models["sam2_model"] = sam2_model
    models["clip_processor"] = clip_processor
    models["clip_model"] = clip_model
    models["dinov2_processor"] = dinov2_processor
    models["dinov2_model"] = dinov2_model
    models["qdrant"] = qdrant

    print("Models loaded. Server ready.\n")
    yield
    models.clear()


app = FastAPI(title="Scenedrobe API", lifespan=lifespan)


@app.post("/detect")
async def detect_endpoint(
    image: UploadFile = File(...),
    prompt: str = Form(...),
    box_threshold: float = Form(0.15),
    text_threshold: float = Form(0.15),
):
    """
    Upload an image + text prompt, get back the annotated image
    with segmentation masks.
    """
    from PIL import Image as PILImage

    start = time.time()

    # read uploaded image
    image_bytes = await image.read()
    pil_image = PILImage.open(io.BytesIO(image_bytes)).convert("RGB")

    # ensure prompt ends with period
    text = prompt.strip()
    if not text.endswith("."):
        text += "."

    device = models["device"]

    # step 1: detect with Grounding DINO (low threshold for recall)
    boxes, scores, labels = detect(
        pil_image, text,
        models["gdino_processor"], models["gdino_model"], device,
        box_threshold=box_threshold,
        text_threshold=text_threshold,
    )

    # step 2: verify with CLIP (filter false positives)
    if len(boxes) > 0:
        boxes, scores, labels = verify_detections(
            pil_image, boxes, scores, labels,
            models["clip_processor"], models["clip_model"], device,
        )

    if len(boxes) == 0:
        buf = io.BytesIO()
        pil_image.save(buf, format="JPEG")
        return Response(
            content=buf.getvalue(),
            media_type="image/jpeg",
            headers={"X-Detections": "0", "X-Time-Ms": str(int((time.time() - start) * 1000))},
        )

    # step 2: segment with SAM 2
    masks = segment(pil_image, boxes, models["sam2_processor"], models["sam2_model"], device)

    # draw masks on image
    annotated = draw_results(pil_image, masks, scores, labels)

    # encode result
    buf = io.BytesIO()
    annotated.save(buf, format="JPEG", quality=90)

    elapsed_ms = int((time.time() - start) * 1000)
    n = len(boxes)
    print(f"  Prompt: {text!r} | {n} detection(s) | {elapsed_ms}ms")

    return Response(
        content=buf.getvalue(),
        media_type="image/jpeg",
        headers={
            "X-Detections": str(n),
            "X-Time-Ms": str(elapsed_ms),
        },
    )


def mask_to_polygon(mask_np: np.ndarray, tolerance: int = 2) -> list[list[int]]:
    """Extract a simplified polygon boundary from a binary mask using basic edge tracing."""
    from PIL import Image as PILImage, ImageFilter

    # find contour by eroding the mask and XOR-ing with the original
    mask_img = PILImage.fromarray(mask_np.astype(np.uint8) * 255, "L")
    eroded = mask_img.filter(ImageFilter.MinFilter(3))
    edge = np.array(mask_img).astype(bool) & ~np.array(eroded).astype(bool)

    ys, xs = np.where(edge)
    if len(ys) == 0:
        return []

    # subsample points for a lighter polygon
    step = max(1, len(ys) // 200)
    points = list(zip(xs[::step].tolist(), ys[::step].tolist()))
    return points


def mask_to_bbox(mask_np: np.ndarray) -> list[int]:
    """Tight bounding box [x_min, y_min, x_max, y_max] from a binary mask."""
    ys, xs = np.where(mask_np)
    return [int(xs.min()), int(ys.min()), int(xs.max()), int(ys.max())]


def mask_to_base64_png(mask_np: np.ndarray, color: tuple = (255, 0, 0, 120)) -> str:
    """Encode a binary mask as a transparent PNG, returned as a base64 string."""
    from PIL import Image as PILImage

    rgba = np.zeros((*mask_np.shape, 4), dtype=np.uint8)
    rgba[mask_np] = color
    img = PILImage.fromarray(rgba, "RGBA")
    buf = io.BytesIO()
    img.save(buf, format="PNG")
    return base64.b64encode(buf.getvalue()).decode()


@app.post("/detect/json")
async def detect_json_endpoint(
    image: UploadFile = File(...),
    prompt: str = Form(...),
    box_threshold: float = Form(0.15),
    text_threshold: float = Form(0.15),
    include_mask_png: bool = Form(False),
):
    """
    Upload an image + text prompt, get back JSON with:
    - bbox: tight bounding box around the mask
    - polygon: contour points for SVG/canvas overlay
    - mask_png: (optional) base64-encoded transparent PNG of the mask
    """
    from PIL import Image as PILImage

    start = time.time()

    image_bytes = await image.read()
    pil_image = PILImage.open(io.BytesIO(image_bytes)).convert("RGB")
    width, height = pil_image.size

    text = prompt.strip()
    if not text.endswith("."):
        text += "."

    device = models["device"]

    # step 1: detect with Grounding DINO (low threshold for recall)
    boxes, scores, labels = detect(
        pil_image, text,
        models["gdino_processor"], models["gdino_model"], device,
        box_threshold=box_threshold,
        text_threshold=text_threshold,
    )

    # step 2: verify with CLIP (filter false positives)
    if len(boxes) > 0:
        boxes, scores, labels = verify_detections(
            pil_image, boxes, scores, labels,
            models["clip_processor"], models["clip_model"], device,
        )

    if len(boxes) == 0:
        return JSONResponse({
            "image_size": {"width": width, "height": height},
            "detections": [],
            "time_ms": int((time.time() - start) * 1000),
        })

    # step 3: segment
    masks = segment(pil_image, boxes, models["sam2_processor"], models["sam2_model"], device)

    COLORS = [
        (255, 0, 0, 180), (0, 200, 0, 180), (0, 100, 255, 180),
        (255, 255, 0, 180), (255, 0, 255, 180), (0, 255, 255, 180),
    ]

    MIN_SCORE = 0.3

    # generate a request ID for caching masks
    request_id = str(uuid.uuid4())
    cached_masks = []

    detections = []
    for i, (mask, score, label, box) in enumerate(zip(masks, scores, labels, boxes)):
        if score.item() < MIN_SCORE:
            continue
        mask_np = mask.numpy().astype(bool)
        mask_bbox = mask_to_bbox(mask_np)
        embedding, prepared_crop = embed_masked_crop(
            pil_image, mask_np, mask_bbox,
            models["dinov2_processor"], models["dinov2_model"], device,
        )

        # thumbnail of the prepared crop (what CLIP actually sees)
        crop_thumb = _make_thumbnail(prepared_crop)

        det_index = len(detections)
        det = {
            "label": label,
            "score": round(score.item(), 3),
            "dino_bbox": [round(v.item(), 1) for v in box],
            "mask_bbox": mask_bbox,
            "polygon": mask_to_polygon(mask_np),
            "embedding": embedding,
            "crop_thumbnail": crop_thumb,
            "detection_id": f"{request_id}:{det_index}",
        }
        if include_mask_png:
            det["mask_png"] = mask_to_base64_png(mask_np, COLORS[i % len(COLORS)])
        detections.append(det)
        cached_masks.append({"mask_np": mask_np, "label": label})

    # cache image + masks for re-embedding with erosion
    _detection_cache[request_id] = {
        "image": pil_image,
        "masks": cached_masks,
        "created": time.time(),
    }
    # evict old caches (keep last 5)
    if len(_detection_cache) > 5:
        oldest = sorted(_detection_cache, key=lambda k: _detection_cache[k]["created"])
        for k in oldest[:-5]:
            del _detection_cache[k]

    elapsed_ms = int((time.time() - start) * 1000)
    print(f"  [JSON] Prompt: {text!r} | {len(detections)} detection(s) | {elapsed_ms}ms")

    return JSONResponse({
        "image_size": {"width": width, "height": height},
        "detections": detections,
        "time_ms": elapsed_ms,
    })


@app.post("/detect/re-embed")
async def detect_re_embed(
    detection_id: str = Form(...),
    mask_erosion: int = Form(0),
):
    """Re-compute embedding for a cached detection with a new erosion value.
    Returns new embedding + crop thumbnail."""
    from PIL import Image as PILImage, ImageFilter

    parts = detection_id.split(":")
    if len(parts) != 2:
        return JSONResponse({"status": "error", "message": "Invalid detection_id"}, status_code=400)
    request_id, det_idx = parts[0], int(parts[1])

    cache = _detection_cache.get(request_id)
    if not cache:
        return JSONResponse({"status": "error", "message": "Detection expired — re-run detection"}, status_code=404)

    if det_idx >= len(cache["masks"]):
        return JSONResponse({"status": "error", "message": "Invalid detection index"}, status_code=400)

    pil_image = cache["image"]
    mask_data = cache["masks"][det_idx]
    mask_np = mask_data["mask_np"].copy()
    label = mask_data["label"]
    device = models["device"]

    # apply erosion
    if mask_erosion > 0:
        kernel_size = mask_erosion * 2 + 1
        mask_img = PILImage.fromarray(mask_np.astype(np.uint8) * 255, "L")
        mask_img = mask_img.filter(ImageFilter.MinFilter(kernel_size))
        mask_np = np.array(mask_img) > 127

    ys, xs = np.where(mask_np)
    if len(ys) == 0:
        return JSONResponse({"status": "error", "message": "Mask empty after erosion"}, status_code=400)

    mask_bbox = [int(xs.min()), int(ys.min()), int(xs.max()), int(ys.max())]

    # re-embed
    embedding, prepared_crop = embed_masked_crop(
        pil_image, mask_np, mask_bbox,
        models["dinov2_processor"], models["dinov2_model"], device,
    )

    crop_thumb = _make_thumbnail(prepared_crop)

    return JSONResponse({
        "status": "ok",
        "embedding": embedding,
        "crop_thumbnail": crop_thumb,
    })


def _make_thumbnail(pil_image, max_size=200, quality=80) -> str:
    """Resize and encode a PIL image as a base64 JPEG thumbnail."""
    thumb = pil_image.copy()
    thumb.thumbnail((max_size, max_size))
    buf = io.BytesIO()
    thumb.save(buf, format="JPEG", quality=quality)
    return base64.b64encode(buf.getvalue()).decode()


def _extract_garment(
    pil_image,
    label: str,
    box_threshold: float = 0.15,
    text_threshold: float = 0.15,
    clip_threshold: float = 0.45,
    mask_erosion: int = 0,
) -> dict:
    """Detect, segment, and embed a garment from an image. Returns extraction data or raises ValueError."""
    from PIL import Image as PILImage

    device = models["device"]
    prompt = label.strip()
    if not prompt.endswith("."):
        prompt += "."

    print(f"  [Extract] label={prompt!r} box_thresh={box_threshold} text_thresh={text_threshold} clip_thresh={clip_threshold}")

    boxes, scores, labels_out = detect(
        pil_image, prompt,
        models["gdino_processor"], models["gdino_model"], device,
        box_threshold=box_threshold, text_threshold=text_threshold,
    )
    print(f"  [Extract] DINO found {len(boxes)} detection(s)")

    if len(boxes) > 0:
        boxes, scores, labels_out = verify_detections(
            pil_image, boxes, scores, labels_out,
            models["clip_processor"], models["clip_model"], device,
            clip_threshold=clip_threshold,
        )

    if len(boxes) == 0:
        raise ValueError(f"Could not detect '{label}' in the image.")

    best = scores.argmax().item()
    masks = segment(pil_image, boxes[best:best+1], models["sam2_processor"], models["sam2_model"], device)
    mask = masks[0]
    mask_np = mask.numpy().astype(bool)

    # erode mask to trim bleeding edges
    if mask_erosion > 0:
        from PIL import ImageFilter
        from PIL import Image as _Img
        kernel_size = mask_erosion * 2 + 1
        mask_img = _Img.fromarray(mask_np.astype(np.uint8) * 255, "L")
        mask_img = mask_img.filter(ImageFilter.MinFilter(kernel_size))
        mask_np = np.array(mask_img) > 127

    ys, xs = np.where(mask_np)
    if len(ys) == 0:
        raise ValueError(f"Mask became empty after erosion — try a lower value.")
    mask_bbox = [int(xs.min()), int(ys.min()), int(xs.max()), int(ys.max())]

    embedding, prepared_crop = embed_masked_crop(
        pil_image, mask_np, mask_bbox,
        models["dinov2_processor"], models["dinov2_model"], device,
    )

    # extracted_thumbnail = white-background crop (for visual display)
    x1, y1, x2, y2 = mask_bbox
    crop = pil_image.crop((x1, y1, x2, y2))
    crop_mask = mask_np[y1:y2, x1:x2]
    crop_np = np.array(crop)
    crop_np[~crop_mask] = 255
    extracted = PILImage.fromarray(crop_np)

    return {
        "embedding": embedding,
        "original_thumbnail": _make_thumbnail(pil_image),
        "extracted_thumbnail": _make_thumbnail(extracted),
        "prepared_thumbnail": _make_thumbnail(prepared_crop),
        "score": round(scores[best].item(), 3),
    }


@app.post("/catalog/preview")
async def catalog_preview(
    image: UploadFile = File(...),
    label: str = Form(...),
    box_threshold: float = Form(0.15),
    text_threshold: float = Form(0.15),
    clip_threshold: float = Form(0.45),
    mask_erosion: int = Form(0),
):
    """
    Run detection + extraction without saving. Returns thumbnails and embedding
    so the admin can verify before committing.
    """
    from PIL import Image as PILImage

    image_bytes = await image.read()
    pil_image = PILImage.open(io.BytesIO(image_bytes)).convert("RGB")

    try:
        result = _extract_garment(
            pil_image, label,
            box_threshold=box_threshold,
            text_threshold=text_threshold,
            clip_threshold=clip_threshold,
            mask_erosion=mask_erosion,
        )
    except ValueError as e:
        return JSONResponse({"status": "error", "message": str(e)}, status_code=400)

    return JSONResponse({
        "status": "ok",
        "label": label,
        "detection_score": result["score"],
        "original_thumbnail": result["original_thumbnail"],
        "extracted_thumbnail": result["extracted_thumbnail"],
        "prepared_thumbnail": result["prepared_thumbnail"],
        "embedding": result["embedding"],
    })


@app.post("/catalog/add")
async def catalog_add(
    label: str = Form(...),
    product_url: str = Form(""),
    price: Optional[str] = Form(None),
    embedding: str = Form(...),
    original_thumbnail: str = Form(...),
    extracted_thumbnail: str = Form(...),
    prepared_thumbnail: str = Form(""),
):
    """
    Save a pre-extracted garment to the catalog.
    Accepts the embedding and thumbnails from a prior /catalog/preview call.
    """
    import json
    embedding_vec = json.loads(embedding) if isinstance(embedding, str) else embedding

    point_id = str(uuid.uuid4())
    models["qdrant"].upsert(
        collection_name=COLLECTION_NAME,
        points=[
            PointStruct(
                id=point_id,
                vector=embedding_vec,
                payload={
                    "label": label,
                    "product_url": product_url,
                    "price": price,
                    "original_thumbnail": original_thumbnail,
                    "extracted_thumbnail": extracted_thumbnail,
                    "prepared_thumbnail": prepared_thumbnail,
                },
            )
        ],
    )

    return JSONResponse({"id": point_id, "label": label, "status": "added"})


@app.post("/catalog/search")
async def catalog_search(
    embedding: str = Form(...),
    limit: int = Form(5),
    min_score: float = Form(0.30),
):
    """Search the catalog for items similar to the given CLIP embedding vector."""
    import json
    embedding_vec = json.loads(embedding)

    results = models["qdrant"].query_points(
        collection_name=COLLECTION_NAME,
        query=embedding_vec,
        limit=limit,
    )

    items = []
    for point in results.points:
        items.append({
            "id": point.id,
            "score": round(point.score, 3),
            "label": point.payload.get("label", ""),
            "product_url": point.payload.get("product_url", ""),
            "price": point.payload.get("price"),
            "original_thumbnail": point.payload.get("original_thumbnail", ""),
            "extracted_thumbnail": point.payload.get("extracted_thumbnail", ""),
            "prepared_thumbnail": point.payload.get("prepared_thumbnail", ""),
        })

    for point in results.points:
        print(f"  [Search] id={point.id} label={point.payload.get('label','')} score={point.score:.4f}")
    print(f"  [Search] {len(results.points)} result(s) above {min_score}, returned {len(items)}")

    return JSONResponse({"results": items})


@app.get("/catalog/count")
async def catalog_count():
    """Return the number of items in the catalog."""
    info = models["qdrant"].get_collection(collection_name=COLLECTION_NAME)
    return {"count": info.points_count}


@app.get("/catalog/list")
async def catalog_list(
    limit: int = 50,
    offset: Optional[str] = None,
):
    """List all catalog items with pagination."""
    result = models["qdrant"].scroll(
        collection_name=COLLECTION_NAME,
        limit=limit,
        offset=offset,
        with_vectors=False,
        with_payload=True,
    )

    points, next_offset = result
    items = []
    for point in points:
        items.append({
            "id": point.id,
            "label": point.payload.get("label", ""),
            "product_url": point.payload.get("product_url", ""),
            "price": point.payload.get("price"),
            "original_thumbnail": point.payload.get("original_thumbnail", ""),
            "extracted_thumbnail": point.payload.get("extracted_thumbnail", ""),
        })

    return JSONResponse({
        "items": items,
        "next_offset": next_offset,
    })


@app.delete("/catalog/{item_id}")
async def catalog_delete(item_id: str):
    """Delete a single item from the catalog."""
    models["qdrant"].delete(
        collection_name=COLLECTION_NAME,
        points_selector=[item_id],
    )
    return JSONResponse({"id": item_id, "status": "deleted"})


@app.get("/health")
async def health():
    return {"status": "ok", "models_loaded": bool(models)}


# serve frontend — must be last so API routes take priority
WEB_DIR = Path(__file__).resolve().parent / "web"
app.mount("/", StaticFiles(directory=WEB_DIR, html=True), name="web")