import gradio as gr
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
import torch

# this model was loaded from https://hf.co/models
device = 'cuda' if torch.cuda.is_available() else 'cpu'
model = AutoModelForSeq2SeqLM.from_pretrained("facebook/nllb-200-distilled-600M").to(device)
tokenizer = AutoTokenizer.from_pretrained("facebook/nllb-200-distilled-600M")


LANGS = ["pes_Arab", "ckb_Arab", "eng_Latn"]
langs_dict = {
    "فارسی": "pes_Arab",
    "کردی": "ckb_Arab",
    "انگلیسی": "eng_Latn"
    }

def translate(text, src_lang, tgt_lang):
    """
    Translate the text from source lang to target lang
    """
    translation_pipeline = pipeline("translation", model=model, tokenizer=tokenizer, src_lang=langs_dict[src_lang], tgt_lang=langs_dict[tgt_lang], max_length=400, device=device)
    result = translation_pipeline(text)
    return result[0]['translation_text']

def file_translate(sorce_file_path, src_lang, tgt_lang, pred_file_path):
    sorce_list = []
    with open(sorce_file_path, "r", encoding="utf-8") as sorce_file:
        for line in sorce_file:
            sorce_list.append(line.strip())

    pred_list = []
    for line in sorce_list:
        pred_list.append(translate(line, src_lang, tgt_lang))

    with open(pred_file_path, "w", encoding="utf-8") as output_file:
        for translation in pred_list:
            output_file.write(translation + "\n")
    return pred_file_path

def add_line(input_path, output_path):
    # خواندن محتوای فایل ورودی
    with open(input_path, encoding="utf-8") as f:
        text = f.read()

    # اضافه کردن خط "سلام" به انتهای متن
    new_text = text + "\nسلام"

    # نوشتن متن جدید در فایل خروجی
    with open(output_path, "w", encoding="utf-8") as f:
        f.write(new_text)

    return output_path

if __name__ == '__main__':

    interface = gr.Interface(
    fn=file_translate,
    inputs=[
        gr.components.File(label="Input File"),
        gr.components.Dropdown(label="زبان مبدا", choices=list(langs_dict.keys())),
        gr.components.Dropdown(label="زبان مقصد", choices=list(langs_dict.keys())),
        gr.components.Textbox(label="Output File Name (optional)"),
    ],
    outputs=[
        gr.components.File(label="Modified File"),
    ],
    title="NLLB 200 - (Translation Demo)",
    description="This Gradio demo translate text files. (CPU)",
    )
    interface.launch()