from fastapi import FastAPI, File, UploadFile
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
import uvicorn
import numpy as np
import cv2
import base64
import tempfile
import os
from ultralytics import YOLO
from PIL import Image
import io

# ── Create the FastAPI app ───────────────────────────────────────
app = FastAPI(
    title       = "Polyp Detection API",
    description = "YOLOv8 polyp detection and segmentation",
    version     = "1.0.0"
)

# ── CORS Middleware ──────────────────────────────────────────────
# This allows your React app (running on a different port)
# to talk to this backend. Without this, browser blocks the request.
app.add_middleware(
    CORSMiddleware,
    allow_origins     = ["*"],  # in production, replace with your React URL
    allow_credentials = True,
    allow_methods     = ["*"],
    allow_headers     = ["*"],
)

# ── Load Model Once at Startup ───────────────────────────────────
# We load the model once when server starts
# Not on every request — that would be very slow
MODEL_PATH = "polyp_best_model.pt"

print(f"Loading model from {MODEL_PATH}...")
model = YOLO(MODEL_PATH)
print("✅ Model loaded successfully!")


# ── Helper: Convert numpy image to base64 string ────────────────
# React can't receive raw image data directly
# We convert it to a base64 string which React can display as <img>
def image_to_base64(img_rgb: np.ndarray) -> str:
    _, buffer = cv2.imencode('.jpg', cv2.cvtColor(img_rgb, cv2.COLOR_RGB2BGR))
    img_bytes = buffer.tobytes()
    img_b64   = base64.b64encode(img_bytes).decode('utf-8')
    return f"data:image/jpeg;base64,{img_b64}"


# ── Route 1: Health Check ────────────────────────────────────────
# React will ping this first to confirm backend is alive
@app.get("/")
def health_check():
    return {
        "status"  : "online",
        "message" : "Polyp Detection API is running"
    }


# ── Route 2: Detect Polyps ───────────────────────────────────────
# This is the main route React sends images to
@app.post("/detect")
async def detect_polyps(file: UploadFile = File(...)):
    """
    Receives an image file from React
    Returns:
    - annotated image (with masks drawn) as base64
    - polyp count
    - confidence scores
    - original image as base64
    """

    try:
        # ── Step 1: Read uploaded image ──────────────────────────
        contents  = await file.read()
        pil_image = Image.open(io.BytesIO(contents)).convert("RGB")
        img_array = np.array(pil_image)

        # ── Step 2: Save to temp file (the fix that worked before)
        img_bgr     = cv2.cvtColor(img_array, cv2.COLOR_RGB2BGR)
        img_resized = cv2.resize(img_bgr, (640, 640))

        tmp      = tempfile.NamedTemporaryFile(suffix='.jpg', delete=False)
        tmp_path = tmp.name
        tmp.close()
        cv2.imwrite(tmp_path, img_resized)

        # ── Step 3: Run YOLO model ────────────────────────────────
        results = model(
            tmp_path,
            conf    = 0.15,
            iou     = 0.4,
            verbose = False
        )
        os.remove(tmp_path)

        # ── Step 4: Draw results on image ─────────────────────────
        annotated     = results[0].plot(conf=True, masks=True, boxes=True)
        annotated_rgb = cv2.cvtColor(annotated, cv2.COLOR_BGR2RGB)

        # ── Step 5: Calculate frame-level cancer spread percentage ──
        frame_height, frame_width = img_resized.shape[:2]
        total_frame_pixels = frame_height * frame_width
        cancer_pixels = 0

        # Count pixels covered by segmentation masks
        if results[0].masks is not None:
            for mask in results[0].masks.data:
                # Convert mask to numpy and count positive pixels
                mask_np = mask.cpu().numpy()
                cancer_pixels += int(np.sum(mask_np))

        cancer_spread_percentage = (cancer_pixels / total_frame_pixels) * 100 if total_frame_pixels > 0 else 0

        # ── Step 6: Extract detection data ───────────────────────
        boxes      = results[0].boxes
        detections = []

        if boxes is not None and len(boxes) > 0:
            for idx, box in enumerate(boxes):
                conf  = float(box.conf[0].cpu().numpy())
                xyxy  = box.xyxy[0].cpu().numpy().tolist()

                detections.append({
                    "polyp_number" : idx + 1,
                    "confidence"   : round(conf * 100, 1),  # as percentage
                    "bbox"         : {
                        "x1" : round(xyxy[0]),
                        "y1" : round(xyxy[1]),
                        "x2" : round(xyxy[2]),
                        "y2" : round(xyxy[3])
                    }
                })

        # ── Step 7: Prepare original image for display ────────────
        orig_resized = cv2.resize(
            cv2.cvtColor(img_array, cv2.COLOR_RGB2BGR),
            (640, 640)
        )
        orig_rgb = cv2.cvtColor(orig_resized, cv2.COLOR_BGR2RGB)

        # ── Step 8: Return everything to React ───────────────────
        return JSONResponse(content={
            "success"                    : True,
            "polyp_count"                : len(detections),
            "cancer_spread_percentage"   : round(cancer_spread_percentage, 2),
            "total_cancer_pixels"        : cancer_pixels,
            "total_frame_pixels"         : total_frame_pixels,
            "detections"                 : detections,
            "annotated_image"            : image_to_base64(annotated_rgb),
            "original_image"             : image_to_base64(orig_rgb),
            "message"                    : f"Found {len(detections)} polyp(s)" if detections else "No polyps detected"
        })

    except Exception as e:
        return JSONResponse(
            status_code = 500,
            content     = {
                "success" : False,
                "error"   : str(e),
                "message" : "Something went wrong processing the image"
            }
        )


# ── Run the server ───────────────────────────────────────────────
if __name__ == "__main__":
    uvicorn.run(
        "main:app",
        host   = "0.0.0.0",
        port   = 8000,
        reload = True    # auto-restarts when you change code
    )