import gradio as gr
from src.translate_any_doc import translate_document, translate
from src.aligner import Aligner
from nltk.tokenize.treebank import TreebankWordDetokenizer


ip='10.192.31.127'
config_folder = 'fast_align_config'
source_lang = 'en'
target_lang = 'ca'
temp_folder = 'tmp'
aligner = Aligner(config_folder, source_lang, target_lang, temp_folder)
detokenizer = TreebankWordDetokenizer()


def upload_file(filepath):
    translated_file_name = translate_document(filepath, source_lang, target_lang, aligner, detokenizer, ip)
    return [gr.UploadButton(visible=False), gr.DownloadButton(label=f"Download {translated_file_name}", value=translated_file_name, visible=True)]

def download_file():
    return [gr.UploadButton(visible=True), gr.DownloadButton(visible=False)]


with gr.Blocks() as demo:

    with gr.Tab("Text"):
        gr.Interface(fn=translate, inputs=["text","text","text"], outputs="text")
    with gr.Tab("Docx documents"):
        gr.Markdown("First upload a file and and then you'll be able download it (but only once!)")
        with gr.Row():
            u = gr.UploadButton("Upload a file", file_count="single")
            d = gr.DownloadButton("Download the file", visible=False)

        u.upload(upload_file, u, [u, d])
        d.click(download_file, None, [u, d])
if __name__ == "__main__":
    demo.launch()