import os
import gradio as gr
import cv2
import numpy as np

def read_files(source_file):
    jpg_files = []
    txt_files = []
    for file in os.listdir(source_file):
        if file.endswith(".jpg"):
            jpg_files.append(os.path.join(source_file, file))
            txt_file = os.path.join(source_file, file[:-4] + ".txt")
            if os.path.exists(txt_file):
                txt_files.append(txt_file)
            else:
                txt_files.append(None)
    return jpg_files, txt_files

def load_jpg(jpg_file):
    img = cv2.imread(jpg_file)
    return cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

def load_txt(txt_file):
    with open(txt_file) as f:
        text = f.read().strip()
    return text

def save_txt(txt_file, text):
    with open(txt_file, "w") as f:
        f.write(text)

def main(source_file, target_file):
    jpg_files, txt_files = read_files(source_file)
    preview_list = gr.inputs.CheckboxGroup(jpg_files, label="Preview List")
    input_box = gr.inputs.Textbox(lines=5, placeholder="Enter text here", label="Text Input")
    gradio_interface = gr.Interface(
        fn=None,
        inputs=[preview_list],
        outputs=[gr.outputs.Image(type="numpy", label="Image Preview"), input_box],
        layout="vertical",
        theme="compact",
        title="Image Text Editor",
        description="Select an image from the preview list and edit the associated text."
    )

    @gradio_interface.func
    def image_text_editor(preview_list):
        selected_file = preview_list[0]
        selected_index = jpg_files.index(selected_file)
        jpg_file = jpg_files[selected_index]
        txt_file = txt_files[selected_index]
        img = load_jpg(jpg_file)
        if txt_file is not None:
            text = load_txt(txt_file)
            input_box.placeholder = text
            suggestions = [f"({i}) {word}" for i, word in enumerate(text.split(","))]
            output_box = gr.outputs.CheckboxGroup(suggestions, label="Text Suggestions")
            gradio_interface.set_output([img, input_box, output_box])
        else:
            gradio_interface.set_output([img, input_box])

    @gradio_interface.func
    def apply_suggestion(suggestion):
        original_text = input_box.value
        selected_index = int(suggestion[1:].split(")")[0])
        words = original_text.split(",")
        if selected_index < len(words):
            words[selected_index] = ""
            new_text = ",".join(words)
            input_box.update(new_text)

    @gradio_interface.func
    def save_image_text():
        selected_file = preview_list[0]
        selected_index = jpg_files.index(selected_file)
        jpg_file = jpg_files[selected_index]
        txt_file = os.path.join(target_file, os.path.basename(jpg_file)[:-4] + ".txt")
        save_txt(txt_file, input_box.value)
        new_jpg_file = os.path.join(target_file, os.path.basename(jpg_file))
        os.makedirs(os.path.dirname(new_jpg_file), exist_ok=True)
        os.rename(jpg_file, new_jpg_file)
        gradio_interface.set_output("Text and image saved successfully.")
    
    gradio_interface.add_output("Text Suggestions", apply_suggestion)
    gradio_interface.add_output("Save", save_image_text)
    gradio_interface.launch()

if __name__ == "__main__":


    source_file = gr.inputs.Textbox(lines=1,label="source_file")
    target_file = gr.inputs.Textbox(lines=1,label="target_file")

    interface = gr.Interface(
        fn=main,
        [source_file,target_file]
    )
    interface.launch()