# from translator import translator
from lang_list import LANGUAGE_NAME_TO_CODE
import runpod

def run_t2tt(input_text: str, source_language: str, target_language: str) -> str:
    source_language_code = LANGUAGE_NAME_TO_CODE[source_language]
    target_language_code = LANGUAGE_NAME_TO_CODE[target_language]
    # out_texts, _ = translator.predict(
    #     input=input_text,
    #     task_str="T2TT",
    #     src_lang=source_language_code,
    #     tgt_lang=target_language_code,
    # )
    # return str(out_texts[0])
    import json
    return json.dumps({"input_text": input_text, "src_code": source_language_code, "tgt_code": target_language_code})

def runpod_handler(job):
    job_input = job['input']
    input_text = job_input["input_text"]
    source_language = job_input["source_language"]
    target_language = job_input["target_language"]
    return run_t2tt(input_text, source_language, target_language)

runpod.serverless.start({"handler": runpod_handler})