from absl import app
from absl import flags
from absl import logging

import euler

from base_thrift import Base
from MetaTranslation_thrift import MetaTranslation, RegularTranslationRequest

FLAGS = flags.FLAGS
flags.DEFINE_string('input', None, 'File to be translated.')
flags.DEFINE_string('output', None, 'File to be store translated sentences.')
flags.DEFINE_string('src_lang', None, 'Source language.')
flags.DEFINE_string('trg_lang', None, 'Target language.')
flags.DEFINE_integer('batch_size', 32, 'Send batch_size sentences at one request.', lower_bound=1)
flags.DEFINE_boolean('use_google', False, 'Only use google translate.')

consul = "sd://lab.mt.chief"
client = euler.Client(
    MetaTranslation, consul, transport="buffered", timeout=10)
psm = "test"
# key = "7f30e9315edc6a877b621139f42e3740"
key = "7f30e9358adb2be77b621139f42e3740"

def do_translate(text_list):
    if FLAGS.src_lang == FLAGS.trg_lang:
        return text_list

    if FLAGS.use_google:
        api_list = ['google_key']
    else:
        api_list = ['lab']
    
    request = RegularTranslationRequest(text_list=text_list,
                                        src_lang = FLAGS.src_lang,
                                        trg_lang= FLAGS.trg_lang,
                                        api_list = api_list,
                                        id=psm, key=key,Base=Base(Caller="", Addr=""))
    try:
        res = client.regular_translation(request)
    except:
        logging.warning("Error return for request with %d examples", len(text_list))
        return text_list
    ret = []
    if res.translation_resp_list == None:
        logging.warning("Empty result for request with %d examples", len(text_list))
        return text_list
    for oneret in res.translation_resp_list:
        ret.append(oneret.translated_text)
    return ret

def main(argv):
    if FLAGS.log_dir:
        logging.get_absl_handler().use_absl_log_file('log', FLAGS.log_dir)
    if FLAGS.src_lang is None or FLAGS.trg_lang is None:
        logging.error('src_lang and trg_lang can\'t be None')
    if FLAGS.src_lang == FLAGS.trg_lang:
        logging.warning('src_lang and trg_lang are the same!')
    if FLAGS.input is None or FLAGS.output is None:
        logging.error('input and output can\'t be None')
        exit(0)

    logging.info('input file: %s\noutfile: %s\nbatch_size: %s\nsrc_lang: %s trg_lang:%s\n',
                    FLAGS.input, FLAGS.output, FLAGS.batch_size, FLAGS.src_lang, FLAGS.trg_lang)
    linenum = 0
    with open(FLAGS.input) as f, open(FLAGS.output, 'w') as fw:
        text_list = []
        for line in f:
            text_list.append(line.strip())
            if len(text_list) >= FLAGS.batch_size:
                res = do_translate(text_list)
                for oneres in res:
                    fw.write(oneres+'\n')
                text_list = []
                linenum += FLAGS.batch_size
                logging.info('finish %d sentences!', linenum)

        if len(text_list) > 0:
            res = do_translate(text_list)
            for oneres in res:
                fw.write(oneres+'\n')
            linenum += len(text_list)
            text_list = []

        logging.info('finish %d sentences!', linenum)

if __name__ == '__main__':
    app.run(main)