import io
import numpy as np
import tensorflow as tf
from fastapi import FastAPI, File, UploadFile
from fastapi.responses import JSONResponse
from PIL import Image
from pydantic import BaseModel

class TranslationResponse(BaseModel):
    prediction: str
    confidence: float

try:
    model = tf.keras.models.load_model('best_model_2.keras')
except Exception as e:
    raise IOError(f"Error al cargar el modelo 'best_model.keras'. Error: {e}")

app = FastAPI(
    title="API keras asl",
    description="Sube una imagen del alfabeto de señas (ASL) para obtener una predicción del modelo.",
    version="1.0.0"
)

CLASS_NAMES = [
    'A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J', 'K', 'L', 'M',
    'N', 'O', 'P', 'Q', 'R', 'S', 'T', 'U', 'V', 'W', 'X', 'Y', 'Z'
]

def preprocess_image(image: Image.Image) -> np.ndarray:
    image = image.resize((96, 96))
    image_array = np.array(image)
    
    if image_array.shape[2] == 4: # Maneja imágenes RGBA
        image_array = image_array[..., :3]
    
    return np.expand_dims(image_array, axis=0)

@app.post("/predict/", response_model=TranslationResponse)
async def predict(file: UploadFile = File(...)):
    contents = await file.read()
    
    try:
        image = Image.open(io.BytesIO(contents)).convert('RGB')
    except Exception as e:
        return JSONResponse(status_code=400, content={"message": f"Error al leer la imagen: {e}"})

    processed_image = preprocess_image(image)

    predictions = model.predict(processed_image)
    
    predicted_index = np.argmax(predictions, axis=1)[0]
    confidence = float(predictions[0][predicted_index])
    prediction_label = CLASS_NAMES[predicted_index]

    return TranslationResponse(prediction=prediction_label, confidence=confidence)

@app.get("/")
def read_root():
    return {"message": "/predict/ para leer imagen."}