from fastapi import FastAPI, UploadFile, File, HTTPException
from fastapi.responses import JSONResponse
from tensorflow.keras.models import load_model
from tensorflow.keras.preprocessing import image
import numpy as np
from PIL import Image
import io

app = FastAPI(title="OralScan Model API")

# Load the model when the app starts
model = None

@app.on_event("startup")
async def load_model_on_startup():
    global model
    try:
        model = load_model("model.keras")
        print("✅ MobileNetV2 model loaded successfully!")
    except Exception as e:
        print(f"❌ Error loading model: {e}")

@app.get("/")
def home():
    return {"message": "OralScan Model API is running! Upload image to /predict"}

@app.post("/predict")
async def predict(file: UploadFile = File(...)):
    if model is None:
        raise HTTPException(status_code=500, detail="Model not loaded yet")

    try:
        # Read uploaded image
        contents = await file.read()
        img = Image.open(io.BytesIO(contents)).convert("RGB")
        
        # Resize to 224x224
        img = img.resize((224, 224))
        
        # Preprocess
        img_array = image.img_to_array(img)
        img_array = np.expand_dims(img_array, axis=0)
        img_array = img_array / 255.0

        # Predict
        predictions = model.predict(img_array)
        predicted_class = int(np.argmax(predictions[0]))
        confidence = float(np.max(predictions[0]) * 100)

        # Your actual class names
        class_names = [
            "Oral Homogenous Leukoplakia",
            "Oral Non-Homogenous Leukoplakia",
            "Other Oral White Lesions"
        ]

        return {
            "predicted_class": predicted_class,
            "class_name": class_names[predicted_class],
            "confidence": round(confidence, 2),
            "message": "Prediction successful"
        }

    except Exception as e:
        raise HTTPException(status_code=400, detail=f"Error processing image: {str(e)}")