from fastapi import FastAPI, File, UploadFile, HTTPException
import numpy as np
import tensorflow as tf
from PIL import Image
import io
from huggingface_hub import hf_hub_download
import logging

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

app = FastAPI()
model = None

def load_model():
    global model
    try:
        model_path = hf_hub_download(
            "ayushirathour/chest-xray-pneumonia-detection", 
            "best_chest_xray_model.h5"
        )
        model = tf.keras.models.load_model(model_path)
        logger.info("✅ Model loaded successfully!")
        return True
    except Exception as e:
        logger.error(f"❌ Model loading failed: {e}")
        return False

def preprocess_image(image_bytes):
    image = Image.open(io.BytesIO(image_bytes))
    image = image.convert("RGB")
    image = image.resize((224, 224))
    array = np.asarray(image, dtype=np.float32) / 255.0
    return np.expand_dims(array, axis=0)

@app.on_event("startup")
async def startup_event():
    load_model()

@app.get("/")
def home():
    return {
        "message": "🏥 AI Pneumonia Detection API", 
        "model_loaded": model is not None,
        "tensorflow_version": tf.__version__
    }

@app.post("/predict")
async def predict(file: UploadFile = File(...)):
    if not model:
        raise HTTPException(status_code=503, detail="Model not loaded")
    
    if not file.content_type or not file.content_type.startswith("image/"):
        raise HTTPException(status_code=400, detail="Invalid file type")
    
    try:
        contents = await file.read()
        input_tensor = preprocess_image(contents)
        prediction = model.predict(input_tensor, verbose=0)
        score = float(prediction[0][0])
        
        if score > 0.5:
            diagnosis = "PNEUMONIA"
            confidence = score * 100
        else:
            diagnosis = "NORMAL" 
            confidence = (1 - score) * 100
            
        return {
            "diagnosis": diagnosis,
            "confidence": round(confidence, 2),
            "raw_score": round(score, 4),
            "filename": file.filename
        }
        
    except Exception as e:
        logger.error(f"Prediction error: {e}")
        raise HTTPException(status_code=500, detail=f"Prediction failed: {str(e)}")

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)