import gradio as gr
import json
import os
from generate_eban import translate_es_to_mapudungun

__TITLE: str = "# Traductor Español - Mapudungun"

__DESCRIPTION: str = """
Ingrese texto en español para traducirlo al mapudungun.
Luego, puede dar feedback sobre la traducción para ayudar a mejorar nuestro modelo.
"""

__FOOTER: str = """
---
(c) 2024 [Inria Chile Research Center](https://inria.cl)
"""

__LOGO_URI: str = """https://encrypted-tbn0.gstatic.com/images?q=tbn:ANd9GcSQHNnGyXMYCS-zyYFq-4NLkRCF9AwQmxwdjw&s"""

# using eban's model
def translate_text(text):
    return translate_es_to_mapudungun(text)

# function to save feedback
def save_feedback(input_text, output_text, feedback, correction=None):
    feedback_data = {
        "input": input_text,
        "output": output_text,
        "feedback": feedback,
        "correction": correction
    }
    
    # ensure the feedback directory exists
    os.makedirs("feedback", exist_ok=True)
    
    # save feedback to a JSON file
    with open(f"feedback/feedback_{len(os.listdir('feedback'))}.json", "w") as f:
        json.dump(feedback_data, f)

# function to initial tests
def translate_and_store(input_text):
    translated_text = translate_text(input_text)
    return translated_text, *clear_feedback()

def clear_feedback():
    return (gr.update(value=""), 
            gr.update(value="", visible=False), 
            gr.update(visible=False), 
            gr.update(visible=False))

# feedback functions
def like_translation(input_text, output_text):
    save_feedback(input_text, output_text, "liked")
    return ("Gracias por tu feedback positivo!", 
            gr.update(visible=False), 
            gr.update(visible=False), 
            gr.update(visible=False))

def dislike_translation(input_text, output_text):
    return ("¿Puedes ingresar la traducción correcta?", 
            gr.update(visible=True), 
            gr.update(visible=True), 
            gr.update(visible=True), 
            gr.update(value="", visible=False), 
            gr.update(visible=False))

def yes_provide_correction():
    return (gr.update(visible=True), 
            gr.update(visible=True), 
            gr.update(visible=False), 
            gr.update(visible=False))

def no_provide_correction(input_text, output_text):
    save_feedback(input_text, output_text, "disliked")
    return ("Gracias por usar la aplicación.", 
            gr.update(visible=False), 
            gr.update(visible=False), 
            gr.update(visible=False))

def submit_correction(input_text, output_text, correction):
    save_feedback(input_text, output_text, "disliked", correction)
    return ("Gracias. Tu corrección nos ayudará a mejorar nuestro modelo de traducción.", 
            gr.update(value="", visible=False), 
            gr.update(visible=False), 
            gr.update(visible=False))

# create the gradio interface
with gr.Blocks() as demo:
    gr.Markdown(__TITLE)
    gr.Markdown(__DESCRIPTION)
    
    with gr.Row():
        input_text = gr.Textbox(
            label="Ingrese texto en español",
            placeholder="Escriba aquí su texto en español...",
            lines=5,  
            max_lines=10,  
        )
        output_text = gr.Textbox(
            label="Traducción al mapudungun",
            lines=5,  
            max_lines=10,  
        )
    
    translate_btn = gr.Button("Translate")

    gr.Markdown("## ¿Cuál es su opinión sobre la traducción?")

    with gr.Row():
        like_btn = gr.Button("Me gusta la traducción")
        dislike_btn = gr.Button("No me gusta")
    
    feedback_output = gr.Textbox(label="Feedback", interactive=False)
    
    with gr.Row(visible=False) as correction_choice_row:
        yes_btn = gr.Button("Sí")
        no_btn = gr.Button("No")
    
    correction_input = gr.Textbox(label="Traducción correcta", visible=False)
    submit_correction_btn = gr.Button("Enviar corrección", visible=False)
    
    translate_btn.click(
        translate_and_store, 
        inputs=input_text, 
        outputs=[output_text, feedback_output, correction_input, submit_correction_btn, correction_choice_row]
    )
    like_btn.click(
        like_translation, 
        inputs=[input_text, output_text], 
        outputs=[feedback_output, correction_input, submit_correction_btn, correction_choice_row]
    )
    dislike_btn.click(
        dislike_translation, 
        inputs=[input_text, output_text], 
        outputs=[feedback_output, correction_choice_row, yes_btn, no_btn, correction_input, submit_correction_btn]
    )
    yes_btn.click(
        yes_provide_correction, 
        outputs=[correction_input, submit_correction_btn, yes_btn, no_btn]
    )
    no_btn.click(
        no_provide_correction, 
        inputs=[input_text, output_text], 
        outputs=[feedback_output, correction_input, submit_correction_btn, correction_choice_row]
    )
    submit_correction_btn.click(
        submit_correction, 
        inputs=[input_text, output_text, correction_input], 
        outputs=[feedback_output, correction_input, submit_correction_btn, correction_choice_row]
    )
    
    
    with gr.Row():
        with gr.Column(scale=2):
            gr.Markdown(__FOOTER)
        with gr.Column(scale=1):
            gr.HTML(value=f"<img width='92%' src='{__LOGO_URI}' alt='logo'/>")
demo.launch()