import argparse
import glob
import os
import time
import argostranslate.package
import argostranslate.translate


def setup_translator(from_code='en', to_code='ro'):
    """Configurează și instalează pachetele necesare pentru ArgosTranslate"""
    # Actualizăm indexul de pachete
    argostranslate.package.update_package_index()
    
    # Găsim și instalăm pachetul necesar
    available_packages = argostranslate.package.get_available_packages()
    package_to_install = next(
        filter(
            lambda x: x.from_code == from_code and x.to_code == to_code, 
            available_packages
        ),
        None
    )
    
    if package_to_install is None:
        print(f"Nu a fost găsit pachet de traducere de la {from_code} la {to_code}")
        return False
    
    # Instalăm pachetul
    argostranslate.package.install_from_path(package_to_install.download())
    print(f"Pachetul de traducere {from_code} -> {to_code} a fost instalat")
    return True


def loop_translation(target_file_path, from_code='en', to_code='ro'):
    """Încearcă traducerea unui fișier cu ArgosTranslate de mai multe ori în caz de eroare"""
    for loop in range(3):
        try:
            with open(target_file_path, 'rt', encoding='utf-8', errors='ignore') as file:
                text = file.read()
            
            # Autodetectarea limbii sursă nu este suportată în ArgosTranslate, 
            # dar o putem seta manual la 'en' sau parametrul from_code
            translated_text = argostranslate.translate.translate(text, from_code, to_code)
            return translated_text
            
        except Exception as exception:
            print(f"Eroare la traducere (încercarea {loop+1}/3): {str(exception)}")
        time.sleep(1)
    return None


def main():
    parser = argparse.ArgumentParser(description='Athos article translator with ArgosTranslate')
    parser.add_argument(
        '--in_dir',
        type=str,
        default='data_en_split',
        required=False,
    )
    parser.add_argument(
        '--out_dir',
        type=str,
        default='data_ro',
        required=False,
    )
    parser.add_argument(
        '--src',
        type=str,
        default='en',
        required=False,
        help='Source language code (default: en)'
    )
    parser.add_argument(
        '--dest',
        type=str,
        default='ro',
        required=False,
        help='Destination language code (default: ro)'
    )

    args = parser.parse_args()

    # Creăm directorul de ieșire dacă nu există
    if not os.path.exists(args.out_dir):
        os.makedirs(args.out_dir)

    # Configurăm și instalăm pachetele pentru ArgosTranslate
    if not setup_translator(args.src, args.dest):
        print("Nu s-a putut configura translația. Se oprește scriptul.")
        return

    # Verificăm dacă putem traduce un text simplu
    test_translation = argostranslate.translate.translate("veritas lux mea", args.src, args.dest)
    print(f'Test de traducere: "veritas lux mea" -> "{test_translation}"')

    # Găsim fișierele pentru traducere
    pattern = f"{args.in_dir}/*.txt"
    txt_files = glob.glob(pattern)
    
    if not txt_files:
        print(f"Nu s-au găsit fișiere text în {args.in_dir}")
        return

    print(f"S-au găsit {len(txt_files)} fișiere pentru traducere.")

    # Procesăm fiecare fișier
    counter = 0
    for file_path in txt_files:
        print(f'Se procesează: {file_path}')

        translated_text = loop_translation(file_path, args.src, args.dest)
        if translated_text is not None:
            output_path = f'{args.out_dir}/{os.path.basename(file_path)}'
            with open(output_path, 'wt', encoding='utf-8', errors='ignore') as file_out:
                file_out.write(translated_text)
            print(f'Traducere salvată în: {output_path}')
        else:
            print(f'Traducerea fișierului {file_path} a eșuat. Se continuă cu următorul.')
        
        time.sleep(1)

        counter += 1
        if counter == 100000:
            break


if __name__ == "__main__":
    main()