translate

#!/usr/bin/env python3
"""
Universal subtitle translator using OPUS-MT (EN→PT only) or NLLB-200/M2M100.

Enhancements:
──────────────
• Only asks once for destination language and model per source language
• Caches selections for full-season batches
• EN→PT offers OPUS-MT, NLLB-1.3B, M2M100-1.2B or NLLB-3.3B
• All other pairs use NLLB-1.3B, M2M100-1.2B or NLLB-3.3B
"""

import os
import re
import time
import torch
from transformers import (
    AutoTokenizer,
    AutoModelForSeq2SeqLM,
    MarianMTModel,
    MarianTokenizer,
)

# ---------- Helpers ----------

def list_srt_files():
    files = [f for f in sorted(os.listdir()) if f.lower().endswith(".srt")]
    for i, f in enumerate(files, 1):
        print(f"{i}: {f}")
    return files


def get_user_selection(files):
    sel = input("Enter numbers of files to translate (comma-separated), or press Enter for all: ").strip()
    if not sel:
        return files
    idxs = [int(i) - 1 for i in sel.split(",") if i.strip().isdigit() and 0 < int(i) <= len(files)]
    return [files[i] for i in idxs]


def detect_lang_from_filename(filename):
    match = re.findall(r"\.([a-z]{2})(?:-[a-z]{2})?\.srt$", filename.lower())
    return match[0] if match else None


def strip_all_lang_suffixes(base):
    return re.sub(r'(\.[a-z]{2}(-[a-z]{2})?)+$', '', base, flags=re.I)

def merge_lines_if_needed(lines):
    """
    Merge lines according to the rule:
    If line[i] does NOT end with punctuation and next line does NOT start with '-',
    merge them into one line.
    """
    merged = []
    i = 0
    
    while i < len(lines):
        line = lines[i].strip()
        
        # last line → nothing to merge
        if i == len(lines) - 1:
            merged.append(line)
            break
        
        next_line = lines[i + 1].strip()

        # condition for merging:
        if (not re.search(r"[.!?…]$", line)) and (not next_line.startswith("-")):
            merged_line = line + " " + next_line
            merged.append(merged_line)
            i += 2  # skip next line
        else:
            merged.append(line)
            i += 1

    return merged


# ---------- Translation ----------

def translate_batch(texts, tokenizer, model, device, is_marian=False, forced_bos_token_id=None):
    inputs = tokenizer(texts, return_tensors="pt", padding=True, truncation=True, max_length=256).to(device)

    outputs = model.generate(
        **inputs,
        forced_bos_token_id=forced_bos_token_id,  # NLLB or M2M100
        num_beams=4,
        max_length=256,
        no_repeat_ngram_size=3
    )

    return [tokenizer.decode(o, skip_special_tokens=True) for o in outputs]


def translate_srt(file_path, tokenizer, model, device, src_short, tgt_short, is_marian, model_name, forced_bos):
    start = time.time()
    base, _ = os.path.splitext(file_path)
    base_no_lang = strip_all_lang_suffixes(base)
    out_path = f"{base_no_lang}.{tgt_short}.ai.srt"

    with open(file_path, "r", encoding="utf-8") as f:
        content = f.read().strip()

    blocks = re.split(r"\n\n+", content)

    # Each entry: (idx, timing, list_of_lines)
    metas = []
    all_lines = []    # flat list of ALL subtitle lines from all blocks

    for block in blocks:
        parts = block.strip().splitlines()
        if len(parts) < 3:
            continue

        idx = parts[0]
        timing = parts[1]
        raw_lines = parts[2:]
        lines = merge_lines_if_needed(raw_lines)

        metas.append((idx, timing, len(lines)))
        all_lines.extend(lines)

    print(f"🌍 Using {model_name}  [{len(all_lines)} total lines]")

    # --- Translate all lines batch-wise ---
    out_lines = []
    batch_size = 16

    for i in range(0, len(all_lines), batch_size):
        batch = all_lines[i:i+batch_size]
        out_lines.extend(translate_batch(
            batch, tokenizer, model, device,
            is_marian=is_marian,
            forced_bos_token_id=forced_bos
        ))
        pct = int((i + len(batch)) / max(1, len(all_lines)) * 100)
        print(f"⏳ {pct}% complete...", end="\r")

    # --- Rebuild structure ---
    final_blocks = []
    cursor = 0

    for (idx, timing, line_count) in metas:
        translated_block_lines = out_lines[cursor: cursor + line_count]
        cursor += line_count

        # Keep exact original structure
        block_text = "\n".join(translated_block_lines)
        final_blocks.append(f"{idx}\n{timing}\n{block_text}")

    with open(out_path, "w", encoding="utf-8") as f:
        f.write("\n\n".join(final_blocks))

    print(f"\n\n✅ Saved translated subtitles to: {out_path}")
    print(f"⏱ Translation completed in {time.time() - start:.1f}s\n")


# ---------------------- MAIN ----------------------

def main():
    files = list_srt_files()
    if not files:
        print("No .srt files found.")
        return

    selected = get_user_selection(files)
    device = "cuda" if torch.cuda.is_available() else "cpu"
    print(f"\n💻 Using device: {device.upper()}")

    previous_settings = {}

    for path in selected:
        print(f"\n🎬 Processing: {path}")

        src_short = detect_lang_from_filename(path)
        if src_short:
            print(f"🌍 Detected source: {src_short.upper()}")
        else:
            src_short = input("Enter source language (2-letter, e.g., en, pt, fr): ").strip().lower() or "en"

        if src_short in previous_settings:
            tgt_short, is_marian, model_name, tokenizer, model, forced_bos = previous_settings[src_short]
            print(f"🔁 Reusing settings for {src_short.upper()} → {tgt_short.upper()} ({model_name})")

        else:
            tgt_short = input("Enter destination language (2-letter, e.g., en, pt, fr): ").strip().lower() or "en"

            if src_short == tgt_short:
                print("⚠️  Source and destination are the same. Skipping.")
                continue

            # Default model
            is_marian = False
            model_name = "facebook/nllb-200-1.3B"

            # Present multilingual engine choices for all language pairs.
            # OPUS-MT (Marian) is only offered for the en→pt combination.
            print("\nChoose translation engine:")
            print("1) NLLB-200 1.3B (Multilingual, large, default)")
            print("2) NLLB-200 3.3B (Multilingual, very large)")
            print("3) M2M100 1.2B (Multilingual, medium)")

            opus_option = (src_short, tgt_short) == ("en", "pt")
            if opus_option:
                print("4) OPUS-MT (European Portuguese, smaller, faster)")

            valid_choices = ["1", "2", "3"] + (["4"] if opus_option else [])

            choice = input("Enter number [1]: ").strip() or "1"
            if choice not in valid_choices:
                print("Invalid choice, defaulting to 1")
                choice = "1"

            if choice == "1":
                model_name = "facebook/nllb-200-1.3B"
                is_marian = False

            elif choice == "2":
                model_name = "facebook/nllb-200-3.3B"
                is_marian = False

            elif choice == "3":
                model_name = "facebook/m2m100_1.2B"
                is_marian = False

            elif choice == "4" and opus_option:
                model_name = "Helsinki-NLP/opus-mt-tc-big-en-pt"
                is_marian = True

            print(f"\n🔁 Loading model: {model_name}")

            # --- Load Marian (OPUS-MT) ---
            if is_marian:
                tokenizer = MarianTokenizer.from_pretrained(model_name)
                model = MarianMTModel.from_pretrained(model_name).to(device)
                forced_bos = None

            else:
                tokenizer = AutoTokenizer.from_pretrained(model_name)
                model = AutoModelForSeq2SeqLM.from_pretrained(
                    model_name,
                    dtype=torch.float16 if device == "cuda" else torch.float32
                ).to(device)

                # --------- M2M100 SPECIFIC FIX ---------
                if "m2m100" in model_name.lower():
                    tokenizer.src_lang = src_short
                    forced_bos = tokenizer.get_lang_id(tgt_short)
                else:
                    # --------- NLLB LANGUAGE CODES ---------
                    NLLB_MAP = {
                        "af":"afr_Latn","am":"amh_Ethi","ar":"arb_Arab","az":"azj_Latn","be":"bel_Cyrl","bg":"bul_Cyrl",
                        "bn":"ben_Beng","bs":"bos_Latn","ca":"cat_Latn","cs":"ces_Latn","cy":"cym_Latn","da":"dan_Latn",
                        "de":"deu_Latn","el":"ell_Grek","en":"eng_Latn","es":"spa_Latn","et":"est_Latn","eu":"eus_Latn",
                        "fa":"pes_Arab","fi":"fin_Latn","fr":"fra_Latn","gl":"glg_Latn","gu":"guj_Gujr","he":"heb_Hebr",
                        "hi":"hin_Deva","hr":"hrv_Latn","hu":"hun_Latn","hy":"hye_Armn","id":"ind_Latn","is":"isl_Latn",
                        "it":"ita_Latn","ja":"jpn_Jpan","jv":"jav_Latn","ka":"kat_Geor","kk":"kaz_Cyrl","km":"khm_Khmr",
                        "kn":"kan_Knda","ko":"kor_Hang","lo":"lao_Laoo","lt":"lit_Latn","lv":"lvs_Latn","mk":"mkd_Cyrl",
                        "ml":"mal_Mlym","mn":"khk_Cyrl","mr":"mar_Deva","ms":"zsm_Latn","my":"mya_Mymr","nl":"nld_Latn",
                        "no":"nob_Latn","or":"ory_Orya","pa":"pan_Guru","pl":"pol_Latn","ps":"pbt_Arab","pt":"por_Latn",
                        "ro":"ron_Latn","ru":"rus_Cyrl","sd":"snd_Arab","si":"sin_Sinh","sk":"slk_Latn","sl":"slv_Latn",
                        "so":"som_Latn","sq":"als_Latn","sr":"srp_Cyrl","sv":"swe_Latn","sw":"swh_Latn","ta":"tam_Taml",
                        "te":"tel_Telu","th":"tha_Thai","tl":"tgl_Latn","tr":"tur_Latn","uk":"ukr_Cyrl","ur":"urd_Arab",
                        "uz":"uzn_Latn","vi":"vie_Latn","yo":"yor_Latn","zh":"zho_Hans","zh-tw":"zho_Hant","zu":"zul_Latn",
                    }

                    tokenizer.src_lang = NLLB_MAP.get(src_short, "eng_Latn")
                    tokenizer.tgt_lang = NLLB_MAP.get(tgt_short, "por_Latn")
                    forced_bos = tokenizer.convert_tokens_to_ids(tokenizer.tgt_lang)

            previous_settings[src_short] = (tgt_short, is_marian, model_name, tokenizer, model, forced_bos)

        translate_srt(path, tokenizer, model, device, src_short, tgt_short, is_marian, model_name, forced_bos)


if __name__ == "__main__":
    main()