from tensorflow.keras.models import load_model # type: ignore
from tensorflow.keras.preprocessing import image as keras_image # type: ignore
import numpy as np
from PIL import Image
import io
import os
import gdown

# === Path model lokal ===
model_path = 'saved_model_palm_disease.keras'

# === Unduh model dari Google Drive jika belum ada ===
if not os.path.exists(model_path):
    url = 'https://drive.google.com/uc?id=1g-QPUIsySVm1oBl0KXpKKlxe7x_JPe7B'
    gdown.download(url, model_path, quiet=False)

# === Load model hanya sekali ===
model = load_model(model_path)

# === Label urutan class_name (dari training) ===
labels = ['Boron Excess', 'Ganoderma', 'Healthy', 'Scale insect']

# === Fungsi preprocessing gambar ===
def preprocess_image(image_bytes):
    img = Image.open(io.BytesIO(image_bytes)).convert("RGB").resize((224, 224))
    img_array = keras_image.img_to_array(img) / 255.0
    img_array = np.expand_dims(img_array, axis=0)
    return img_array

# === Fungsi prediksi utama ===
def predict_image(image_bytes):
    img_array = preprocess_image(image_bytes)
    predictions = model.predict(img_array)
    class_index = int(np.argmax(predictions))
    confidence = float(np.max(predictions))
    return {
        'class': labels[class_index],
        'confidence': confidence
    }

# === Fungsi handler yang dipanggil Hugging Face API ===
def handler(inputs):
    try:
        # Input berupa base64 atau byte dari gambar
        image_bytes = inputs['inputs']

        # Jika data berupa string base64, ubah ke byte
        if isinstance(image_bytes, str):
            import base64
            image_bytes = base64.b64decode(image_bytes)

        result = predict_image(image_bytes)
        return result

    except Exception as e:
        return {"error": str(e)}