import os
import json
import cv2
import gradio as gr
import numpy as np
from PIL import Image

# Path to images folder
image_folder = "./images"
output_folder = "./jsons"

# Ensure output folder exists
os.makedirs(output_folder, exist_ok=True)

# Get all image files
image_files = sorted([f for f in os.listdir(image_folder) if f.lower().endswith(('.png', '.jpg', '.jpeg'))])
image_index = 0  # Start at first image

# JSON template for annotation
default_json_template = {
    "image_path": "",
    "text": {
        "TITLE": "",
        "TEMPAT-TANGGAL": "",
        "NOMOR REGISTRASI": "",
        "No.": "",
        "NAMA PEMILIK": "",
        "ALAMAT": "",
        "MERK": "",
        "TYPE": "",
        "JENIS": "",
        "MODEL": "",
        "TAHUN PEMBUATAN": "",
        "ISI SILINDER/DAYA LISTRIK": "",
        "NOMOR RANGKA": "",
        "NOMOR MESIN": "",
        "NIK/TDP/NIB/KITAS/KITAP": "",
        "WARNA": "",
        "BAHAN BAKAR": "",
        "WARNA TNKB": "",
        "TAHUN REGISTRASI": "",
        "NOMOR BPKB": "",
        "KODE LOKASI": "",
        "NO URUT PENDAFTARAN": "",
        "BERLAKU SAMPAI": ""
    }
}

def load_image():
    global image_index
    if 0 <= image_index < len(image_files):
        image_path = os.path.join(image_folder, image_files[image_index])
        image = Image.open(image_path)
        
        json_filename = f"license_{image_index+1:03d}.json"
        json_path = os.path.join(output_folder, json_filename)
        annotation_data = default_json_template.copy()

        if os.path.exists(json_path):
            with open(json_path, "r") as f:
                annotation_data = json.load(f)

        return image, annotation_data["text"], image_path
    return None, {}, ""

def save_annotation(image_path, annotation_text):
    global image_index
    if 0 <= image_index < len(image_files):
        json_filename = f"license_{image_index+1:03d}.json"
        json_path = os.path.join(output_folder, json_filename)
        
        annotation_data = {"image_path": image_path, "text": annotation_text}
        
        with open(json_path, "w") as f:
            json.dump(annotation_data, f, indent=4)
        
        return f"✅ Saved: {json_filename}"
    return "❌ Error saving annotation"

def next_image():
    global image_index
    if image_index < len(image_files) - 1:
        image_index += 1
    return load_image()

def prev_image():
    global image_index
    if image_index > 0:
        image_index -= 1
    return load_image()

# Gradio Interface
with gr.Blocks() as demo:
    with gr.Row():
        image_display = gr.Image()
        
        with gr.Column():
            text_inputs = {key: gr.Textbox(label=key) for key in default_json_template["text"]}
            save_button = gr.Button("Save")
    
    status = gr.Textbox(label="Status", interactive=False)
    
    with gr.Row():
        prev_button = gr.Button("Previous")
        next_button = gr.Button("Next")
    
    # Load first image on startup
    image, annotation, img_path = load_image()
    image_display.value = image
    for key, textbox in text_inputs.items():
        textbox.value = annotation.get(key, "")
    
    # Button Actions
    def save_action():
        annotation_text = {key: text_inputs[key].value for key in text_inputs}
        return save_annotation(img_path, annotation_text)
    
    save_button.click(save_action, outputs=status)
    next_button.click(lambda: next_image(), outputs=[image_display, *text_inputs.values()])
    prev_button.click(lambda: prev_image(), outputs=[image_display, *text_inputs.values()])

demo.launch()