from fastapi import FastAPI, UploadFile, File, Query
import tifffile
import numpy as np
from cellpose import models
from fastapi.responses import FileResponse
from huggingface_hub import snapshot_download
import os
from io import BytesIO
import tempfile

app = FastAPI(
    title="Cellpose API",
    docs_url="/docs",
    redoc_url=None
)

@app.get("/")
def root():
    return {"status": "ok"}
    
MODEL_REPO = "Shalani08/cellpose_3.1.1_finetuned"

repo_dir = snapshot_download(
    repo_id=MODEL_REPO,
    repo_type="model",
    revision="8a020c0",
    token=os.environ.get("HF_TOKEN")
)

model_dir = f"{repo_dir}/ddq_model"

model = models.CellposeModel(
    pretrained_model=model_dir,
    gpu=False
)



@app.post("/segment")
async def segment(
    image: UploadFile = File(...),
    diameter: float = Query(0, description="Cell diameter (0 = auto)"),
    channels: str = Query("0,0", description="Cellpose channels"),
    flow_threshold: float = Query(0.4),
    cellprob_threshold: float = Query(0.0)
):
    img_bytes = await image.read()
    img = tifffile.imread(BytesIO(img_bytes))

    if img.ndim == 3 and img.shape[-1] == 3:
        img = img[..., 0]

    ch = [int(c) for c in channels.split(",")]

    masks, _, _ = model.eval(
        img,
        diameter=diameter,
        channels=ch,
        flow_threshold=flow_threshold,
        cellprob_threshold=cellprob_threshold
    )

    tmp = tempfile.NamedTemporaryFile(suffix=".tif", delete=False)
    tifffile.imwrite(tmp.name, masks.astype(np.uint16))

    return FileResponse(tmp.name, media_type="image/tiff")