from src.aligner import Aligner
from src.salamandraTA7b_translator import SalamandraTA7bTranslator, SalamandraTA7bQTranslator
from src.salamandraTA7b_translator_HF import SalamandraTA7bTranslatorHF
import os
import time
import argparse

if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        prog='main',
        description='Translate a file')
    parser.add_argument("-s", '--source_lang', type=str, required=True)
    parser.add_argument("-t", '--target_lang', type=str, required=True)
    parser.add_argument("-f", '--file_path', type=str, required=True)

    parser.add_argument("-m", '--model_path', type=str, required=True)
    parser.add_argument("-tt", '--translator_type', type=str,
                        choices=["normal", "quantized"], default="none",
                        help="normal=regular model; quantized=quantized model")

    parser.add_argument('--fastalign_config_folder', type=str, default="fast_align_config")
    parser.add_argument('--temp_folder', type=str, default="tmp")

    args = parser.parse_args()

    os.makedirs(args.temp_folder, exist_ok=True)

    if args.translator_type == "normal":
        translator = SalamandraTA7bTranslator(args.model_path)
    elif args.translator_type == "quantized":
        translator = SalamandraTA7bQTranslator(args.model_path)
    else:
        raise NotImplementedError(f"Option {args.translator_type} is not implemented.")

    aligner = Aligner(args.fastalign_config_folder, args.source_lang, args.target_lang, args.temp_folder)

    start_time = time.time()
    for status, translated_file_name in translator.translate_document(args.file_path, args.source_lang,
                                                                      args.target_lang):
        if translated_file_name:  # finished
            break
        else:
            print(status)
    print(f"Finished document in {time.time() - start_time} seconds")