app/engine/paraphrase/__init__.py

"""CPU paraphraser with MiniLM-ranked candidate selection."""

from __future__ import annotations

import logging
import re
import threading
from dataclasses import dataclass
from difflib import SequenceMatcher
from functools import lru_cache
from typing import Any

from app.config import (
    ENGINE_PARAPHRASE_MAX_NEW_TOKENS,
    ENGINE_PARAPHRASE_MAX_SURFACE,
    ENGINE_PARAPHRASE_MIN_DIVERGENCE,
    ENGINE_PARAPHRASE_MIN_SIM,
    ENGINE_PARAPHRASE_MODEL,
    ENGINE_PARAPHRASE_NUM_RETURN,
    ENGINE_PARAPHRASE_PRIMARY,
)
from app.pipeline.minilm import pick_best_candidate, score_candidate

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

_lock = threading.Lock()
_tokenizer = None
_model = None
_failed = False
_WORD = re.compile(r"[A-Za-z']+")


@dataclass
class ParaphraseResult:
    text: str
    confidence: float = 0.0
    reason: str = ""
    candidates: list[str] | None = None


def paraphrase_resource_available() -> bool:
    return _get_pipeline() is not None


@lru_cache(maxsize=1)
def _torch_device() -> str:
    try:
        import torch

        return "cpu"
    except Exception:
        return "cpu"


def _get_pipeline() -> tuple[Any, Any] | None:
    global _tokenizer, _model, _failed
    if _failed:
        return None
    if _tokenizer is not None and _model is not None:
        return _tokenizer, _model
    with _lock:
        if _failed:
            return None
        if _tokenizer is not None and _model is not None:
            return _tokenizer, _model
        try:
            import torch
            from transformers import AutoModelForSeq2SeqLM, AutoTokenizer

            name = ENGINE_PARAPHRASE_MODEL
            _tokenizer = AutoTokenizer.from_pretrained(name)
            _model = AutoModelForSeq2SeqLM.from_pretrained(name)
            _model.to(_torch_device())
            _model.eval()
            logger.info("Paraphrase model ready (%s)", name)
            return _tokenizer, _model
        except Exception as exc:
            logger.warning("Paraphrase model unavailable: %s", exc)
            _failed = True
            _tokenizer = None
            _model = None
            return None


def warm_paraphrase() -> bool:
    try:
        pipe = _get_pipeline()
        if pipe is None:
            return False
        paraphrase_sentence("Warm up the rewrite model.", num_return=1)
        return True
    except Exception as exc:
        logger.warning("Paraphrase warm failed: %s", exc)
        return False


def _normalize_candidate(text: str) -> str:
    value = re.sub(r"\s+", " ", (text or "").strip())
    value = value.strip(" \"'")
    if not value:
        return ""
    if value[-1] not in ".!?":
        value += "."
    return value[:1].upper() + value[1:]


def _content_overlap(source: str, candidate: str) -> float:
    src = {w.lower() for w in _WORD.findall(source) if len(w) >= 3}
    cand = {w.lower() for w in _WORD.findall(candidate) if len(w) >= 3}
    if not src:
        return 0.0
    return len(src & cand) / max(1, len(src))


def surface_similarity(source: str, candidate: str) -> float:
    left = re.sub(r"\s+", " ", (source or "").strip().lower()).rstrip(".!?")
    right = re.sub(r"\s+", " ", (candidate or "").strip().lower()).rstrip(".!?")
    if not left or not right:
        return 0.0
    # autojunk discards any character filling >1% of a sequence longer than 200,
    # which for prose means every space and common letter. The ratio would then
    # be decided by rare letters alone and swing wildly on long inputs.
    return SequenceMatcher(None, left, right, autojunk=False).ratio()


def sufficiently_changed(
    source: str,
    candidate: str,
    *,
    max_surface: float | None = None,
    min_divergence: float | None = None,
) -> bool:
    """Require a real wording or order change, not punctuation/near-copy edits."""
    src = (source or "").strip()
    cand = (candidate or "").strip()
    if not src or not cand:
        return False
    limit = (
        max_surface
        if max_surface is not None
        else ENGINE_PARAPHRASE_MAX_SURFACE
    )
    src_tokens = [w.lower() for w in _WORD.findall(src)]
    cand_tokens = [w.lower() for w in _WORD.findall(cand)]
    if not src_tokens or src_tokens == cand_tokens:
        return False
    similarity = surface_similarity(src, cand)
    divergence = 1.0 - similarity
    needed = (
        min_divergence
        if min_divergence is not None
        else 0.0
    )
    if needed > 0 and divergence < needed:
        return False
    if similarity >= limit:
        # High surface overlap is OK only when several content words differ
        # (clear synonym/reorder), not a one-character or tiny tweak.
        src_set = {w for w in src_tokens if len(w) >= 3}
        cand_set = {w for w in cand_tokens if len(w) >= 3}
        if len(src_set ^ cand_set) < 2:
            return False
    return True


def _generate_raw(
    text: str,
    *,
    num_return: int,
    max_surface: float,
    min_overlap: float = 0.32,
    prompts: list[str] | None = None,
    max_new_tokens: int | None = None,
) -> list[str]:
    loaded = _get_pipeline()
    if loaded is None:
        return []
    tokenizer, model = loaded
    import torch

    prompt_list = prompts or [
        f"paraphrase: {text.strip()}",
        f"rewrite with different wording but same meaning: {text.strip()}",
    ]
    returns = max(1, num_return)
    per_prompt = max(1, returns // len(prompt_list))
    if per_prompt * len(prompt_list) < returns:
        per_prompt += 1
    token_budget = (
        max_new_tokens
        if max_new_tokens is not None
        else ENGINE_PARAPHRASE_MAX_NEW_TOKENS
    )
    decoded: list[str] = []
    for prompt in prompt_list:
        encoded = tokenizer(
            prompt,
            return_tensors="pt",
            truncation=True,
            max_length=256,
        )
        prompt_returns = per_prompt
        groups = min(prompt_returns, 4)
        beams = max(6, prompt_returns * 2)
        beams = max(groups, (beams // groups) * groups)
        with torch.no_grad():
            try:
                outputs = model.generate(
                    **encoded,
                    max_new_tokens=token_budget,
                    num_beams=beams,
                    num_beam_groups=groups,
                    diversity_penalty=1.0,
                    num_return_sequences=prompt_returns,
                    do_sample=False,
                    early_stopping=True,
                )
            except Exception:
                # Fall back to plain beam search if diverse beams are unsupported.
                outputs = model.generate(
                    **encoded,
                    max_new_tokens=token_budget,
                    num_beams=max(4, prompt_returns),
                    num_return_sequences=prompt_returns,
                    do_sample=False,
                    early_stopping=True,
                )
        decoded.extend(tokenizer.batch_decode(outputs, skip_special_tokens=True))
    cleaned: list[str] = []
    seen: set[str] = set()
    source_key = text.strip().lower().rstrip(".!?")
    for item in decoded:
        value = _normalize_candidate(item)
        key = value.lower().rstrip(".!?")
        if not value or key == source_key or key in seen:
            continue
        if _content_overlap(text, value) < min_overlap:
            continue
        if not sufficiently_changed(text, value, max_surface=max_surface):
            continue
        seen.add(key)
        cleaned.append(value)
    return cleaned


def paraphrase_sentence(
    text: str,
    *,
    min_sim: float | None = None,
    num_return: int | None = None,
    prefer_divergent: bool | None = None,
    max_surface: float | None = None,
    min_divergence: float | None = None,
) -> ParaphraseResult:
    """Generate a meaning-preserving paraphrase when the model is available.

    When prefer_divergent is on (primary rewrite mode), choose the most
    surface-different candidate that still clears the MiniLM meaning floor.
    """
    source = (text or "").strip()
    if not source:
        return ParaphraseResult(text=text, reason="empty")
    threshold = (
        min_sim if min_sim is not None else ENGINE_PARAPHRASE_MIN_SIM
    )
    divergent = (
        ENGINE_PARAPHRASE_PRIMARY
        if prefer_divergent is None
        else prefer_divergent
    )
    surface_limit = (
        max_surface
        if max_surface is not None
        else ENGINE_PARAPHRASE_MAX_SURFACE
    )
    divergence_floor = (
        min_divergence
        if min_divergence is not None
        else (ENGINE_PARAPHRASE_MIN_DIVERGENCE if divergent else 0.0)
    )
    returns = (
        num_return
        if num_return is not None
        else max(ENGINE_PARAPHRASE_NUM_RETURN, 5 if divergent else 3)
    )
    candidates = _generate_raw(
        source,
        num_return=returns,
        max_surface=surface_limit,
        min_overlap=0.30 if divergent else 0.35,
    )
    if not candidates:
        return ParaphraseResult(text=source, reason="no_candidates")

    surface_scores = {
        candidate: surface_similarity(source, candidate)
        for candidate in candidates
    }
    # Prefer MiniLM ranking when available; otherwise keep the first beam.
    best = pick_best_candidate(
        source,
        candidates,
        min_meaning=threshold,
        prefer_divergent=divergent,
        surface_scores=surface_scores,
    )
    if best is None:
        scored = []
        for candidate in candidates:
            meaning = score_candidate(source, candidate)
            if meaning is None or meaning <= 0.0:
                # MiniLM unavailable — use lexical overlap as a soft ranker.
                meaning = _content_overlap(source, candidate)
            if meaning >= threshold * 0.85 and sufficiently_changed(
                source,
                candidate,
                max_surface=surface_limit,
                min_divergence=divergence_floor if divergent else 0.0,
            ):
                scored.append(
                    (
                        surface_scores.get(candidate, 1.0) if divergent else -meaning,
                        -meaning if divergent else surface_scores.get(candidate, 1.0),
                        meaning,
                        candidate,
                    )
                )
        if not scored:
            return ParaphraseResult(
                text=source,
                reason="below_similarity",
                candidates=candidates,
            )
        scored.sort()
        best = scored[0][3]
        confidence = scored[0][2]
    else:
        if not sufficiently_changed(
            source,
            best,
            max_surface=surface_limit,
            min_divergence=divergence_floor if divergent else 0.0,
        ):
            # Ranker picked a near-copy; try the next sufficiently changed option.
            alternates = [
                candidate
                for candidate in candidates
                if candidate != best
                and sufficiently_changed(
                    source,
                    candidate,
                    max_surface=surface_limit,
                    min_divergence=divergence_floor if divergent else 0.0,
                )
            ]
            if not alternates:
                return ParaphraseResult(
                    text=source,
                    reason="near_copy",
                    candidates=candidates,
                )
            reranked = pick_best_candidate(
                source,
                alternates,
                min_meaning=threshold,
                prefer_divergent=divergent,
                surface_scores=surface_scores,
            )
            best = reranked or alternates[0]
        scored_best = score_candidate(source, best)
        confidence = (
            scored_best
            if scored_best is not None
            else _content_overlap(source, best)
        )

    if not sufficiently_changed(
        source,
        best,
        max_surface=surface_limit,
        min_divergence=divergence_floor if divergent else 0.0,
    ):
        return ParaphraseResult(
            text=source,
            reason="near_copy",
            candidates=candidates,
        )

    return ParaphraseResult(
        text=best,
        confidence=round(float(confidence), 4),
        candidates=candidates,
    )


def _normalize_span_candidate(text: str) -> str:
    value = re.sub(r"\s+", " ", (text or "").strip())
    value = value.strip(" \"'")
    # Spans are mid-sentence fragments — keep lowercase lead when source is.
    return value.rstrip(".!?")


_PROMPT_LEAK = re.compile(
    r"\b(?:paraphrase|rewrite|different wording|same meaning|rephrase)\b",
    re.I,
)


def _span_candidate_ok(source: str, candidate: str) -> bool:
    """Reject prompt echoes and structurally invalid mid-sentence spans."""
    src = (source or "").strip()
    cand = (candidate or "").strip()
    if not src or not cand:
        return False
    low = cand.lower()
    if _PROMPT_LEAK.search(low):
        return False
    if ":" in cand or ";" in cand:
        return False
    src_tokens = [token.lower() for token in _WORD.findall(src)]
    cand_tokens = [token.lower() for token in _WORD.findall(cand)]
    if not src_tokens or not cand_tokens:
        return False
    # Keep span length close so splicing stays grammatical.
    if len(cand_tokens) > len(src_tokens) + 1 or len(cand_tokens) + 2 < len(src_tokens):
        return False
    if len(cand) > max(12, int(len(src) * 1.45)):
        return False
    # Mid-sentence VO spans usually start with a verb; reject full-clause flips
    # like "Customer loyalty strengthens" / "customer loyalty strengthens".
    if src[0].islower() and cand[0].isupper():
        return False
    src_tail = set(src_tokens[1:])
    if cand_tokens[0] in src_tail:
        return False
    # Reject duplicated content words ("…experience to create").
    from collections import Counter

    src_counts = Counter(src_tokens)
    cand_counts = Counter(cand_tokens)
    for token, count in cand_counts.items():
        if len(token) < 4:
            continue
        if count > max(1, src_counts.get(token, 0)):
            return False
    # Reject dangling infinitive tails invented by T5.
    if re.search(r"\bto\s+[a-z]{3,}$", low) and " to " not in src.lower():
        return False
    return True


def paraphrase_span(
    text: str,
    *,
    num_return: int | None = None,
    min_sim: float | None = None,
) -> ParaphraseResult:
    """Paraphrase a short verb–object span for phrase-level rewriting."""
    source = (text or "").strip()
    if not source:
        return ParaphraseResult(text=text, reason="empty")
    if _get_pipeline() is None:
        return ParaphraseResult(text=source, reason="unavailable")

    returns = num_return if num_return is not None else 4
    threshold = min_sim if min_sim is not None else max(0.70, ENGINE_PARAPHRASE_MIN_SIM - 0.02)
    # Spans: use only the short paraphrase prompt and a tight token budget so
    # the model does not echo instruction text into the fragment.
    raw = _generate_raw(
        source,
        num_return=returns,
        max_surface=0.92,
        min_overlap=0.40,
        prompts=[f"paraphrase: {source}"],
        max_new_tokens=min(24, ENGINE_PARAPHRASE_MAX_NEW_TOKENS),
    )
    cleaned: list[str] = []
    seen: set[str] = set()
    source_key = source.lower().rstrip(".!?")
    source_lead_lower = bool(source[:1].islower())
    for item in raw:
        value = _normalize_span_candidate(item)
        if source_lead_lower and value[:1].isupper():
            value = value[:1].lower() + value[1:]
        key = value.lower()
        if not value or key == source_key or key in seen:
            continue
        if not _span_candidate_ok(source, value):
            continue
        if _content_overlap(source, value) < 0.40:
            continue
        if not sufficiently_changed(source, value, max_surface=0.94, min_divergence=0.06):
            continue
        seen.add(key)
        cleaned.append(value)
    if not cleaned:
        return ParaphraseResult(text=source, reason="no_candidates", candidates=[])

    surface_scores = {
        candidate: surface_similarity(source, candidate) for candidate in cleaned
    }
    best = pick_best_candidate(
        source,
        cleaned,
        min_meaning=threshold,
        prefer_divergent=True,
        surface_scores=surface_scores,
    )
    if best is None:
        ranked = sorted(
            cleaned,
            key=lambda c: (surface_scores.get(c, 1.0), -_content_overlap(source, c)),
        )
        best = ranked[0]
    if not _span_candidate_ok(source, best):
        return ParaphraseResult(text=source, reason="no_candidates", candidates=cleaned)
    meaning = score_candidate(source, best)
    confidence = meaning if meaning is not None else _content_overlap(source, best)
    return ParaphraseResult(
        text=best,
        confidence=round(float(confidence), 4),
        candidates=cleaned,
        reason="span_paraphrase",
    )