import gradio as gr
from dotenv import load_dotenv
from roboflow import Roboflow
import tempfile
import os
import requests

# Muat variabel lingkungan dari file .env
load_dotenv()
api_key = os.getenv("ROBOFLOW_API_KEY")
workspace = os.getenv("ROBOFLOW_WORKSPACE")
project_name = os.getenv("ROBOFLOW_PROJECT")
model_version = int(os.getenv("ROBOFLOW_MODEL_VERSION"))

# Inisialisasi Roboflow menggunakan data yang diambil dari secrets
rf = Roboflow(api_key=api_key)
project = rf.workspace(workspace).project(project_name)
model = project.version(model_version).model

# Fungsi untuk menangani input dan output gambar
def detect_objects(image):
    # Simpan gambar yang diupload sebagai file sementara
    with tempfile.NamedTemporaryFile(delete=False, suffix=".jpg") as temp_file:
        image.save(temp_file, format="JPEG")
        temp_file_path = temp_file.name

    try:
        # Lakukan prediksi pada gambar
        predictions = model.predict(temp_file_path, confidence=55, overlap=80).json()

        # Menghitung jumlah objek per kelas
        class_count = {}
        total_count = 0  # Menyimpan total jumlah objek

        for prediction in predictions['predictions']:
            class_name = prediction['class']
            class_count[class_name] = class_count.get(class_name, 0) + 1
            total_count += 1  # Tambah jumlah objek untuk setiap prediksi

        # Menyusun output berupa string hasil perhitungan
        result_text = "Product Nestle\n\n"
        for class_name, count in class_count.items():
            result_text += f"{class_name}: {count}\n"
        result_text += f"\nTotal Product Nestle: {total_count}"

        # Menyimpan gambar dengan prediksi
        output_image_path = "/tmp/prediction.jpg"
        model.predict(temp_file_path, confidence=55, overlap=80).save(output_image_path)
        
    except requests.exceptions.HTTPError as http_err:
        # Menangani kesalahan HTTP
        result_text = f"HTTP error occurred: {http_err}"
        output_image_path = temp_file_path  # Kembalikan gambar asli jika terjadi error
    except Exception as err:
        # Menangani kesalahan lain
        result_text = f"An error occurred: {err}"
        output_image_path = temp_file_path  # Kembalikan gambar asli jika terjadi error

    # Hapus file sementara setelah prediksi
    os.remove(temp_file_path)
    
    return output_image_path, result_text

# Membuat antarmuka Gradio dengan tata letak fleksibel
with gr.Blocks() as iface:
    with gr.Row():
        with gr.Column():
            input_image = gr.Image(type="pil", label="Input Image")
        with gr.Column():
            output_image = gr.Image(label="Detect Object")
        with gr.Column():
            output_text = gr.Textbox(label="Counting Object")
    
    # Tombol untuk memproses input
    detect_button = gr.Button("Detect")
    
    # Hubungkan tombol dengan fungsi deteksi
    detect_button.click(
        fn=detect_objects, 
        inputs=input_image, 
        outputs=[output_image, output_text]
    )

# Menjalankan antarmuka
iface.launch()