office_translator/translators/segments.py

"""
Segment recording & overrides — the review foundation.

During a translation job, every (source → translation) pair is recorded so
it can be persisted per job (translation_segments table), reviewed and
edited side-by-side, and re-applied when rebuilding the document.

Overrides are the inverse path: human-reviewed translations (approved or
edited segments) applied verbatim on rebuild — no provider call, no drift.
"""

from typing import Dict, List, Optional, Tuple

from core.logging import get_logger

logger = get_logger(__name__)


class SegmentRecorder:
    """Collects unique (source, translation) pairs in document order."""

    def __init__(self):
        self._pairs: List[Tuple[str, str]] = []
        self._seen: set = set()

    def record_pair(self, source: str, translation: str) -> None:
        if not source or not source.strip():
            return
        stripped = source.strip()
        # Identity pairs carry no review value ( untranslated or unchanged
        # text) — skip them so the review list stays actionable.
        if not translation or not translation.strip() or translation.strip() == stripped:
            return
        if stripped in self._seen:
            return
        self._seen.add(stripped)
        self._pairs.append((stripped, translation.strip()))

    def record_pairs(self, pairs) -> None:
        for source, translation in pairs:
            self.record_pair(source, translation)

    def get_pairs(self) -> List[Tuple[str, str]]:
        return list(self._pairs)

    def __len__(self) -> int:
        return len(self._pairs)


def apply_overrides(
    texts: List[str],
    overrides: Optional[Dict[str, str]],
) -> Tuple[Dict[int, str], List[int]]:
    """Split a batch: ({index: override_translation}, [indices to translate]).

    Overrides are keyed by the stripped source text (that is how segments
    are persisted). Identity overrides are treated as "no override" so the
    normal pipeline (TM → provider) stays in charge.
    """
    if not overrides:
        return {}, list(range(len(texts)))

    hits: Dict[int, str] = {}
    misses: List[int] = []
    for i, text in enumerate(texts):
        if not text or not text.strip():
            misses.append(i)
            continue
        override = overrides.get(text.strip())
        if override is not None and override.strip() and override.strip() != text.strip():
            hits[i] = override.strip()
        else:
            misses.append(i)
    return hits, misses


# ---------------------------------------------------------------------------
# Wrong-script retry — turn the script detector into an actionable guard
# ---------------------------------------------------------------------------

def _script_retry_hint(target_language: str, reason: str) -> str:
    """Extra instructions appended to the retry request's custom prompt.

    Tells the model, in its own working language (English), that the
    previous attempt was rejected for the wrong script/variant and what
    the output must look like this time.
    """
    from core.languages import language_name
    from services.quality import config as _qa_config

    name = language_name(target_language) or target_language
    lines = [
        "CRITICAL RETRY: your previous attempt was rejected because the "
        f"output was not written in {name}.",
        f"Rejected output reason: {reason}",
        f"Write the translation strictly in {name}.",
    ]
    target_chars = _qa_config.get_discriminating_chars(
        (target_language or "").lower().split("-")[0]
    )
    if target_chars:
        sample = " ".join(sorted(target_chars))
        lines.append(
            f"It must be {name}, not plain Arabic: use the letters specific "
            f"to {name} ({sample}) wherever the words require them."
        )
    return " ".join(lines)


def retry_wrong_script(
    sources: List[str],
    translations: List[str],
    target_language: str,
    translate_fn,
) -> Tuple[List[str], int]:
    """Retry the translations delivered in the wrong script, once each.

    The most damaging silent failure of a translation pipeline is the
    right language family with the wrong alphabet: Arabic delivered for
    Persian looks almost correct to a non-reader and ships unnoticed.
    This guard checks every delivered translation against the target
    language's expected script (and, for Arabic-script languages, the
    expected variant), and re-asks the model once for each offender with
    an explicit instruction. A retry only replaces the first attempt when
    it passes the same check; otherwise the first attempt is kept.

    translate_fn(texts, extra_prompt) must translate a list of texts in
    one call, appending extra_prompt to the request's custom prompt.

    Returns (translations, number_of_replacements). Never raises: any
    error in detection or retry keeps the original translations.
    """
    if not translations or not target_language:
        return translations, 0

    try:
        from services.quality.script_detector import script_issue
    except Exception:
        return translations, 0

    # Only genuinely translated segments can be "wrong": identity output
    # (numbers, proper names kept as-is) must not trigger a retry.
    offenders = []
    reasons = []
    for i, (src, trans) in enumerate(zip(sources, translations)):
        if not trans or trans.strip() == (src or "").strip():
            continue
        try:
            reason = script_issue(trans, target_language)
        except Exception:
            continue
        if reason:
            offenders.append(i)
            reasons.append(reason)

    if not offenders:
        return translations, 0

    hint = _script_retry_hint(target_language, reasons[0])
    try:
        retried = translate_fn([sources[i] for i in offenders], hint)
    except Exception as e:
        logger.warning("script_retry_failed", error=str(e)[:200], count=len(offenders))
        return translations, 0

    fixed = 0
    result = list(translations)
    for pos, i in enumerate(offenders):
        if pos >= len(retried):
            break
        candidate = retried[pos]
        if not candidate or not candidate.strip():
            continue
        try:
            still_wrong = script_issue(candidate, target_language)
        except Exception:
            still_wrong = True  # do not swap in an unverifiable retry
        if not still_wrong:
            result[i] = candidate
            fixed += 1

    logger.info(
        "script_retry",
        target=target_language,
        offenders=len(offenders),
        fixed=fixed,
    )
    return result, fixed