services/translation.py

"""Translation service: Transcript -> Translation (text only).

Per the EchoScript v1.0 architecture decision, translation is ALWAYS
derived from the canonical Transcript's text, never from the original
audio.

Two interchangeable backends, chosen per-request based on whether an
Anthropic API key is supplied -- never stored:

- "anthropic": uses Claude for translation via the supplied key. No
  missing-language-pair failure mode. Key is used only for this call.
- "marian": fully offline, no API key needed. Uses Helsinki-NLP MarianMT
  models. Available targets are determined from KNOWN_MARIAN_PAIRS below.

IMPORTANT -- Helsinki-NLP model naming quirks addressed here:
  - Some pairs use multilingual group models (e.g. "en-iir" for all
    Indo-Iranian languages) rather than a direct bilingual model.
  - Those multilingual models require a >>langcode<< token prepended to
    every source segment so the model knows which target language to use.
  - The repo suffix and the optional prefix token are stored together in
    KNOWN_MARIAN_PAIRS so the translation code can apply them correctly.
  - Japanese uses repo code "jap" not "ja" for the en->ja direction.
  - Korean (en->ko) has no confirmed Helsinki-NLP model; it is omitted
    from the Marian pair table so it only appears when an API key exists.
"""

from __future__ import annotations

import os
import re
from functools import lru_cache
from typing import Optional

from models.transcript import Segment, Transcript, Translation

# ---------------------------------------------------------------------------
# Language catalog
# ---------------------------------------------------------------------------

LANGUAGE_NAMES: dict[str, str] = {
    "fr": "French",
    "en": "English",
    "de": "German",
    "fa": "Persian",
    "es": "Spanish",
    "it": "Italian",
    "pt": "Portuguese",
    "nl": "Dutch",
    "ar": "Arabic",
    "ru": "Russian",
    "tr": "Turkish",
    "ja": "Japanese",
    "zh": "Chinese",
    "ko": "Korean",
}

ANTHROPIC_TARGET_LANGUAGES: dict[str, str] = dict(LANGUAGE_NAMES)

# ---------------------------------------------------------------------------
# Verified Marian pair table
#
# Format: (source_code, target_code) -> (repo_suffix, prefix_token)
#
# repo_suffix: appended to "Helsinki-NLP/opus-mt-" to form the model name.
# prefix_token: prepended to every source segment for multilingual models
#   that need a >>langcode<< token to select the target language, or None
#   for standard bilingual models.
#
# Verified against the Helsinki-NLP catalog. Key findings:
#   - opus-mt-en-fa does NOT exist. Use opus-mt-en-iir (Indo-Iranian
#     group) with >>pes<< (Western Persian, ISO 639-3).
#   - opus-mt-en-pt does NOT exist. Use opus-mt-en-roa (Romance group)
#     with >>por<< (Portuguese, ISO 639-3).
#   - opus-mt-en-tr does NOT exist. Use opus-mt-en-trk (Turkic group)
#     with >>tur<< (Turkish, ISO 639-3).
#   - opus-mt-en-ko has no confirmed model; Korean removed from en->X.
#   - opus-mt-en-jap (note: "jap" not "ja") for English->Japanese.
#   - fa-en uses direct opus-mt-fa-en (not the group model).
#   - pt-en uses opus-mt-roa-en (Romance group; pt-en not on Hub).
# ---------------------------------------------------------------------------
KNOWN_MARIAN_PAIRS: dict[tuple[str, str], tuple[str, Optional[str]]] = {
    # English -> X
    ("en", "fr"): ("en-fr",  None),
    ("en", "de"): ("en-de",  None),
    ("en", "fa"): ("en-iir", ">>pes<<"),  # Indo-Iranian group, Western Persian
    ("en", "es"): ("en-es",  None),
    ("en", "it"): ("en-it",  None),
    ("en", "pt"): ("en-roa", ">>por<<"),  # Romance group, Portuguese
    ("en", "nl"): ("en-nl",  None),
    ("en", "ar"): ("en-ar",  None),
    ("en", "ru"): ("en-ru",  None),
    ("en", "tr"): ("en-trk", ">>tur<<"),  # Turkic group, Turkish
    ("en", "zh"): ("en-zh",  None),
    ("en", "ja"): ("en-jap", None),       # Note: "jap" not "ja"
    # en->ko omitted: no confirmed Helsinki-NLP model

    # X -> English
    ("fr", "en"): ("fr-en", None),
    ("de", "en"): ("de-en", None),
    ("fa", "en"): ("fa-en", None),
    ("es", "en"): ("es-en", None),
    ("it", "en"): ("it-en", None),
    ("pt", "en"): ("roa-en", None),   # Romance group -> English
    ("nl", "en"): ("nl-en", None),
    ("ar", "en"): ("ar-en", None),
    ("ru", "en"): ("ru-en", None),
    ("tr", "en"): ("tr-en", None),
    ("zh", "en"): ("zh-en", None),
    ("ja", "en"): ("ja-en", None),
    ("ko", "en"): ("ko-en", None),

    # Selected direct non-English pairs (avoids double pivot hop)
    ("fr", "de"): ("fr-de", None),
    ("de", "fr"): ("de-fr", None),
    ("fr", "es"): ("fr-es", None),
    ("es", "fr"): ("es-fr", None),
    ("de", "es"): ("de-es", None),
    ("es", "de"): ("es-de", None),
}


def _marian_entry(src: str, tgt: str) -> Optional[tuple[str, Optional[str]]]:
    return KNOWN_MARIAN_PAIRS.get((src, tgt))


def _marian_path(src: str, tgt: str) -> Optional[list[tuple[str, Optional[str]]]]:
    """Return the list of (repo_suffix, prefix_token) steps for src->tgt.

    Returns a 1-step list for a direct (or multilingual-model) pair, a
    2-step list for an English pivot, or None if no path is known.
    """
    if src == tgt:
        return None
    entry = _marian_entry(src, tgt)
    if entry:
        return [entry]
    if src != "en" and tgt != "en":
        hop1 = _marian_entry(src, "en")
        hop2 = _marian_entry("en", tgt)
        if hop1 and hop2:
            return [hop1, hop2]
    return None


def available_marian_targets(source_language: str) -> dict[str, str]:
    """Every language MarianMT can reach from `source_language`.

    Based on the verified KNOWN_MARIAN_PAIRS table. No network calls.
    """
    return {
        code: name
        for code, name in LANGUAGE_NAMES.items()
        if code != source_language and _marian_path(source_language, code) is not None
    }


# ---------------------------------------------------------------------------
# Anthropic backend
# ---------------------------------------------------------------------------

_ANTHROPIC_MODEL = "claude-haiku-4-5-20251001"
_ANTHROPIC_BATCH_SIZE = 40
_NUMBERED_LINE_RE = re.compile(r"^\s*(\d+)[.\)]\s?(.*)$")


class TranslationError(RuntimeError):
    pass


def _anthropic_client(api_key: str):
    import anthropic
    return anthropic.Anthropic(api_key=api_key)


def _translate_batch_via_anthropic(
    texts: list[str], source_language: str, target_language: str, api_key: str
) -> list[str]:
    source_name = LANGUAGE_NAMES.get(source_language, source_language)
    target_name = LANGUAGE_NAMES.get(target_language, target_language)
    numbered_input = "\n".join(f"{i + 1}. {text}" for i, text in enumerate(texts))
    try:
        response = _anthropic_client(api_key).messages.create(
            model=_ANTHROPIC_MODEL,
            max_tokens=4096,
            system=(
                f"You translate transcript lines from {source_name} to {target_name}. "
                "You will be given a numbered list of lines, one sentence or fragment "
                "per line. Reply with the same numbers, translated, one per line, in "
                "the same order. Keep the same number of lines as the input. "
                "Output only the numbered translated lines, no preamble or commentary."
            ),
            messages=[{"role": "user", "content": numbered_input}],
        )
    except Exception as exc:
        raise TranslationError(
            f"Anthropic request failed for '{source_language}'->'{target_language}': {exc}"
        ) from exc

    raw = "".join(
        b.text for b in response.content if getattr(b, "type", None) == "text"
    )
    parsed: dict[int, str] = {}
    for line in raw.splitlines():
        m = _NUMBERED_LINE_RE.match(line)
        if m:
            parsed[int(m.group(1))] = m.group(2).strip()

    if len(parsed) != len(texts) or any((i + 1) not in parsed for i in range(len(texts))):
        raise TranslationError(
            f"Anthropic response line count mismatch for "
            f"'{source_language}'->'{target_language}' "
            f"(expected {len(texts)}, got {len(parsed)})."
        )
    return [parsed[i + 1] for i in range(len(texts))]


def _translate_segments_via_anthropic(
    segments: list[Segment], src: str, tgt: str, api_key: str
) -> list[Segment]:
    non_empty = [(i, seg) for i, seg in enumerate(segments) if seg.text]
    by_index: dict[int, str] = {}
    for start in range(0, len(non_empty), _ANTHROPIC_BATCH_SIZE):
        chunk = non_empty[start: start + _ANTHROPIC_BATCH_SIZE]
        translated = _translate_batch_via_anthropic(
            [s.text for _, s in chunk], src, tgt, api_key
        )
        for (i, _), text in zip(chunk, translated):
            by_index[i] = text
    return [
        Segment(index=seg.index, start=seg.start, end=seg.end,
                text=by_index.get(i, seg.text))
        for i, seg in enumerate(segments)
    ]


# ---------------------------------------------------------------------------
# Marian (offline) backend
# ---------------------------------------------------------------------------


@lru_cache(maxsize=None)
def _load_marian_engine(repo_suffix: str):
    """Load and cache a MarianMT model+tokenizer by repo suffix."""
    from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
    model_name = f"Helsinki-NLP/opus-mt-{repo_suffix}"
    try:
        tokenizer = AutoTokenizer.from_pretrained(model_name)
        model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
    except Exception as exc:
        raise TranslationError(
            f"Failed to load MarianMT model '{model_name}': {exc}"
        ) from exc
    return tokenizer, model


def _run_marian(tokenizer, model, text: str, prefix: Optional[str]) -> str:
    """Translate one text segment, optionally prepending a language token."""
    source = f"{prefix} {text}" if prefix else text
    inputs = tokenizer(source, return_tensors="pt", truncation=True)
    generated = model.generate(**inputs, max_new_tokens=512)
    return tokenizer.decode(generated[0], skip_special_tokens=True).strip()


def _translate_segments_via_marian(
    segments: list[Segment], src: str, tgt: str
) -> list[Segment]:
    steps = _marian_path(src, tgt)
    if not steps:
        raise TranslationError(
            f"No Marian translation path known for '{src}'->'{tgt}'."
        )
    engines = [(_load_marian_engine(suffix), prefix) for suffix, prefix in steps]
    result = []
    for seg in segments:
        if not seg.text:
            result.append(seg)
            continue
        text = seg.text
        for (tokenizer, model), prefix in engines:
            text = _run_marian(tokenizer, model, text, prefix)
        result.append(Segment(index=seg.index, start=seg.start, end=seg.end, text=text))
    return result


# ---------------------------------------------------------------------------
# Public service
# ---------------------------------------------------------------------------


class TranslationService:
    def translate(
        self,
        transcript: Transcript,
        target_language: str,
        api_key: Optional[str] = None,
    ) -> Translation:
        if target_language == transcript.language:
            return Translation(
                source_language=transcript.language,
                target_language=target_language,
                segments=list(transcript.segments),
            )
        effective_key = (api_key or "").strip() or os.environ.get("ANTHROPIC_API_KEY")
        if effective_key:
            translated = _translate_segments_via_anthropic(
                transcript.segments, transcript.language, target_language, effective_key
            )
        else:
            translated = _translate_segments_via_marian(
                transcript.segments, transcript.language, target_language
            )
        return Translation(
            source_language=transcript.language,
            target_language=target_language,
            segments=translated,
        )