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

app = FastAPI(title="Facial Emotion Recognition API")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

MODEL_PATH = "Fer2013.h5"
EMOTION_LABELS = ["Angry", "Disgust", "Fear", "Happy", "Sad", "Surprise", "Neutral"]

model = None


@app.on_event("startup")
def load_fer_model():
    global model
    if not os.path.exists(MODEL_PATH):
        raise FileNotFoundError(
            f"Model file '{MODEL_PATH}' not found. Make sure it is uploaded to the Space root."
        )
    model = load_model(MODEL_PATH, compile=False)
    print("Model loaded successfully.")


def preprocess_image(image_bytes: bytes) -> np.ndarray:
    image = Image.open(io.BytesIO(image_bytes)).convert("L")  # grayscale
    image = image.resize((48, 48))
    array = np.array(image, dtype=np.float32) / 255.0
    array = array.reshape(1, 48, 48, 1)
    return array


@app.get("/")
def root():
    return {"status": "ok", "message": "Facial Emotion Recognition API is running."}


@app.post("/predict")
async def predict(file: UploadFile = File(...)):
    contents = await file.read()
    input_array = preprocess_image(contents)

    predictions = model.predict(input_array)[0]
    predicted_idx = int(np.argmax(predictions))

    result = {
        "predicted_emotion": EMOTION_LABELS[predicted_idx],
        "confidence": float(predictions[predicted_idx]),
        "all_probabilities": {
            EMOTION_LABELS[i]: float(predictions[i]) for i in range(len(EMOTION_LABELS))
        },
    }
    return result