import io
import json
import numpy as np
import torch
from fastapi import FastAPI, UploadFile, File, Form, HTTPException
from fastapi.responses import StreamingResponse
from PIL import Image, ImageDraw
from transformers import SamModel, SamProcessor

app = FastAPI(title="MedSAM Image Segmentation Service")

# Global device initialization (Fall back gracefully to CPU on standard HF spaces)
device = "cuda" if torch.cuda.is_available() else "cpu"

print(f"Loading MedSAM model onto: {device}...")
processor = SamProcessor.from_pretrained("flaviagiammarino/medsam-vit-base")
model = SamModel.from_pretrained("flaviagiammarino/medsam-vit-base").to(device)
model.eval()
print("Model loaded successfully.")

@app.get("/")
def health_check():
    return {"status": "healthy", "device": device}

@app.post("/segment")
async def segment_image(
    file: UploadFile = File(...),
    bbox: str = Form(..., description="JSON string array format: '[xmin, ymin, xmax, ymax]'")
):
    # 1. Parse bounding box coordinates
    try:
        box_coords = json.loads(bbox)
        if not isinstance(box_coords, list) or len(box_coords) != 4:
            raise ValueError
    except ValueError:
        raise HTTPException(status_code=400, detail="Bounding box format must be '[xmin, ymin, xmax, ymax]'")

    # 2. Read image
    try:
        image_bytes = await file.read()
        raw_image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
    except Exception:
        raise HTTPException(status_code=400, detail="Invalid image file format uploaded.")

    # 3. Process image and run inference
    try:
        inputs = processor(
            raw_image, 
            input_boxes=[[box_coords]], 
            return_tensors="pt"
        ).to(device)
        
        with torch.no_grad():
            outputs = model(**inputs)
        
        # Post-process low resolution logits into a concrete mask array
        masks = processor.image_processor.post_process_masks(
            outputs.pred_masks.cpu(), 
            inputs["original_sizes"].cpu(), 
            inputs["reshaped_input_sizes"].cpu()
        )
        
        # Binary mask array (True/False boolean configuration)
        mask_array = masks[0][0][0].numpy()
        
    except Exception as e:
        raise HTTPException(
            status_code=500, 
            detail=f"MedSAM core processing exception: {str(e)}"
        )

    # 4. Generate the composite visual output (matching reference target)
    # Create an overlay layout where True items highlight yellow at an alpha state
    yellow_overlay = Image.new(
        "RGBA", 
        raw_image.size, 
        (251, 252, 30, 150)
    ) 
    mask_image = Image.fromarray(
        (mask_array * 255).astype(np.uint8), 
        mode="L"
    )
    
    # Composite the isolated mask layout with the underlying source picture
    segmented_image = Image.composite(
        yellow_overlay, 
        raw_image.convert("RGBA"), 
        mask_image
    )
    
    # Draw the diagnostic bounding box outline border matching your input configuration
    draw = ImageDraw.Draw(segmented_image)
    draw.rectangle(
        box_coords, 
        outline="blue", 
        width=3
    )
    
    # Stream back final constructed image payload 
    buffer = io.BytesIO()
    segmented_image.convert("RGB").save(
        buffer, format="PNG"
    )
    buffer.seek(0)
    
    return StreamingResponse(
        buffer, 
        media_type="image/png"
    )