rewrite.py

"""Any2human rewrite pipeline — claim-atom content creation flow."""

from __future__ import annotations

import os
import re
import time
from dataclasses import dataclass
from typing import Callable, Optional

from openai import OpenAI

import prompts
from mechanics import mechanical_humanize

ProgressCallback = Optional[Callable[[str], None]]

CHUNK_TARGET_WORDS = 600
CHUNK_OVERLAP_WORDS = 40
MAX_INPUT_WORDS = 4500
MAX_RETRIES = 4


@dataclass
class RewriteResult:
    text: str
    mode: str
    chunks: int
    api_calls: int
    input_words: int
    output_words: int
    model: str
    warnings: list[str]


def _env(name: str, default: str = "") -> str:
    return (os.environ.get(name) or default).strip()


def get_client() -> OpenAI:
    api_key = _env("OPENROUTER_API_KEY")
    if not api_key:
        raise RuntimeError(
            "Missing OPENROUTER_API_KEY. Add it as a Space secret or local env var."
        )
    return OpenAI(
        base_url="https://openrouter.ai/api/v1",
        api_key=api_key,
        default_headers={
            "HTTP-Referer": _env("SPACE_HOST", "https://huggingface.co/spaces"),
            "X-Title": _env("APP_TITLE", "Any2human"),
        },
    )


def default_model() -> str:
    return _env("OPENROUTER_MODEL", "openrouter/free")


def _notify(cb: ProgressCallback, message: str) -> None:
    if cb:
        cb(message)


def _chat(

    client: OpenAI,

    *,

    model: str,

    system: str,

    user: str,

    temperature: float = 0.85,

    presence_penalty: float = 0.35,

    frequency_penalty: float = 0.35,

    progress: ProgressCallback = None,

) -> str:
    last_error: Exception | None = None
    for attempt in range(1, MAX_RETRIES + 1):
        try:
            response = client.chat.completions.create(
                model=model,
                temperature=temperature,
                presence_penalty=presence_penalty,
                frequency_penalty=frequency_penalty,
                messages=[
                    {"role": "system", "content": system},
                    {"role": "user", "content": user},
                ],
            )
            content = (response.choices[0].message.content or "").strip()
            if not content:
                raise RuntimeError("Empty response from the model.")
            return _strip_fences(content)
        except Exception as exc:  # noqa: BLE001
            last_error = exc
            msg = str(exc).lower()
            retryable = any(
                token in msg
                for token in ("429", "rate", "timeout", "temporar", "503", "502", "overloaded")
            )
            if not retryable or attempt == MAX_RETRIES:
                break
            wait = min(2**attempt, 20)
            _notify(progress, f"Rate limited / busy — retrying in {wait}s…")
            time.sleep(wait)
    raise RuntimeError(f"OpenRouter request failed: {last_error}")


def _strip_fences(text: str) -> str:
    text = text.strip()
    if text.startswith("```"):
        text = re.sub(r"^```(?:\w+)?\s*", "", text)
        text = re.sub(r"\s*```$", "", text)
    return text.strip()


def _split_into_chunks(text: str, target: int = CHUNK_TARGET_WORDS) -> list[str]:
    words = text.split()
    if len(words) <= target:
        return [text.strip()]

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

    chunks: list[str] = []
    current: list[str] = []
    current_words = 0

    for para in paragraphs:
        p_words = len(para.split())
        if current and current_words + p_words > target:
            chunks.append("\n\n".join(current))
            if current and len(current[-1].split()) <= CHUNK_OVERLAP_WORDS:
                current = [current[-1], para]
                current_words = len(current[-1].split()) + p_words
            else:
                current = [para]
                current_words = p_words
        else:
            current.append(para)
            current_words += p_words

    if current:
        chunks.append("\n\n".join(current))
    return chunks


def _merge_chunks(parts: list[str]) -> str:
    if len(parts) == 1:
        return parts[0].strip()
    merged: list[str] = []
    for part in parts:
        cleaned = part.strip()
        if not cleaned:
            continue
        if merged:
            prev_tail = " ".join(merged[-1].split()[-25:]).lower()
            lead = " ".join(cleaned.split()[:25]).lower()
            if lead and lead in prev_tail:
                sentences = re.split(r"(?<=[.!?])\s+", cleaned)
                cleaned = " ".join(sentences[1:]).strip() or cleaned
        merged.append(cleaned)
    return "\n\n".join(merged).strip()


def _maybe_compress(

    client: OpenAI,

    *,

    model: str,

    text: str,

    max_words: int,

    progress: ProgressCallback,

    api_calls: int,

) -> tuple[str, int]:
    if len(text.split()) <= int(max_words * 1.12):
        return text, api_calls
    _notify(progress, "Compressing to length budget…")
    out = _chat(
        client,
        model=model,
        system=prompts.system_compress(),
        user=prompts.user_compress(text, max_words),
        temperature=0.35,
        presence_penalty=0.15,
        frequency_penalty=0.15,
        progress=progress,
    )
    return out, api_calls + 1


def _pipeline_chunk(

    client: OpenAI,

    *,

    chunk: str,

    chunk_index: int,

    tone: str,

    voice_sample: str,

    quality: str,

    preserve_length: bool,

    model: str,

    progress: ProgressCallback,

) -> tuple[str, int]:
    """New flow per chunk. Returns (text, api_calls_used)."""
    calls = 0
    label = f"section {chunk_index}"
    chunk_words = len(chunk.split())
    _, max_words, length_rule = prompts.length_budget(chunk_words, preserve_length)
    seed = prompts.style_seed(chunk_index - 1)

    if quality == "Fast":
        _notify(progress, f"Fast rewrite {label}…")
        out = _chat(
            client,
            model=model,
            system=prompts.system_fast(tone),
            user=prompts.user_fast(chunk, tone, seed, length_rule, voice_sample),
            temperature=0.95,
            presence_penalty=0.55,
            frequency_penalty=0.5,
            progress=progress,
        )
        calls += 1
        out, calls = _maybe_compress(
            client, model=model, text=out, max_words=max_words, progress=progress, api_calls=calls
        )
        return out, calls

    # Shared: atomize (throws away AI sentence skeleton)
    _notify(progress, f"Atomizing claims for {label}…")
    atoms = _chat(
        client,
        model=model,
        system=prompts.system_atomize(),
        user=prompts.user_atomize(chunk),
        temperature=0.15,
        presence_penalty=0.0,
        frequency_penalty=0.0,
        progress=progress,
    )
    calls += 1

    if quality == "Best":
        # Interview → answers → weave (rebuilds discourse from scratch)
        _notify(progress, f"Building questions for {label}…")
        questions = _chat(
            client,
            model=model,
            system=prompts.system_interview(),
            user=prompts.user_interview(atoms),
            temperature=0.5,
            presence_penalty=0.2,
            frequency_penalty=0.2,
            progress=progress,
        )
        calls += 1

        _notify(progress, f"Answering in human bursts ({label})…")
        answers = _chat(
            client,
            model=model,
            system=prompts.system_answer(tone),
            user=prompts.user_answer(questions, atoms),
            temperature=0.95,
            presence_penalty=0.65,
            frequency_penalty=0.55,
            progress=progress,
        )
        calls += 1
        material = answers
    else:
        # Balanced: weave directly from atoms
        material = atoms

    _notify(progress, f"Weaving prose for {label} (seed: {seed[:28]}…)…")
    woven = _chat(
        client,
        model=model,
        system=prompts.system_weave(tone),
        user=prompts.user_weave(
            material=material,
            tone=tone,
            seed=seed,
            length_rule=length_rule,
            voice_sample=voice_sample,
            anti_source=chunk,
        ),
        temperature=0.92,
        presence_penalty=0.6,
        frequency_penalty=0.5,
        progress=progress,
    )
    calls += 1

    if quality == "Best":
        _notify(progress, f"Bridging seams for {label}…")
        woven = _chat(
            client,
            model=model,
            system=prompts.system_bridge(tone),
            user=prompts.user_bridge(woven, max_words=max_words),
            temperature=0.7,
            presence_penalty=0.4,
            frequency_penalty=0.35,
            progress=progress,
        )
        calls += 1

    woven, calls = _maybe_compress(
        client,
        model=model,
        text=woven,
        max_words=max_words,
        progress=progress,
        api_calls=calls,
    )
    return woven, calls


def rewrite_document(

    text: str,

    *,

    tone: str = "Neutral",

    voice_sample: str = "",

    quality: str = "Best",

    preserve_length: bool = True,

    model: str | None = None,

    progress: ProgressCallback = None,

) -> RewriteResult:
    text = (text or "").strip()
    if not text:
        raise ValueError("Paste or upload some text first.")

    words = len(text.split())
    warnings: list[str] = []
    if words > MAX_INPUT_WORDS:
        raise ValueError(
            f"Input is {words:,} words. Please keep under {MAX_INPUT_WORDS:,} words "
            "on the free tier (split long documents)."
        )

    model = (model or default_model()).strip() or default_model()
    client = get_client()
    api_calls = 0
    chunks = _split_into_chunks(text)

    if quality == "Fast":
        mode = "fast: meaning rewrite + mechanics"
    elif quality == "Best":
        mode = "best: atoms → interview → answers → weave → bridge + mechanics"
    else:
        mode = "balanced: atoms → weave → mechanics"

    if len(chunks) > 1:
        warnings.append(f"Split into {len(chunks)} sections; each uses its own style seed.")

    outputs: list[str] = []
    for i, chunk in enumerate(chunks, start=1):
        out, used = _pipeline_chunk(
            client,
            chunk=chunk,
            chunk_index=i,
            tone=tone,
            voice_sample=voice_sample,
            quality=quality,
            preserve_length=preserve_length,
            model=model,
            progress=progress,
        )
        api_calls += used
        outputs.append(out)

    final = _merge_chunks(outputs)
    _, global_max, _ = prompts.length_budget(words, preserve_length)
    final, api_calls = _maybe_compress(
        client,
        model=model,
        text=final,
        max_words=global_max,
        progress=progress,
        api_calls=api_calls,
    )
    if len(final.split()) > int(global_max * 1.05):
        warnings.append(f"Trimmed toward ~{global_max} words length budget.")

    _notify(progress, "Applying mechanical human rhythm…")
    final = mechanical_humanize(final, tone=tone)
    _notify(progress, "Done.")
    return RewriteResult(
        text=final,
        mode=mode,
        chunks=len(chunks),
        api_calls=api_calls,
        input_words=words,
        output_words=len(final.split()),
        model=model,
        warnings=warnings,
    )