import os
import io
import numpy as np
from PIL import Image
import tensorflow as tf
from tensorflow import keras
from fastapi import FastAPI, UploadFile, File, HTTPException
from fastapi.middleware.cors import CORSMiddleware

# Initialize FastAPI app
app = FastAPI(
    title="PneuVision AI API",
    description="Backend API for MobileNetV2 Pneumonia Detection model",
    version="1.0.0"
)

# Enable CORS so frontend deployed on Cloudflare can access this API
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

MODEL_PATH = "MobileNetV2_final.keras"
model = None

@app.on_event("startup")
def load_model():
    global model
    if os.path.exists(MODEL_PATH):
        try:
            # Load the Keras model. Keras 3 format handles both weights and architecture.
            model = keras.models.load_model(MODEL_PATH)
            print("[BACKEND] Model loaded successfully.")
        except Exception as e:
            print(f"[BACKEND] Error loading model: {e}")
    else:
        print(f"[BACKEND] Model file not found at {MODEL_PATH}")

@app.get("/")
def health_check():
    """Health check endpoint to see if API and model are active."""
    return {
        "status": "ready" if model is not None else "model_not_loaded",
        "model_architecture": "MobileNetV2",
        "input_shape": [224, 224, 3],
        "classes": ["NORMAL", "PNEUMONIA"]
    }

@app.post("/predict")
async def predict(file: UploadFile = File(...)):
    """Receives a chest X-ray image and predicts the probability of pneumonia."""
    global model
    if model is None:
        raise HTTPException(
            status_code=503, 
            detail="Model is not initialized on the server. Please check server logs."
        )
    
    if not file.content_type.startswith("image/"):
        raise HTTPException(
            status_code=400, 
            detail="Invalid file format. Please upload an image (JPEG/PNG)."
        )
    
    try:
        # Read file contents and open image
        contents = await file.read()
        image = Image.open(io.BytesIO(contents)).convert("RGB")
        
        # Preprocessing:
        # 1. Resize image to 224x224 using Bilinear interpolation (matching Colab training setup)
        image = image.resize((224, 224), Image.Resampling.BILINEAR)
        img_array = np.array(image, dtype=np.float32)
        
        # 2. Add batch dimension: shape (1, 224, 224, 3)
        img_array = np.expand_dims(img_array, axis=0)
        
        # Note: Rescaling (1.0/255.0) and Normalization (ImageNet mean/std) 
        # are defined inside the Keras model as layers, so we don't apply them here.
        # Running model inference (returns shape [1, 2])
        predictions = model.predict(img_array)
        
        # Class 0: NORMAL, Class 1: PNEUMONIA
        prob_normal = float(predictions[0][0])
        prob_pneumonia = float(predictions[0][1])
        
        predicted_class = "PNEUMONIA" if prob_pneumonia > prob_normal else "NORMAL"
        confidence = max(prob_normal, prob_pneumonia)
        
        return {
            "prediction": predicted_class,
            "confidence": confidence,
            "probabilities": {
                "NORMAL": prob_normal,
                "PNEUMONIA": prob_pneumonia
            },
            "info": {
                "model": "MobileNetV2",
                "processed_shape": [224, 224, 3]
            }
        }
    except Exception as e:
        raise HTTPException(
            status_code=500, 
            detail=f"Inference error: {str(e)}"
        )