from fastapi import FastAPI, File, UploadFile, HTTPException
from tensorflow.keras.models import load_model
import numpy as np
from PIL import Image
import io
from fastapi.middleware.cors import CORSMiddleware
from llm_client import LLMClient
import json
# Initialize FastAPI app
app = FastAPI(title="Image Classification API")
import logging

# Add CORS middleware
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)



# Load the Keras model once at startup
try:
    model = load_model('IAPLD.h5')
except Exception as e:
    raise RuntimeError(f"Failed to load model 'IAPLD.h5': {str(e)}")

# Define class names (adjust if model outputs 3 classes instead of 4)
CLASS_NAMES = ['Potato___healthy', 'Potato___Early_blight','Potato___Late_blight']

# Function to preprocess the uploaded image
def preprocess_image(image: Image.Image) -> np.ndarray:
    
    # Resize to match model input shape (250, 250 as per your code)
    image = image.resize((250, 250))  # Adjust to (256, 256) if model expects that
    
    # Convert to NumPy array and normalize to 0-1 range
    image_array = np.array(image) / 255.0
    
    # Add batch dimension (1, 250, 250, 3)
    image_array = np.expand_dims(image_array, axis=0)
    
    return image_array

# Root endpoint
@app.get("/")
async def root():
    return {"message": "Welcome to the Image Classification API. Use POST /predict/ to upload an image."}
from recomm import Redommend

@app.get("/recommendation")
async def recommendation(disease: str):
    if not disease:
        raise HTTPException(status_code=400, detail="Disease parameter is required")

    try:
        llm_client = LLMClient()
        recommender = Redommend(llm_client)

        raw_response = recommender._run(disease).strip()

        if raw_response.startswith("```json"):
            raw_response = raw_response.replace("```json", "").replace("```", "").strip()

        data = json.loads(raw_response)
        return data

    except json.JSONDecodeError:
        raise HTTPException(
            status_code=500,
            detail=f"Le LLM n’a pas renvoyé un JSON valide : {raw_response}"

        )
    except Exception as e:
        logging.error(f"Error in recommendation endpoint: {str(e)}")
        raise HTTPException(status_code=500, detail=f"Error processing request: {str(e)}")

# Prediction endpoint
@app.post("/predict/")
async def predict(file: UploadFile = File(...)):
    if not file.content_type.startswith('image/'):
        raise HTTPException(status_code=400, detail="Uploaded file must be an image")

    try:
        # Read the image bytes
        contents = await file.read()
        
        # Open as PIL image
        image = Image.open(io.BytesIO(contents))
        print("Image size:", image.size)  # Debug: Check image size
        # Preprocess the image
        image_array = preprocess_image(image)
        print("Image shape:", image_array.shape)  # Debug: Check input shape
        
        # Make prediction (model outputs probabilities directly)
        predictions = model.predict(image_array)
        print("Probabilities:", predictions)  # Debug: Direct probabilities
        
        # Get predicted class and confidence
        class_index = np.argmax(predictions[0])
        class_name = CLASS_NAMES[class_index]
        probability = float(predictions[0][class_index])
        
        # Return prediction result
        return {
            "predicted_class": class_name,
            "confidence": probability
        }
    
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error processing image: {str(e)}")

# Run the app with: uvicorn main:app --reload
if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)