"""Light post-rewrite cleanup and length control."""

from __future__ import annotations

import re


def tidy(text: str) -> str:
    """Normalize whitespace without flattening document structure."""
    value = re.sub(r"\n{3,}", "\n\n", text or "")
    value = re.sub(r"[ \t]{2,}", " ", value)
    value = re.sub(r'([,;:])(["“])', r"\1 \2", value)
    paragraphs = [part.strip() for part in re.split(r"\n\s*\n", value) if part.strip()]
    return "\n\n".join(paragraphs).strip()


def enforce_length_budget(
    original: str,
    rewritten: str,
    *,
    preserve_length: bool = True,
) -> str:
    """Trim only when the rewrite clearly ballooned past the original.

    Preserves paragraph breaks (\\n\\n). Never trims so hard that most of the
    source content disappears.
    """
    source = original or ""
    output = rewritten or ""
    source_words = len(source.split())
    output_words = len(output.split())
    if source_words == 0:
        return output

    # Slightly looser ceiling so normal phrase/lexical growth is kept.
    max_ratio = 1.20 if preserve_length else 1.6
    max_words = max(1, int(source_words * max_ratio))
    # Floor: do not destroy content just to hit the ceiling.
    min_words = max(1, int(source_words * 0.90))
    if output_words <= max_words:
        return output

    paragraphs = [
        part.strip() for part in re.split(r"\n\s*\n", output.strip()) if part.strip()
    ]
    if not paragraphs:
        return output

    kept_paragraphs: list[str] = []
    count = 0
    for paragraph in paragraphs:
        parts = re.split(r"(?<=[.!?])\s+", paragraph.strip())
        kept_sentences: list[str] = []
        for part in parts:
            width = len(part.split())
            if kept_sentences and count + width > max_words:
                break
            if not kept_sentences and kept_paragraphs and count + width > max_words:
                break
            kept_sentences.append(part)
            count += width
        if kept_sentences:
            kept_paragraphs.append(" ".join(kept_sentences))
        if count >= max_words:
            break
    trimmed = "\n\n".join(kept_paragraphs).strip()
    if not trimmed or len(trimmed.split()) < min_words:
        # Trimming would drop too much meaning — keep the full rewrite.
        return output
    return trimmed