app/engine/orchestrator.py

"""Batch/stream orchestrator for the structural rewrite engine."""

from __future__ import annotations

import logging
import time
from difflib import SequenceMatcher

from app.config import (
    ENGINE_BATCH_PARAS,
    ENGINE_CLASSICAL_AGGRESSIVE,
    ENGINE_FORCE_REWRITE,
    ENGINE_LEXICAL_MIN_WSD,
    ENGINE_LEXICAL_REFINEMENT,
    ENGINE_MIN_CONFIDENCE,
    ENGINE_PARAPHRASE,
    ENGINE_PARAPHRASE_MIN_SIM,
    ENGINE_PARAPHRASE_PRIMARY,
    ENGINE_PHRASE_MAX_CHANGES,
    ENGINE_PHRASE_MIN_SIM,
    ENGINE_PHRASE_REWRITE,
    ENGINE_PHRASE_USE_T5,
    ENGINE_PRESERVE_LENGTH,
    ENGINE_REQUIRE_WORDING_CHANGE,
    ENGINE_SAFETY_MIN,
    ENGINE_SPLIT_LONG,
    ENGINE_SPLIT_MIN_WORDS,
    ENGINE_STRUCTURAL_VARIATION,
    ENGINE_USE_MINILM_SAFETY,
    GRAMMAR_FIX_OUTPUT,
)
from app.engine.classify import classify_sentence
from app.engine.consistency import apply_consistency
from app.engine.force import force_cleft_rewrite
from app.engine.grammar import repair_sentence
from app.engine.ingest import ingest_text
from app.engine.lexical import (
    dynamic_lexical_budget,
    ensure_wording_change,
    refine_sentence,
)
from app.engine.mechanics import enforce_length_budget, tidy
from app.engine.phrase import rewrite_phrases
from app.engine.models import (
    DocumentBlock,
    EngineResult,
    EngineStats,
    SentenceRecord,
)
from app.engine.normalize import (
    detect_blocks,
    mask_protected_spans,
    normalize_text,
    unmask_protected_spans,
)
from app.engine.fallback import structural_fallback_candidates
from app.engine.paraphrase import (
    paraphrase_sentence,
    sufficiently_changed,
    surface_similarity,
)
from app.engine.plan import build_plan
from app.engine.rewrite import generate_from_plan
from app.engine.safety import check_safety
from app.engine.segment import iter_paragraph_batches, split_sentences, word_count
from app.engine.stitch import join_sentences, stitch_blocks
from app.engine.variation import new_seed, rotate
from app.engine.voice import active_to_passive

logger = logging.getLogger("plainrewrite.engine")

_REWRITEABLE_TYPES = frozenset(
    {"simple_declarative", "compound", "complex", "because_clause"}
)


def _classical_modes(use_minilm: bool) -> tuple[bool, bool]:
    """Return (classical_strict, classical_aggressive) for wording paths."""
    if use_minilm:
        return False, False
    return True, bool(ENGINE_CLASSICAL_AGGRESSIVE)


def _similarity(a: str, b: str) -> float:
    return SequenceMatcher(None, (a or "").lower(), (b or "").lower()).ratio()


def _apply_lexical_refinement(
    record: SentenceRecord,
    *,
    enabled: bool,
    min_wsd: float,
    max_changes: int | None,
    safety_min: float,
    min_confidence: float,
    use_minilm: bool,
    polish: bool = False,
) -> SentenceRecord:
    if not enabled or record.sentence_type not in _REWRITEABLE_TYPES:
        return record
    before = record.rewritten
    classical_strict, classical_aggressive = _classical_modes(use_minilm)
    refined = refine_sentence(
        before,
        min_wsd=min_wsd,
        max_changes=max_changes,
        polish=polish,
        classical_strict=classical_strict,
        classical_aggressive=classical_aggressive,
    )
    if not refined.changes or refined.text == before:
        return record
    safety = check_safety(
        before,
        refined.text,
        min_meaning=safety_min,
        min_confidence=min_confidence,
        use_minilm=use_minilm,
        protected_entities=None,
        protected_auxiliaries=None,
        structural_validation=False,
    )
    if not safety.ok:
        return record

    prior_status = record.status
    prior_reasons = list(record.reasons)
    prior_lex = list(record.lexical_changes)
    record.rewritten = refined.text
    record.lexical_changes = prior_lex + refined.changes
    record.status = "rewritten"
    if not record.template_id:
        record.template_id = "lexical_refine"
    elif "+lexical" not in record.template_id:
        record.template_id = f"{record.template_id}+lexical"
    if prior_status in {"skipped", "reverted"}:
        record.reasons = [
            f"structural_{prior_status}:{reason}" for reason in prior_reasons
        ]
    record.confidence = min(
        safety.confidence,
        record.confidence if record.confidence > 0 else max(min_confidence, 0.55),
    )
    return record


def _apply_extra_polish(
    record: SentenceRecord,
    *,
    enabled: bool,
    min_wsd: float,
    max_changes: int | None,
    safety_min: float,
    min_confidence: float,
    use_minilm: bool,
) -> SentenceRecord:
    """Polish-only extra pass so the UI toggle has a real end-to-end effect."""
    if not enabled or record.sentence_type not in _REWRITEABLE_TYPES:
        return record
    classical_strict, classical_aggressive = _classical_modes(use_minilm)
    # Tight classical (no MiniLM, not aggressive): skip extra polish cycling.
    if classical_strict and not classical_aggressive:
        return record
    before = record.rewritten
    target_budget = (
        max(1, min(int(max_changes), 15))
        if max_changes is not None
        else dynamic_lexical_budget(before, polish=True)
    )
    remaining = max(0, target_budget - len(record.lexical_changes))
    if remaining <= 0:
        return record
    refined = refine_sentence(
        before,
        min_wsd=min(min_wsd, 0.10),
        max_changes=remaining,
        aggressive=True,
        classical_strict=classical_strict,
        classical_aggressive=classical_aggressive,
    )
    if not refined.changes or refined.text == before:
        return record
    safety = check_safety(
        before,
        refined.text,
        min_meaning=safety_min,
        min_confidence=min_confidence,
        use_minilm=use_minilm,
        protected_entities=None,
        protected_auxiliaries=None,
        structural_validation=False,
    )
    if not safety.ok:
        return record

    prior_lex = list(record.lexical_changes)
    record.rewritten = refined.text
    record.lexical_changes = prior_lex + refined.changes
    record.status = "rewritten"
    if not record.template_id:
        record.template_id = "lexical_refine"
    elif "+lexical" not in record.template_id:
        record.template_id = f"{record.template_id}+lexical"
    record.confidence = min(
        safety.confidence,
        record.confidence if record.confidence > 0 else max(min_confidence, 0.55),
    )
    return record


def _apply_phrase_rewrite(
    record: SentenceRecord,
    *,
    enabled: bool,
    polish: bool,
    safety_min: float,
    min_confidence: float,
    use_minilm: bool,
) -> SentenceRecord:
    """Rewrite verb–object phrases before single-word lexical polish."""
    if not enabled or record.sentence_type not in _REWRITEABLE_TYPES:
        return record
    before = record.rewritten
    max_changes = (
        ENGINE_PHRASE_MAX_CHANGES
        if ENGINE_PHRASE_MAX_CHANGES > 0
        else (2 if polish else 1)
    )
    classical_strict, classical_aggressive = _classical_modes(use_minilm)
    refined = rewrite_phrases(
        before,
        max_changes=max_changes,
        polish=polish,
        min_sim=min(ENGINE_PHRASE_MIN_SIM, safety_min),
        use_t5=ENGINE_PHRASE_USE_T5 and use_minilm,
        classical_strict=classical_strict,
        classical_aggressive=classical_aggressive,
    )
    if not refined.changes or refined.text == before:
        return record
    safety = check_safety(
        before,
        refined.text,
        min_meaning=min(safety_min, ENGINE_PHRASE_MIN_SIM),
        min_confidence=min_confidence,
        use_minilm=use_minilm,
        protected_entities=None,
        protected_auxiliaries=None,
        structural_validation=False,
    )
    if not safety.ok:
        return record

    prior_lex = list(record.lexical_changes)
    record.rewritten = refined.text
    record.lexical_changes = prior_lex + refined.changes
    record.status = "rewritten"
    if not record.template_id:
        record.template_id = "phrase_rewrite"
    elif "+phrase" not in record.template_id:
        record.template_id = f"{record.template_id}+phrase"
    record.confidence = min(
        safety.confidence,
        record.confidence if record.confidence > 0 else max(min_confidence, 0.55),
    )
    return record


def _unchanged(record: SentenceRecord) -> bool:
    return (
        record.rewritten.strip().lower().rstrip(".!?")
        == record.original.strip().lower().rstrip(".!?")
    )


def _try_accept_rewrite(
    record: SentenceRecord,
    candidate: str,
    *,
    template_id: str,
    confidence: float,
    safety_min: float,
    min_confidence: float,
    use_minilm: bool,
    reason_prefix: str,
) -> SentenceRecord | None:
    """Accept a candidate when safety passes and wording truly changed."""
    if not sufficiently_changed(record.original, candidate):
        return None
    safety = check_safety(
        record.original,
        candidate,
        min_meaning=min(safety_min, ENGINE_PARAPHRASE_MIN_SIM),
        min_confidence=min_confidence,
        use_minilm=use_minilm,
        protected_entities=None,
        protected_auxiliaries=None,
        structural_validation=False,
    )
    if not safety.ok:
        hard = check_safety(
            record.original,
            candidate,
            hard_invariants_only=True,
        )
        if not hard.ok:
            return None
        safety = hard

    repaired = repair_sentence(candidate) if GRAMMAR_FIX_OUTPUT else candidate
    if repaired != candidate:
        repaired_safety = check_safety(
            record.original,
            repaired,
            hard_invariants_only=True,
        )
        if repaired_safety.ok and sufficiently_changed(
            record.original, repaired
        ):
            candidate = repaired
            safety = repaired_safety

    prior_status = record.status
    prior_reasons = list(record.reasons)
    record.rewritten = candidate
    record.status = "rewritten"
    record.template_id = template_id
    record.confidence = min(0.75, max(confidence, safety.confidence))
    record.reasons = [
        f"{reason_prefix}_from_{prior_status}:{reason}"
        for reason in prior_reasons
    ] or [reason_prefix]
    return record


def _apply_paraphrase_fallback(
    record: SentenceRecord,
    *,
    enabled: bool,
    safety_min: float,
    min_confidence: float,
    use_minilm: bool,
    primary: bool | None = None,
    variation_seed: int | None = None,
) -> SentenceRecord:
    """Meaning-safe paraphrase rewrite, then structural fallbacks if needed.

    When primary mode is on, paraphrase can replace an existing light structural
    rewrite if the paraphrase keeps meaning and looks substantially newer.
    """
    if record.sentence_type not in _REWRITEABLE_TYPES:
        return record
    prefer_primary = ENGINE_PARAPHRASE_PRIMARY if primary is None else primary

    options: list[tuple[str, str, float, str]] = []
    if enabled:
        result = paraphrase_sentence(
            record.original,
            prefer_divergent=prefer_primary,
        )
        if (
            result.text
            and result.text != record.original
            and sufficiently_changed(record.original, result.text)
        ):
            para_sim = surface_similarity(record.original, result.text)
            current_sim = surface_similarity(record.original, record.rewritten)
            # Always keep paraphrase when the sentence is still unchanged.
            # In primary mode, also prefer it when it is clearly more divergent
            # than the current structural rewrite.
            take_paraphrase = _unchanged(record) or (
                prefer_primary and para_sim + 0.03 < current_sim
            )
            if take_paraphrase:
                options.append(
                    (
                        "paraphrase",
                        result.text,
                        max(0.55, result.confidence),
                        "paraphrase_primary"
                        if prefer_primary
                        else "paraphrase_fallback",
                    )
                )

    # Structural fallbacks only fill gaps. If the plan path already rewrote the
    # sentence, overwriting it here would discard per-request rotation and pin
    # every sentence to the single highest-ranked fallback.
    if _unchanged(record):
        structural = rotate(
            structural_fallback_candidates(record.original),
            seed=variation_seed,
            position=record.index,
            confidence_of=lambda item: item[2],
        )
        for template_id, candidate, confidence in structural:
            options.append(
                (template_id, candidate, confidence, "structure_fallback")
            )

    for template_id, candidate, confidence, reason_prefix in options:
        accepted = _try_accept_rewrite(
            record,
            candidate,
            template_id=template_id,
            confidence=confidence,
            safety_min=safety_min,
            min_confidence=min_confidence,
            use_minilm=use_minilm,
            reason_prefix=reason_prefix,
        )
        if accepted is not None:
            return accepted
    return record


def _apply_forced_rewrite(record: SentenceRecord, *, enabled: bool) -> SentenceRecord:
    """Optional legacy cleft rewrite; disabled by default."""
    if not enabled or record.sentence_type not in _REWRITEABLE_TYPES:
        return record
    if not _unchanged(record):
        return record

    candidate = force_cleft_rewrite(record.original)
    if not candidate:
        return record
    safety = check_safety(
        record.original,
        candidate,
        hard_invariants_only=True,
    )
    if not safety.ok:
        return record
    repaired = repair_sentence(candidate) if GRAMMAR_FIX_OUTPUT else candidate
    if repaired != candidate:
        repaired_safety = check_safety(
            record.original,
            repaired,
            hard_invariants_only=True,
        )
        if not repaired_safety.ok:
            repaired = candidate
        else:
            safety = repaired_safety

    prior_status = record.status
    prior_reasons = list(record.reasons)
    record.rewritten = repaired
    record.status = "rewritten"
    record.template_id = "forced_cleft"
    record.confidence = min(0.60, safety.confidence)
    record.reasons = [
        f"forced_from_{prior_status}:{reason}" for reason in prior_reasons
    ] or ["forced_rewrite"]
    return record


def _finalize_sentence(
    record: SentenceRecord,
    *,
    lexical_enabled: bool,
    force_enabled: bool,
    paraphrase_enabled: bool,
    require_wording: bool,
    lexical_min_wsd: float,
    lexical_max_changes: int | None,
    lexical_polish: bool,
    safety_min: float,
    min_confidence: float,
    use_minilm: bool,
    variation_seed: int | None = None,
) -> SentenceRecord:
    # Paraphrase first as the main rewrite when enabled: pick a meaning-safe
    # but surface-divergent wording. Structural fallbacks cover gaps. Synonym
    # refinement then polishes the chosen rewrite.
    record = _apply_paraphrase_fallback(
        record,
        enabled=paraphrase_enabled,
        safety_min=safety_min,
        min_confidence=min_confidence,
        use_minilm=use_minilm,
        variation_seed=variation_seed,
    )
    record = _apply_forced_rewrite(record, enabled=force_enabled)
    record = _apply_phrase_rewrite(
        record,
        enabled=ENGINE_PHRASE_REWRITE,
        polish=lexical_polish,
        safety_min=safety_min,
        min_confidence=min_confidence,
        use_minilm=use_minilm,
    )
    record = _apply_lexical_refinement(
        record,
        enabled=lexical_enabled,
        min_wsd=lexical_min_wsd,
        max_changes=lexical_max_changes,
        safety_min=safety_min,
        min_confidence=min_confidence,
        use_minilm=use_minilm,
        polish=lexical_polish,
    )
    record = _apply_extra_polish(
        record,
        enabled=lexical_enabled and lexical_polish,
        min_wsd=lexical_min_wsd,
        max_changes=lexical_max_changes,
        safety_min=safety_min,
        min_confidence=min_confidence,
        use_minilm=use_minilm,
    )
    return _ensure_wording_changed(
        record,
        enabled=require_wording,
        lexical_min_wsd=lexical_min_wsd,
        lexical_max_changes=lexical_max_changes,
        lexical_polish=lexical_polish,
        safety_min=safety_min,
        min_confidence=min_confidence,
        use_minilm=use_minilm,
    )


def _ensure_wording_changed(
    record: SentenceRecord,
    *,
    enabled: bool,
    lexical_min_wsd: float,
    lexical_max_changes: int | None,
    lexical_polish: bool,
    safety_min: float,
    min_confidence: float,
    use_minilm: bool,
) -> SentenceRecord:
    """Guarantee a meaning-safe wording change when the sentence is still identical."""
    if not enabled or record.sentence_type not in _REWRITEABLE_TYPES:
        return record
    if not _unchanged(record):
        return record

    # Voice shift is a structural last resort before synonym forcing.
    passive = active_to_passive(record.original)
    if passive:
        accepted = _try_accept_rewrite(
            record,
            passive,
            template_id="ensure_passive",
            confidence=0.62,
            safety_min=safety_min,
            min_confidence=min_confidence,
            use_minilm=use_minilm,
            reason_prefix="ensure_passive",
        )
        if accepted is not None:
            return accepted

    ensured = ensure_wording_change(
        record.original,
        min_wsd=min(lexical_min_wsd, 0.10),
        max_changes=lexical_max_changes,
        polish=lexical_polish,
        classical_strict=(not use_minilm),
        classical_aggressive=(not use_minilm) and ENGINE_CLASSICAL_AGGRESSIVE,
    )
    if not ensured.changes or ensured.text == record.original:
        return record
    if not sufficiently_changed(record.original, ensured.text):
        return record
    safety = check_safety(
        record.original,
        ensured.text,
        min_meaning=min(safety_min, ENGINE_PARAPHRASE_MIN_SIM),
        min_confidence=min_confidence,
        use_minilm=use_minilm,
        protected_entities=None,
        protected_auxiliaries=None,
        structural_validation=False,
    )
    if not safety.ok:
        hard = check_safety(
            record.original,
            ensured.text,
            hard_invariants_only=True,
        )
        if not hard.ok:
            return record
        safety = hard

    prior_status = record.status
    prior_reasons = list(record.reasons)
    record.rewritten = ensured.text
    record.lexical_changes = ensured.changes
    record.status = "rewritten"
    record.template_id = "ensure_wording"
    record.confidence = min(0.70, max(0.55, safety.confidence, ensured.confidence))
    record.reasons = [
        f"ensure_from_{prior_status}:{reason}" for reason in prior_reasons
    ] or ["ensure_wording_change"]
    return record


def _process_sentence(
    text: str,
    *,
    index: int,
    block_index: int,
    min_confidence: float,
    safety_min: float,
    use_minilm: bool,
    lexical_enabled: bool,
    force_enabled: bool,
    paraphrase_enabled: bool,
    require_wording: bool,
    lexical_min_wsd: float,
    lexical_max_changes: int | None,
    lexical_polish: bool,
    last_template: str,
    stats: EngineStats,
    variation_seed: int | None = None,
) -> tuple[SentenceRecord, str]:
    """Rewrite one sentence with plan → generate → grammar → safety → fallback."""
    original = (text or "").strip()
    kind = classify_sentence(original)

    if not original:
        rec = SentenceRecord(
            index=index,
            original=original,
            rewritten=original,
            confidence=1.0,
            status="passthrough",
            sentence_type=kind,
            block_index=block_index,
        )
        stats.passthrough += 1
        return rec, ""

    plan = build_plan(original, min_confidence=min_confidence)
    if not plan.safe:
        reason = plan.skip_reason or kind
        stats.skipped += 1
        stats.bump(reason)
        rec = SentenceRecord(
            index=index,
            original=original,
            rewritten=original,
            confidence=plan.confidence,
            status="skipped",
            sentence_type=kind,
            reasons=[reason],
            block_index=block_index,
        )
        rec = _finalize_sentence(
            rec,
            lexical_enabled=lexical_enabled,
            force_enabled=force_enabled,
            paraphrase_enabled=paraphrase_enabled,
            require_wording=require_wording,
            lexical_min_wsd=lexical_min_wsd,
            lexical_max_changes=lexical_max_changes,
            lexical_polish=lexical_polish,
            safety_min=safety_min,
            min_confidence=min_confidence,
            use_minilm=use_minilm,
            variation_seed=variation_seed,
        )
        return rec, rec.template_id if rec.status == "rewritten" else ""

    # Prefer alternate template if last sentence used the same time-front family
    time_fronts = {
        "time_subj_manner_verb_place",
        "time_subj_verb_place",
        "time_subj_verb_object",
        "time_subj_verb_object_place",
        "time_front",
    }
    # Slot templates alone rarely offer more than one survivor, so fold in the
    # clause/voice/split transforms. Rotation then has real alternatives to pick
    # between instead of always landing on the same fronting.
    pool: list[tuple[str, str | None, float]] = [
        (cand.template_id, None, cand.confidence) for cand in plan.candidates
    ]
    if variation_seed is not None:
        planned = {template_id for template_id, _text, _conf in pool}
        pool += [
            (template_id, text, confidence)
            for template_id, text, confidence in structural_fallback_candidates(
                original, include_plan=False
            )
            # Passivizing reads worse than any reorder, so it stays a last
            # resort in the fallback path rather than a rotation option.
            if template_id not in planned and template_id != "active_to_passive"
        ]

    candidates = rotate(
        pool,
        seed=variation_seed,
        position=index,
        confidence_of=lambda item: item[2],
    )
    if last_template in time_fronts and candidates:
        reordered = [c for c in candidates if c[0] != last_template]
        reordered += [c for c in candidates if c[0] == last_template]
        candidates = reordered

    best: SentenceRecord | None = None
    for template_id, prepared, cand_confidence in candidates:
        generated = (
            prepared
            if prepared is not None
            else generate_from_plan(plan, template_id=template_id)
        )
        if not generated:
            continue
        # Voice flips legitimately move the agent and rebuild the auxiliary
        # chain, so slot-level protections do not apply to them.
        voice_flip = template_id in {"active_to_passive", "passive_to_active"}
        safety = check_safety(
            original,
            generated,
            min_meaning=safety_min,
            min_confidence=min_confidence,
            use_minilm=use_minilm,
            protected_entities=(
                None if voice_flip else plan.slots.entities if plan.slots else ()
            ),
            protected_auxiliaries=(
                None if voice_flip else plan.slots.auxiliaries if plan.slots else ()
            ),
            structural_validation=not voice_flip,
        )
        if not safety.ok and (not use_minilm) and ENGINE_CLASSICAL_AGGRESSIVE:
            # Aggressive classical: accept structure when hard facts still hold.
            hard = check_safety(
                original,
                generated,
                hard_invariants_only=True,
            )
            if hard.ok:
                safety = hard
            else:
                stats.bump(safety.reasons[0] if safety.reasons else "safety")
                continue
        elif not safety.ok:
            stats.bump(safety.reasons[0] if safety.reasons else "safety")
            continue
        repaired = repair_sentence(generated) if GRAMMAR_FIX_OUTPUT else generated
        if repaired != generated:
            repaired_safety = check_safety(
                original,
                repaired,
                min_meaning=safety_min,
                min_confidence=min_confidence,
                use_minilm=use_minilm,
                protected_entities=(
                    None if voice_flip else plan.slots.entities if plan.slots else ()
                ),
                protected_auxiliaries=(
                    None
                    if voice_flip
                    else plan.slots.auxiliaries if plan.slots else ()
                ),
                structural_validation=False,
            )
            if not repaired_safety.ok:
                stats.bump(
                    repaired_safety.reasons[0]
                    if repaired_safety.reasons
                    else "grammar_safety"
                )
                continue
            safety = repaired_safety
        conf = min(cand_confidence, safety.confidence + 0.2)
        best = SentenceRecord(
            index=index,
            original=original,
            rewritten=repaired,
            confidence=conf,
            status="rewritten",
            template_id=template_id,
            sentence_type=kind,
            reasons=[],
            block_index=block_index,
        )
        break

    if best is None:
        stats.reverted += 1
        stats.bump("safety_fallback")
        rec = SentenceRecord(
            index=index,
            original=original,
            rewritten=original,
            confidence=0.0,
            status="reverted",
            sentence_type=kind,
            reasons=["safety_fallback"],
            block_index=block_index,
        )
        rec = _finalize_sentence(
            rec,
            lexical_enabled=lexical_enabled,
            force_enabled=force_enabled,
            paraphrase_enabled=paraphrase_enabled,
            require_wording=require_wording,
            lexical_min_wsd=lexical_min_wsd,
            lexical_max_changes=lexical_max_changes,
            lexical_polish=lexical_polish,
            safety_min=safety_min,
            min_confidence=min_confidence,
            use_minilm=use_minilm,
            variation_seed=variation_seed,
        )
        return rec, rec.template_id if rec.status == "rewritten" else ""

    stats.rewritten += 1
    best = _finalize_sentence(
        best,
        lexical_enabled=lexical_enabled,
        force_enabled=force_enabled,
        paraphrase_enabled=paraphrase_enabled,
        require_wording=require_wording,
        lexical_min_wsd=lexical_min_wsd,
        lexical_max_changes=lexical_max_changes,
        lexical_polish=lexical_polish,
        safety_min=safety_min,
        min_confidence=min_confidence,
        use_minilm=use_minilm,
        variation_seed=variation_seed,
    )
    return best, best.template_id


def _process_block(
    block: DocumentBlock,
    *,
    sent_offset: int,
    min_confidence: float,
    safety_min: float,
    use_minilm: bool,
    lexical_enabled: bool,
    force_enabled: bool,
    paraphrase_enabled: bool,
    require_wording: bool,
    lexical_min_wsd: float,
    lexical_max_changes: int | None,
    lexical_polish: bool,
    last_template: str,
    stats: EngineStats,
    variation_seed: int | None = None,
) -> tuple[DocumentBlock, list[SentenceRecord], str]:
    if not block.rewriteable or block.kind != "paragraph":
        stats.passthrough += 1
        rec = SentenceRecord(
            index=sent_offset,
            original=block.text,
            rewritten=block.text,
            confidence=1.0,
            status="passthrough",
            sentence_type=block.kind,
            block_index=block.index,
            reasons=[f"block:{block.kind}"],
        )
        return block, [rec], last_template

    masked, prot = mask_protected_spans(block.text)
    sentences = split_sentences(masked)
    if not sentences:
        return block, [], last_template

    records: list[SentenceRecord] = []
    out_sents: list[str] = []
    tid = last_template
    for i, sent in enumerate(sentences):
        rec, tid = _process_sentence(
            sent,
            index=sent_offset + i,
            block_index=block.index,
            min_confidence=min_confidence,
            safety_min=safety_min,
            use_minilm=use_minilm,
            lexical_enabled=lexical_enabled,
            force_enabled=force_enabled,
            paraphrase_enabled=paraphrase_enabled,
            require_wording=require_wording,
            lexical_min_wsd=lexical_min_wsd,
            lexical_max_changes=lexical_max_changes,
            lexical_polish=lexical_polish,
            last_template=tid,
            stats=stats,
            variation_seed=variation_seed,
        )
        # Unmask protected spans in both sides of the record
        rec.original = unmask_protected_spans(rec.original, prot)
        rec.rewritten = unmask_protected_spans(rec.rewritten, prot)
        records.append(rec)
        out_sents.append(rec.rewritten)

    new_text = unmask_protected_spans(join_sentences(out_sents), prot)
    new_block = DocumentBlock(
        text=new_text,
        kind=block.kind,
        rewriteable=block.rewriteable,
        index=block.index,
        meta=dict(block.meta),
    )
    return new_block, records, tid


def rewrite_document(
    source: str | bytes | None,
    *,
    batch_paras: int | None = None,
    min_confidence: float | None = None,
    safety_min: float | None = None,
    use_minilm_safety: bool | None = None,
    use_lexical_refinement: bool | None = None,
    force_rewrite: bool | None = None,
    use_paraphrase: bool | None = None,
    require_wording_change: bool | None = None,
    lexical_min_wsd: float | None = None,
    lexical_max_changes: int | None = None,
    lexical_polish: bool | None = None,
    preserve_length: bool | None = None,
    variation_seed: int | None = None,
) -> EngineResult:
    """
    Structural rewrite with local paraphrase and low-impact synonyms.

    Safe spaCy templates still run first. When paraphrase primary mode is on,
    a CPU T5 model then prefers a meaning-safe but surface-divergent rewrite so
    the text looks substantially new. MiniLM guards meaning. Synonym refinement
    runs after paraphrase. When require_wording_change is on, a final ensure
    pass forces a meaning-safe wording change whenever a safe option exists.

    Synonym budget is dynamic from sentence length unless lexical_max_changes
    is set explicitly. lexical_polish densifies that budget (UI "word polish").

    ``variation_seed`` rotates between equally-ranked structural options so the
    same text rewrites differently on each request. Pass an explicit seed for
    reproducible output; pass 0 to disable rotation entirely.
    """
    started = time.perf_counter()
    batch_paras = batch_paras if batch_paras is not None else ENGINE_BATCH_PARAS
    min_confidence = (
        min_confidence if min_confidence is not None else ENGINE_MIN_CONFIDENCE
    )
    safety_min = safety_min if safety_min is not None else ENGINE_SAFETY_MIN
    use_minilm = (
        use_minilm_safety
        if use_minilm_safety is not None
        else ENGINE_USE_MINILM_SAFETY
    )
    lexical_enabled = (
        use_lexical_refinement
        if use_lexical_refinement is not None
        else ENGINE_LEXICAL_REFINEMENT
    )
    force_enabled = (
        force_rewrite if force_rewrite is not None else ENGINE_FORCE_REWRITE
    )
    paraphrase_enabled = (
        use_paraphrase if use_paraphrase is not None else ENGINE_PARAPHRASE
    )
    require_wording = (
        require_wording_change
        if require_wording_change is not None
        else ENGINE_REQUIRE_WORDING_CHANGE
    )
    lexical_min_wsd = (
        lexical_min_wsd
        if lexical_min_wsd is not None
        else ENGINE_LEXICAL_MIN_WSD
    )
    # None => fully dynamic per-sentence budget inside refine_sentence.
    if lexical_max_changes is not None and lexical_max_changes <= 0:
        lexical_max_changes = None
    polish_enabled = bool(lexical_polish)
    keep_length = (
        ENGINE_PRESERVE_LENGTH if preserve_length is None else bool(preserve_length)
    )
    if variation_seed is None:
        variation_seed = new_seed() if ENGINE_STRUCTURAL_VARIATION else None
    elif variation_seed == 0:
        variation_seed = None
    if (force_enabled or paraphrase_enabled or require_wording) and min_confidence > 0.40:
        # Prefer more structural templates before generative fallback.
        min_confidence = 0.40

    raw = ingest_text(source)
    if not raw.strip():
        raise ValueError("Paste some text first.")

    normalized = normalize_text(raw)

    blocks = detect_blocks(normalized)
    stats = EngineStats(blocks=len(blocks))
    all_records: list[SentenceRecord] = []
    out_blocks: list[DocumentBlock] = []
    sent_offset = 0
    last_template = ""

    for batch in iter_paragraph_batches(blocks, batch_paras=batch_paras):
        stats.batches += 1
        for block in batch:
            new_block, records, last_template = _process_block(
                block,
                sent_offset=sent_offset,
                min_confidence=min_confidence,
                safety_min=safety_min,
                use_minilm=use_minilm,
                lexical_enabled=lexical_enabled,
                force_enabled=force_enabled,
                paraphrase_enabled=paraphrase_enabled,
                require_wording=require_wording,
                lexical_min_wsd=lexical_min_wsd,
                lexical_max_changes=lexical_max_changes,
                lexical_polish=polish_enabled,
                last_template=last_template,
                stats=stats,
                variation_seed=variation_seed,
            )
            out_blocks.append(new_block)
            all_records.extend(records)
            # Advance offset by sentence count for rewriteable paras
            if block.rewriteable and block.kind == "paragraph":
                sent_offset += max(1, len(records))
            else:
                sent_offset += 1

    stats.sentences = len(all_records)
    stitched = stitch_blocks(out_blocks)

    # Consistency: suppress streaks then rebuild from records if needed
    stitched, all_records = apply_consistency(stitched, all_records)
    streak_reverts = [r for r in all_records if "template_streak" in r.reasons]
    if streak_reverts:
        # Rebuild rewriteable paragraph texts from records
        by_block: dict[int, list[SentenceRecord]] = {}
        for rec in all_records:
            by_block.setdefault(rec.block_index, []).append(rec)
        rebuilt: list[DocumentBlock] = []
        for block in out_blocks:
            recs = by_block.get(block.index, [])
            if block.rewriteable and block.kind == "paragraph" and recs:
                text = join_sentences([r.rewritten for r in recs])
                rebuilt.append(
                    DocumentBlock(
                        text=text,
                        kind=block.kind,
                        rewriteable=True,
                        index=block.index,
                    )
                )
            else:
                rebuilt.append(block)
        stitched = stitch_blocks(rebuilt)
        stitched, all_records = apply_consistency(stitched, all_records)

    stitched = tidy(stitched)
    if keep_length:
        trimmed = enforce_length_budget(
            normalized,
            stitched,
            preserve_length=True,
        )
        if trimmed != stitched:
            stitched = tidy(trimmed)
            notes_trim = "length_budget"
        else:
            notes_trim = ""
    else:
        notes_trim = ""

    stats.rewritten = sum(r.status == "rewritten" for r in all_records)
    stats.skipped = sum(r.status == "skipped" for r in all_records)
    stats.reverted = sum(r.status == "reverted" for r in all_records)
    stats.passthrough = sum(r.status == "passthrough" for r in all_records)
    stats.lexical_refined = sum(bool(r.lexical_changes) for r in all_records)
    stats.forced_rewrites = sum(
        (r.template_id or "").startswith("forced_cleft") for r in all_records
    )
    stats.paraphrased = sum(
        (r.template_id or "").startswith("paraphrase")
        or any(
            reason.startswith("structure_fallback") for reason in (r.reasons or [])
        )
        for r in all_records
    )
    stats.seconds = round(time.perf_counter() - started, 3)
    mapping = [(r.original, r.rewritten) for r in all_records]
    skipped = [r for r in all_records if r.status in {"skipped", "reverted"}]
    in_words = word_count(normalized)
    out_words = word_count(stitched)
    changed = stitched.strip() != normalized.strip()
    sim = _similarity(normalized, stitched)

    notes = (
        f"structural: rewritten={stats.rewritten} skipped={stats.skipped} "
        f"reverted={stats.reverted} lexical={stats.lexical_refined} "
        f"paraphrased={stats.paraphrased} forced={stats.forced_rewrites} "
        f"batches={stats.batches}"
    )
    if notes_trim:
        notes = f"{notes}; {notes_trim}"
    logger.info(
        "engine done words=%s→%s rewritten=%s skipped=%s reverted=%s "
        "paraphrased=%s lexical=%s batches=%s t=%.2fs",
        in_words,
        out_words,
        stats.rewritten,
        stats.skipped,
        stats.reverted,
        stats.paraphrased,
        stats.lexical_refined,
        stats.batches,
        stats.seconds,
    )

    return EngineResult(
        text=stitched,
        sentences=all_records,
        skipped=skipped,
        mapping=mapping,
        stats=stats,
        notes=notes,
        engine="structural-reorder",
        input_words=in_words,
        output_words=out_words,
        changed=changed,
        similarity=round(sim, 4),
    )