from fastapi import FastAPI, File, UploadFile
from fastapi.responses import JSONResponse
from transformers import pipeline
from PIL import Image
import io
from contextlib import asynccontextmanager

@asynccontextmanager
async def lifespan(app: FastAPI):
    global model
    print("🚀 Memuat model ResNet50 dari Hugging Face...")
    
    # Muat model tanpa argumen tambahan
    model = pipeline(
        "image-classification",
        model="SanketJadhav/PlantDiseaseClassifier-Resnet50"
    )
    
    print("✅ Model siap digunakan (CPU mode)")
    yield
    print("🧹 Server FastAPI dimatikan.")

app = FastAPI(lifespan=lifespan)

@app.post("/predict")
async def predict(file: UploadFile = File(...)):
    try:
        # Baca file gambar dari request
        image_bytes = await file.read()
        image = Image.open(io.BytesIO(image_bytes)).convert("RGB")

        # Jalankan prediksi
        results = model(image)

        # Ambil 3 hasil teratas
        top_results = sorted(results, key=lambda x: x['score'], reverse=True)[:3]
        formatted = [
            {"label": res['label'], "score": round(res['score'], 3)}
            for res in top_results
        ]

        return JSONResponse({
            "status": "success",
            "predictions": formatted
        })

    except Exception as e:
        return JSONResponse({
            "status": "error",
            "message": str(e)
        }, status_code=500)

@app.get("/")
async def root():
    return {"message": "🌱 Plant Disease API is running!"}

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