import os

import gradio as gr

import subprocess

from transformers import AutoModel

from huggingface_hub import snapshot_download

MODELS_PATH = "./models"
HF_CACHE_DIR = "./hf_cache"

def download_model(repo_id, revision="main"):
    return snapshot_download(repo_id=repo_id, revision=revision, local_dir=os.path.join(MODELS_PATH, repo_id), cache_dir=HF_CACHE_DIR)
def write_text_to_file(filename, text):
    with open(filename, 'w') as file:
        file.write(text)

model_dir_gl_en = download_model("proxectonos/Nos_MT-OpenNMT-gl-en", revision="main")
model_dir_en_gl = download_model("proxectonos/Nos_MT-OpenNMT-en-gl", revision="main")

def translate(input_text):
    write_text_to_file('input.txt', input_text)
    command = f"onmt_translate -src input.txt -model {os.path.join(MODELS_PATH, model_dir_en_gl)}/NOS-MT-OpenNMT-en-gl --output ./output_file.txt --replace_unk"
    process = subprocess.Popen(command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    stdout, stderr = process.communicate()
    if process.returncode != 0:
        raise Exception(f"Error occurred: {stderr.decode().strip()}")
    with open ('./output_file.txt','r') as f:
        resultado= f.read()
    return resultado

demo = gr.Interface(fn=translate, inputs="textbox", outputs="textbox")
demo.launch()