import gradio as gr
from src.translate_any_doc import translate_document
from src.salamandraTA7b_translator import SalamandraTA7bTranslator
from src.aligner import Aligner
import os

config_folder = 'fast_align_config'
temp_folder = 'tmp'


def upload_file(filepath, source_lang, target_lang, user_token):
    hf_token = os.environ.get('HF_TOKEN')
    if user_token:
        hf_token = user_token
    translator = SalamandraTA7bTranslator(hf_token)
    aligner = Aligner(config_folder, source_lang, target_lang, temp_folder)
    for status, translated_file_name in translate_document(filepath, source_lang, target_lang, translator, aligner):
        if translated_file_name:  # finished
            yield [gr.UploadButton(visible=False),
                   gr.DownloadButton(label=f"Download {translated_file_name}", value=translated_file_name,
                                     visible=True, interactive=True),
                   gr.Textbox(visible=False)]
        else:
            yield [gr.UploadButton(visible=False),
                   gr.DownloadButton(visible=False),
                   gr.Textbox(value=status, visible=True)]


def before_processing():
    return [
        gr.UploadButton(visible=False),
        gr.Textbox(value="Processing...", 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=translator.translate, inputs=["text", "text", "text"], outputs="text")
    with gr.Tab("Documents"):
        with gr.Row():
            dropdown1 = gr.Dropdown(label="Source language", choices=["en", "ca"], value=None, interactive=True)
            dropdown2 = gr.Dropdown(label="Target language", choices=["en", "ca"], value=None, interactive=True)
            token_textbox = gr.Textbox(label="Huggingface token (optional)", visible=True)
        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)
            status_text = gr.Textbox(label="Status", visible=False)

        u.upload(fn=before_processing, inputs=None, outputs=[u, status_text]).then(upload_file,
                                                                                   [u, dropdown1, dropdown2,
                                                                                    token_textbox],
                                                                                   [u, d, status_text])
        d.click(download_file, None, [u, d])
if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)