MeshChatX/scripts/argos_translate.py

#!/usr/bin/env python3
"""Argos Translate JSON localization script.

This script provides an automated workflow to translate JSON localization files
(such as `en.json`) to target languages using Argos Translate. It ensures that
interpolated variables (e.g., `{count}`, `{status}`) are preserved and not
altered during the translation process.

Requirements:
    - argostranslate (pip install argostranslate)

Usage:
    python scripts/argos_translate.py --from en --to zh --input locales/en.json --output locales/zh.json

    You can also use environment variables instead of CLI flags:
    ARGOS_FROM_LANG="en"
    ARGOS_TO_LANG="zh"
    ARGOS_INPUT_FILE="locales/en.json"
    ARGOS_OUTPUT_FILE="locales/zh.json"
"""

import argparse
import json
import os
import re
import sys
import time


# ANSI color codes for terminal output
class Colors:
    HEADER = "\033[95m"
    OKBLUE = "\033[94m"
    OKCYAN = "\033[96m"
    OKGREEN = "\033[92m"
    WARNING = "\033[93m"
    FAIL = "\033[91m"
    ENDC = "\033[0m"
    BOLD = "\033[1m"
    UNDERLINE = "\033[4m"


def print_info(msg):
    print(f"{Colors.OKCYAN}[INFO]{Colors.ENDC} {msg}")


def print_success(msg):
    print(f"{Colors.OKGREEN}[SUCCESS]{Colors.ENDC} {msg}")


def print_warning(msg):
    print(f"{Colors.WARNING}[WARNING]{Colors.ENDC} {msg}")


def print_error(msg):
    print(f"{Colors.FAIL}[ERROR]{Colors.ENDC} {msg}")


try:
    import argostranslate.package
    import argostranslate.translate
except ImportError:
    print_error("The 'argostranslate' module is not installed.")
    print_info("Please install it using: pip install argostranslate")
    sys.exit(1)


def ensure_package_installed(from_code, to_code):
    """Ensure the translation package from `from_code` to `to_code` is installed.

    If not installed, attempts to download and install it automatically.
    """
    installed = argostranslate.translate.get_installed_languages()
    installed_dict = {lang.code: lang for lang in installed}

    from_lang = installed_dict.get(from_code)
    to_lang = installed_dict.get(to_code)

    if not from_lang or not to_lang or to_lang not in from_lang.translations_from:
        print_warning(
            f"Translation package {from_code} -> {to_code} not found. Attempting to install..."
        )
        try:
            argostranslate.package.update_package_index()
            available_packages = argostranslate.package.get_available_packages()

            pkg_to_install = None
            for pkg in available_packages:
                if pkg.from_code == from_code and pkg.to_code == to_code:
                    pkg_to_install = pkg
                    break

            if pkg_to_install:
                print_info(f"Downloading package: {pkg_to_install}")
                argostranslate.package.install_from_path(pkg_to_install.download())
                print_success(
                    f"Successfully installed package: {from_code} -> {to_code}"
                )

                # Refresh installed languages
                installed = argostranslate.translate.get_installed_languages()
                installed_dict = {lang.code: lang for lang in installed}
            else:
                print_error(
                    f"Could not find a translation package for {from_code} -> {to_code}"
                )
                sys.exit(1)
        except Exception as e:
            print_error(f"Failed to install language package: {e}")
            sys.exit(1)

    return installed_dict.get(from_code), installed_dict.get(to_code)


def get_translation_func(from_code, to_code):
    """Returns a translation function for the specified language pair."""
    from_lang, to_lang = ensure_package_installed(from_code, to_code)
    translator = from_lang.get_translation(to_lang)
    if not translator:
        print_error(f"No translation available from {from_code} to {to_code}")
        sys.exit(1)
    return translator.translate


def replace_vars_with_tokens(text):
    """Replaces `{variable}` patterns with a standard token like `XVAR0X`.

    So the translation engine doesn't attempt to translate variable names.
    Returns the modified text and the list of found variables.
    """
    vars_found = []

    def replacer(match):
        vars_found.append(match.group(0))
        return f" XVAR{len(vars_found) - 1}X "

    replaced_text = re.sub(r"\{+.*?\}+", replacer, text)
    return replaced_text, vars_found


def restore_vars_from_tokens(text, vars_found):
    """Restores the original `{variable}` patterns back into the translated text.

    Looks for the `XVAR0X` tokens.
    """
    for i, var in enumerate(vars_found):
        # The translation engine might change case or spacing around the token
        text = re.sub(rf"\s*[xX][vV][aA][rR]{i}[xX]\s*", var, text)
    return text.strip()


def translate_dict(data, translate_func, target_name=None):
    """Recursively iterates over a dictionary and translates all string values.

    Skips the `_languageName` key, which can be explicitly set.
    """
    if isinstance(data, dict):
        new_dict = {}
        for k, v in data.items():
            if k == "_languageName" and target_name:
                new_dict[k] = target_name
                continue
            if k == "_languageName":
                # Keep original if no target name provided
                new_dict[k] = v
                continue

            new_dict[k] = translate_dict(v, translate_func, target_name)
        return new_dict
    if isinstance(data, list):
        return [translate_dict(item, translate_func, target_name) for item in data]
    if isinstance(data, str):
        if not data.strip():
            return data

        temp_text, vars_found = replace_vars_with_tokens(data)
        try:
            translated_temp = translate_func(temp_text)
            return restore_vars_from_tokens(translated_temp, vars_found)
        except Exception as e:
            print_warning(
                f"Failed to translate '{data}': {e}. Falling back to original."
            )
            return data
    else:
        return data


def main():
    parser = argparse.ArgumentParser(
        description="Translate JSON localization files using Argos Translate."
    )
    parser.add_argument(
        "--from", dest="from_lang", help="Source language code (e.g. 'en')"
    )
    parser.add_argument("--to", dest="to_lang", help="Target language code (e.g. 'zh')")
    parser.add_argument("--input", dest="input_file", help="Path to input JSON file")
    parser.add_argument("--output", dest="output_file", help="Path to output JSON file")
    parser.add_argument(
        "--name",
        dest="target_name",
        help="Native name of the target language (e.g. '中文' for Chinese)",
    )

    args = parser.parse_args()

    # Fallback to Environment Variables if arguments aren't provided
    from_lang = args.from_lang or os.environ.get("ARGOS_FROM_LANG")
    to_lang = args.to_lang or os.environ.get("ARGOS_TO_LANG")
    input_file = args.input_file or os.environ.get("ARGOS_INPUT_FILE")
    output_file = args.output_file or os.environ.get("ARGOS_OUTPUT_FILE")
    target_name = args.target_name or os.environ.get("ARGOS_TARGET_NAME")

    if not all([from_lang, to_lang, input_file, output_file]):
        parser.print_help()
        print_error(
            "Missing required arguments. Please provide --from, --to, --input, and --output."
        )
        sys.exit(1)

    if not os.path.exists(input_file):
        print_error(f"Input file not found: {input_file}")
        sys.exit(1)

    print_info(f"Source Language: {Colors.BOLD}{from_lang}{Colors.ENDC}")
    print_info(f"Target Language: {Colors.BOLD}{to_lang}{Colors.ENDC}")
    print_info(f"Input File:      {Colors.BOLD}{input_file}{Colors.ENDC}")
    print_info(f"Output File:     {Colors.BOLD}{output_file}{Colors.ENDC}")

    # Load JSON
    try:
        with open(input_file, encoding="utf-8") as f:
            source_data = json.load(f)
    except json.JSONDecodeError as e:
        print_error(f"Invalid JSON in input file: {e}")
        sys.exit(1)

    start_time = time.time()

    # Get Translator
    translate_func = get_translation_func(from_lang, to_lang)

    print_info(
        "Starting translation. This may take a moment depending on the file size..."
    )
    translated_data = translate_dict(source_data, translate_func, target_name)

    # Ensure output directory exists
    os.makedirs(os.path.dirname(os.path.abspath(output_file)), exist_ok=True)

    # Save translated JSON
    try:
        with open(output_file, "w", encoding="utf-8") as f:
            json.dump(translated_data, f, ensure_ascii=False, indent=4)
            f.write("\n")
    except OSError as e:
        print_error(f"Could not write to output file: {e}")
        sys.exit(1)

    elapsed_time = time.time() - start_time
    print_success(f"Translation complete in {elapsed_time:.2f} seconds!")
    print_success(f"Output saved to {Colors.BOLD}{output_file}{Colors.ENDC}")


if __name__ == "__main__":
    main()