from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.responses import JSONResponse
from typing import List
from io import BytesIO
from PIL import Image
import numpy as np
import tensorflow as tf
from fastapi.middleware.cors import CORSMiddleware

# Inisialisasi aplikasi FastAPI
app = FastAPI()
app.add_middleware(
    CORSMiddleware,
    allow_origins=["https://trash-scan-fe.vercel.app"],  # Gantilah dengan URL frontend Anda
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Muat model yang sudah dilatih sebelumnya
MODEL_PATH = "models/model.h5"  # Ganti dengan path model Anda
try:
    model = tf.keras.models.load_model(MODEL_PATH)
except Exception as e:
    raise RuntimeError(f"Gagal memuat model: {e}")

# Label kelas (sesuaikan dengan model Anda)
CLASS_NAMES = ["battery", "biological", "clothes", "metal", "plastic"]  # Ganti dengan label model Anda


def read_imagefile(file: bytes) -> Image.Image:
    """Membaca file gambar dari bytes"""
    try:
        image = Image.open(BytesIO(file)).convert("RGB")
        return image
    except Exception as e:
        raise HTTPException(status_code=400, detail=f"Gagal membaca file gambar: {e}")


def preprocess_image(image: Image.Image, target_size: tuple) -> np.ndarray:
    """Preproses gambar sesuai dengan kebutuhan model"""
    image = image.resize(target_size)
    image_array = np.array(image) / 255.0  # Normalisasi
    image_array = np.expand_dims(image_array, axis=0)  # Tambahkan batch dimension
    return image_array


@app.post("/predict")
async def predict_image(file: UploadFile = File(...)):
    """Endpoint untuk prediksi gambar"""
    # Tambahkan log awal untuk debugging
    print(f"File diterima: {file.filename}, Content-Type: {file.content_type}")

    # Validasi tipe file
    if not file.content_type.startswith("image"):
        raise HTTPException(status_code=400, detail="File yang diunggah bukan gambar")

    # Baca dan preproses gambar
    try:
        image = read_imagefile(await file.read())
        image_array = preprocess_image(image, target_size=(224, 224))  # Sesuaikan ukuran dengan model Anda
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Gagal memproses gambar: {e}")

    # Lakukan prediksi menggunakan model
    try:
        predictions = model.predict(image_array)
        predicted_class = CLASS_NAMES[np.argmax(predictions)]
        confidence = float(np.max(predictions))
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Gagal melakukan prediksi: {e}")

    # Kembalikan hasil prediksi
    return JSONResponse(content={
        "predicted_class": predicted_class,
        "confidence": confidence,
        "filename": file.filename  # Tambahkan informasi nama file untuk debugging
    })


# Jalankan server
# .venv\Scripts\activate  # Aktifkan virtual environment
# Gunakan `uvicorn main:app --reload` untuk menjalankan API