#!/usr/bin/env python3
"""
Translate missing keys in Verso locale files using Google Translate (free, no API key).

Usage:
    python3 translate_missing.py [--lang fr,de,it,es] [--batch 80] [--dry-run]

Re-run safely: already-translated keys are skipped.
"""

import json
import re
import sys
import time
import argparse
from pathlib import Path
from deep_translator import GoogleTranslator

LOCALES_DIR = Path(__file__).parent
EN_FILE = LOCALES_DIR / "en.json"

LANG_NAMES = {
    "fr": "French",
    "de": "German",
    "it": "Italian",
    "es": "Spanish",
}

# Placeholder pattern: __varName__ and <0>…</N> tags
# We replace them with stable tokens before translating and restore after.
INTERP_RE = re.compile(r'__\w+__|<\d+>.*?</\d+>|<\d+\s*/>', re.DOTALL)


def tokenise(text: str) -> tuple[str, list[str]]:
    """Replace variable/tag fragments with @@0@@, @@1@@, … Return (masked, originals)."""
    originals: list[str] = []
    def replace(m):
        idx = len(originals)
        originals.append(m.group(0))
        return f"@@{idx}@@"
    return INTERP_RE.sub(replace, text), originals


def restore(text: str, originals: list[str]) -> str:
    for i, orig in enumerate(originals):
        text = text.replace(f"@@{i}@@", orig)
    return text


SEP = "\n||||\n"
MAX_CHARS = 4500  # stay well under Google's 5000-char limit


def _call_google(text: str, target: str) -> str:
    """Single Google Translate call with one retry on transient errors."""
    try:
        return GoogleTranslator(source="en", target=target).translate(text)
    except Exception as e:
        print(f"\n    [retry] {type(e).__name__} — waiting 6s…", file=sys.stderr)
        time.sleep(6)
        return GoogleTranslator(source="en", target=target).translate(text)


def _translate_chunk(masked: list[str], orig_lists: list[list[str]], target: str) -> list[str]:
    """Translate one chunk (guaranteed to fit within MAX_CHARS)."""
    if len(masked) == 1:
        translated_text = _call_google(masked[0], target)
        return [restore(translated_text.strip(), orig_lists[0])]

    joined = SEP.join(masked)
    translated_joined = _call_google(joined, target)
    parts = translated_joined.split("||||")

    if len(parts) == len(masked):
        return [restore(p.strip(), origs) for p, origs in zip(parts, orig_lists)]

    # Split-count mismatch: fall back to individual calls
    print(f"\n    [warn] split mismatch ({len(parts)} vs {len(masked)}), going one-by-one", file=sys.stderr)
    results = []
    for m, origs in zip(masked, orig_lists):
        try:
            t = _call_google(m, target)
            results.append(restore(t.strip(), origs))
            time.sleep(0.4)
        except Exception as e:
            print(f"    [skip] could not translate: {e}", file=sys.stderr)
            results.append(restore(m, origs))  # fall back to English
    return results


def translate_batch(texts: list[str], target: str) -> list[str]:
    """Translate a list of strings, auto-splitting to respect MAX_CHARS."""
    masked_all = []
    orig_lists_all = []
    for t in texts:
        m, origs = tokenise(t)
        masked_all.append(m)
        orig_lists_all.append(origs)

    # Group into sub-chunks that fit within MAX_CHARS
    results: list[str] = []
    chunk_m: list[str] = []
    chunk_o: list[list[str]] = []
    chunk_chars = 0

    for m, origs in zip(masked_all, orig_lists_all):
        entry_chars = len(m) + len(SEP)
        if chunk_m and chunk_chars + entry_chars > MAX_CHARS:
            results.extend(_translate_chunk(chunk_m, chunk_o, target))
            chunk_m, chunk_o, chunk_chars = [], [], 0
            time.sleep(0.5)
        chunk_m.append(m)
        chunk_o.append(origs)
        chunk_chars += entry_chars

    if chunk_m:
        results.extend(_translate_chunk(chunk_m, chunk_o, target))

    return results


def process_language(lang: str, en: dict, batch_size: int, dry_run: bool) -> int:
    """Fill missing keys for one language. Returns number of keys translated."""
    target_file = LOCALES_DIR / f"{lang}.json"
    existing: dict = {}
    if target_file.exists():
        with open(target_file, encoding="utf-8") as f:
            existing = json.load(f)

    missing_keys = [k for k in en if k not in existing or not existing[k]]
    if not missing_keys:
        print(f"  {lang}: nothing to translate.")
        return 0

    print(f"  {lang} ({LANG_NAMES[lang]}): {len(missing_keys)} keys to translate")
    translated_count = 0

    for batch_start in range(0, len(missing_keys), batch_size):
        batch_keys = missing_keys[batch_start:batch_start + batch_size]
        batch_values = [en[k] for k in batch_keys]

        end = min(batch_start + batch_size, len(missing_keys))
        print(f"    batch {batch_start + 1}–{end} / {len(missing_keys)}…", end=" ", flush=True)

        if dry_run:
            print("(dry-run, skipped)")
            continue

        translations = translate_batch(batch_values, lang)

        for key, translation in zip(batch_keys, translations):
            existing[key] = translation
        translated_count += len(batch_keys)

        # Write after every batch so progress is saved if interrupted
        sorted_existing = dict(sorted(existing.items()))
        with open(target_file, "w", encoding="utf-8") as f:
            json.dump(sorted_existing, f, ensure_ascii=False, indent=2)
            f.write("\n")

        print(f"done, saved.")
        # Polite pause to avoid rate-limiting
        time.sleep(1.2)

    return translated_count


def main():
    parser = argparse.ArgumentParser(description="Translate missing Verso locale keys.")
    parser.add_argument("--lang", default="fr,de,it,es",
                        help="Comma-separated language codes (default: fr,de,it,es)")
    parser.add_argument("--batch", type=int, default=80,
                        help="Keys per Google Translate call (default: 80)")
    parser.add_argument("--dry-run", action="store_true",
                        help="Report what would be translated without writing files")
    args = parser.parse_args()

    langs = [l.strip() for l in args.lang.split(",")]

    with open(EN_FILE, encoding="utf-8") as f:
        en = json.load(f)
    print(f"Source: en.json — {len(en)} keys\n")

    total = 0
    for lang in langs:
        if lang not in LANG_NAMES:
            print(f"  [skip] unknown language: {lang}")
            continue
        total += process_language(lang, en, args.batch, args.dry_run)
        print()

    print(f"Done. {total} keys translated across {len(langs)} language(s).")


if __name__ == "__main__":
    main()