# coding: utf-8

import gradio as gr
import os
import json
from pathlib import Path
from uuid import uuid4
from datetime import datetime
from huggingface_hub import CommitScheduler
from translation import Translator, LANGUAGES


LANGUAGES_LIST = list(LANGUAGES.keys())
HF_TOKEN = os.environ.get('HF_TOKEN')

JSON_DATASET_DIR = Path("flagged_data")
JSON_DATASET_DIR.mkdir(exist_ok=True, parents=True)
JSON_DATASET_PATH = JSON_DATASET_DIR / f"dataset-session-{uuid4()}.json"




def translate_wrapper(text, src_lang, tgt_lang, by_sentence=True, clean=True, num_beams=4):
    if text in ["", None, []]:
        return "Մուտքագրումը պարապ է։ | Input is empty."

    if src_lang in ["", None, []] or tgt_lang in ["", None, []]:
        return "Ընտրեցէք թարգմանութեան կողմերը | Please select source and target languages"

    if src_lang == tgt_lang:
        return "Ընտրուած լեզուները նոյնն են։ | Source and target languages are identical."

    src_lang = LANGUAGES.get(src_lang)
    tgt_lang = LANGUAGES.get(tgt_lang)

    result = translator.translate(text, src_lang, tgt_lang, by_sentence=by_sentence, clean=clean, num_beams=num_beams)
    return result


# hf_writer = gr.HuggingFaceDatasetSaver(HF_TOKEN, "AriNubar/hyw-en-crowd-source")

scheduler = CommitScheduler(
    repo_id="AriNubar/hyw-en-crowd-source",
    repo_type="dataset",
    folder_path=JSON_DATASET_DIR,
    path_in_repo="collected_data",
    token=HF_TOKEN,
    every=30 # every 30 minute
)

TQS = {
    "😊 Լաւ | Good": "good",
    "😐 Միջակ | Average": "average",
    "☹️ Վատ | Bad": "bad"
}

def save_json(src_lang, tgt_lang, input_text, output_text, by_sentence, clean, num_beams, translation_quality):
    if any([src_lang in ["", None, []], tgt_lang in ["", None, []], input_text in ["", None, []], output_text in ["", None, []]]):
        gr.Warning("Տուեալին մէկ մասը պարապ է։ Ձեր գնահատութիւնը չպահուեցաւ։ | Some part of the data is missing. Your feedback has not been saved.")
        return 
    
    src_lang = LANGUAGES.get(src_lang)
    tgt_lang = LANGUAGES.get(tgt_lang)
    translation_quality = TQS.get(translation_quality)
    
    # print(src_lang, tgt_lang, input_text, output_text, by_sentence, clean, num_beams, translation_quality)
    # print(type(src_lang), type(tgt_lang), type(input_text), type(output_text), type(by_sentence), type(clean), type(num_beams), type(translation_quality))
    data = {
        "src_lang": src_lang,
        "tgt_lang": tgt_lang,
        "original": input_text,
        "translation": output_text,
        "by_sentence": by_sentence,
        "clean": clean,
        "num_beams": num_beams,
        "translation_quality": translation_quality,
        "timestamp": datetime.now().isoformat()
    }
    with scheduler.lock:
        with open(JSON_DATASET_PATH, "a", encoding="utf8") as f:
            f.write(json.dumps(data, ensure_ascii=False) + "\n")

    gr.Info("Ձեր գնահանութիւնը պահուեցաւ։ Շատ շնորհակալութի՛ւն։ | Your feedback has been saved. Thank you.")


theme = gr.themes.Default().set(
    block_info_text_size="*text_xxs" # for info text
)

with gr.Blocks(title="Արեւմտահայերէն-Անգլերէն Մեքենական Թարգմանիչ | Western Armenian-English Machine Translation", 
               theme=theme,
               ) as demo:
    
    gr.HTML("""
        <h2 style='margin-bottom: 5px'>Արեւմտահայերէն-Անգլերէն Մեքենական Թարգմանիչ | Western Armenian-English Machine Translation</h2>
        <h3 style='margin-bottom: 5px'>Ստեղծող՝ | Created By: <a href='https://www.arinubar.com' target='_blank'>Ari Nubar Boyacıoğlu</a></h3>
        <p style="font-size: 0.7rem">Եթէ այս գործիքը կարողացաւ ձեզ օգտակար հանդիսանալ, բարելաւելու համար հաճեցէք սուրճի մը փոխարժէքը նուիրել․ | If this tool has proven useful to you, please consider making a donation. <a href='https://www.paypal.com/donate/?hosted_button_id=RRBCV3GQJ7D8N' target='_blank'>PayPal</a> | <a href='https://buymeacoffee.com/arinubar' target='_blank'>Buy Me a Coffee</a></p>
        <hr style='margin-top: 5px; margin-bottom: 5px'>
        <p>Հոս կը ցուցադրուի առաջին արեւմտահայերէն-անգլերէն մեքենական թարգմանիչը, որ կարուցուած է Մեթայի (Ֆեյսպուքի) 'No Language Left Behind' տիպարի հիման վրայ։ Թարգմանութեան տիպարը կ'աշխատի CPU-ի մը մէջ, ուրեմն նախադասութեան մը թարգմանութիւնը կրնայ տեւել մօտաւորապէս <strong>40-60 երկվայրկեան</strong>։ Ձեր գնահատութիւնները եւ քննադատութիւնները շատ կարեւոր են տիպարի թարգմանութեան որակը բարելաւելու համար։</p>
        <p>Դուք թարգմանութեան որակին մասին ձեր գնահատութիւնը կրնաք տալ երեք գնահատութեան կոճակներէ մէկուն սեղմելով։ Լեզուի, մուտքի եւ ելքի գրութիւններու, յարաչափերու եւ ձեր գնահատութեան մասին տուեալները պիտի պահուին։ Գնահատութիւնը պարտաւոր չէ։</p>
        <p>This is the demo of the first Western Armenian-English neural machine translation system which is based on Meta's 'No Language Left Behind' model. The model runs on a CPU, so it might take approximately <strong>40-60 seconds</strong> to translate a single sentence. Your feedback and comments are very important for us to improve the quality of the translation.</p>
        <p>You can give your feedback about the quality of the translation by clicking one of the three feedback buttons. Information about source, target languages, input and output texts, parameters and your feedback about quality will be saved. It is not mandatory to give feedback.</p>
    """)
    
    with gr.Row():
        with gr.Column():
            text = gr.Textbox(lines=5, label="Մուտքագրում | Input Text")
            with gr.Row():
                src_lang = gr.Dropdown(LANGUAGES_LIST, type="value", label="Թարգմանէ Այս Լեզուէ | Source Language")
                tgt_lang = gr.Dropdown(LANGUAGES_LIST, type="value", label="Թարգմանէ Այս Լեզուի | Target Language")
        with gr.Column():
            translated = gr.Textbox(lines=5, label="Ելքագրում | Output Text", interactive=False)
            translate_btn = gr.Button(value="Թարգմանէ | Translate", variant="primary")
            with gr.Row():
                with gr.Column():
                    gr.Markdown("""
                    ### Թարգմանութեան Որակ | Translation Quality
                    """)                
                    flag_good_btn = gr.Button(value="😊 Լաւ | Good", size="sm")
                    flag_average_btn = gr.Button(value="😐 Միջակ | Average", size="sm")
                    flag_bad_btn = gr.Button(value="☹️ Վատ | Bad", size="sm")
    with gr.Row():
        gr.Markdown("""
            ## Յարաչափեր | Parameters
        """
        )
        by_sentence = gr.Checkbox(label="Նախադասութիւններու Բաժնէ | Split into Sentences", value=True, info="Տուփը նշանագրեցէք եթէ կ'ուզէք ձեր մուտքագրումը թարգմանուի նախադասութիւն առ նախադասութիւն։ Այս կերպով թարգմանուած նախադասութիւններուն որակը ընդհանրապէս աւելի լաւ կ'ըլլան։ | Check this box if you want to split your input text into sentences. This way the quality of the translation will be better.")
        clean = gr.Checkbox(label="Մշակէ | Preprocess", value=True, info="Տուփը նշանագրեցէք եթէ կ'ուզէք ձեր մուտքագրումը կանոնաւորուի ծրագրի կողմէ թարգմանութենէ առաջ։ Կանոնաւորումը թարգմանութեան որակի բարելաւման համար օգտակար է։ | Check this box if you want to preprocess your input text before translation. This way the quality of the translation will be better.")
        num_beams = gr.Dropdown([1, 2, 3, 4, 5], type="value", label="Որոնման Շողեր | Number of Beams", value=4, info="Աւելի բարձր թիւը ընդհանրապէս կը պատճառէ աւելի բարձր որակի, բայց նոյնիսկ երկարատեւ թարգմանութեան։ | Higher beam size will result in better quality translation, but also longer translation time.")

    translate_btn.click(translate_wrapper, inputs=[text, src_lang, tgt_lang, by_sentence, clean, num_beams], outputs=translated)
    
    # hf_writer.setup([src_lang, tgt_lang, text, translated, by_sentence, clean, num_beams, flag_bad_btn], "flagged_data_points")
    
    flag_good_btn.click(save_json, inputs=[src_lang, tgt_lang, text, translated, by_sentence, clean, num_beams, flag_good_btn], outputs=None)
    flag_average_btn.click(save_json, inputs=[src_lang, tgt_lang, text, translated, by_sentence, clean, num_beams, flag_average_btn], outputs=None)
    flag_bad_btn.click(save_json, inputs=[src_lang, tgt_lang, text, translated, by_sentence, clean, num_beams, flag_bad_btn], outputs=None)

    visitor_badge_html = """
        <a href="https://visitorbadge.io/status?path=https%3A%2F%2Farinubar-hyw-en-demo.hf.space%2F">
            <img src="https://api.visitorbadge.io/api/visitors?path=https%3A%2F%2Farinubar-hyw-en-demo.hf.space%2F&label=%D4%B1%D5%B5%D6%81%D5%A5%D5%AC%D5%B8%D6%82%D5%B6%D5%A5%D6%80%20%7C%20Visitors&countColor=%23f97316&style=flat" />
        </a>
    """
    
    gr.HTML(visitor_badge_html)



    sponsors_html = """
        <div style="display: flex; justify-content: center; align-items: center; margin-bottom: 5px;">
            <h3>Աջակցութեամբ՝ | Supported By: </h3>
        </div>
        <div style="display: flex; justify-content: center; align-items: center;">
            <a href="https://gulbenkian.pt/armenian-communities/">
                <img src="file/img/gulbenkian.svg" alt="Calouste Gulbenkian Foundation - Armenian Communities" style="background-color: #414142; padding: 10px" />
            </a>
        </div>

    """
    gr.HTML(sponsors_html)


if __name__ == "__main__":
    translator = Translator()
    demo.launch(favicon_path="img/translate.png", share=True, allowed_paths=["./img"])