from model.vi_lo import translator_vi_lo
from model.vi_en import translator_vi_en
from model.vi_zh import translator_vi_zh
from model.vi_km import translator_vi_km


class Translator:
    def __init__(self):
        self.vi_en_model = translator_vi_en.Translator_Vi_En()
        self.vi_lo_model = translator_vi_lo.Translator_Vi_Lo()
        self.vi_zh_model = translator_vi_zh.Translator_Vi_Zh()
        self.vi_km_model = translator_vi_km.Translator_Vi_Km()

    def translate(self, sen, src_lang, tgt_lang, model_type):
        if src_lang == tgt_lang:
            return sen
        
        mid_sen = ''
        if src_lang == 'vi':
            mid_sen = sen
        elif src_lang == 'en':
            mid_sen = self.vi_en_model.translate(sen, src_lang, 'vi', model_type)
        elif src_lang == 'lo':
            mid_sen = self.vi_lo_model.translate(sen, src_lang, 'vi', model_type)
        elif src_lang == 'zh':
            mid_sen = self.vi_zh_model.translate(sen, src_lang, 'vi', model_type)
        elif src_lang == 'km':
            mid_sen = self.vi_km_model.translate(sen, src_lang, 'vi', model_type)

        if tgt_lang == 'vi':
            return mid_sen
        elif tgt_lang == 'en':
            return self.vi_en_model.translate(sen, 'vi', tgt_lang, model_type)
        elif tgt_lang == 'lo':
            return self.vi_lo_model.translate(sen, 'vi', tgt_lang, model_type)
        elif tgt_lang == 'zh':
            return self.vi_zh_model.translate(sen, 'vi', tgt_lang, model_type)
        elif tgt_lang == 'km':
            return self.vi_km_model.translate(sen, 'vi', tgt_lang, model_type)

        return "None"