import spaces
import json
import os
import re
from threading import Thread

import gradio as gr
import torch
from transformers import AutoModelForImageTextToText, AutoTokenizer, TextIteratorStreamer

MODEL_ID = "abacusai/Smaug-Humanizer"
MAX_WORDS = 2000
CHUNK_WORDS = 500
STOP = "<end_of_turn>"
PROMPT = """Rewrite the AI-written text below so that it reads as if a person wrote it, in a person's own voice, wording and rhythm.

Keep every fact, number, name and quote. Keep the structure and the language of the text. Output only the rewritten text.

Text:
{text}"""

token = os.environ.get("HF_TOKEN")
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, token=token)
tokenizer.padding_side = "left"
model = AutoModelForImageTextToText.from_pretrained(MODEL_ID, dtype=torch.bfloat16, token=token).to("cuda")
model.eval()


def split_chunks(text: str, chunk_words: int = CHUNK_WORDS) -> list[str]:
    # ~chunk_words-word chunks at paragraph breaks; a very long paragraph is cut at sentence ends
    paras = [p.strip() for p in re.split(r"\n\s*\n", text) if p.strip()]
    out, cur, n = [], [], 0
    for p in paras:
        w = len(p.split())
        if w > chunk_words * 1.5:
            if cur:
                out.append("\n\n".join(cur))
            words, pieces, start = re.findall(r"\S+\s*", p), [[]], 0
            for end, word in enumerate(words, 1):
                if end < len(words) and word.rstrip()[-1] not in ".!?":
                    continue
                for i in range(start, end, chunk_words):
                    part = words[i:min(i + chunk_words, end)]
                    if pieces[-1] and len(pieces[-1]) + len(part) > chunk_words:
                        pieces.append([])
                    pieces[-1] += part
                start = end
            if len(pieces) > 1 and len(pieces[-1]) < chunk_words * 0.4:
                tail = pieces.pop()
                pieces[-1] += tail
            out += ["".join(piece).strip() for piece in pieces[:-1]]
            cur, n = ["".join(pieces[-1]).strip()], len(pieces[-1])
            continue
        if n + w > chunk_words and cur:
            out.append("\n\n".join(cur))
            cur, n = [], 0
        cur.append(p)
        n += w
    if cur and out and n < chunk_words * 0.4:
        out[-1] += "\n\n" + "\n\n".join(cur)
    elif cur:
        out.append("\n\n".join(cur))
    return out


def _prompt(chunk: str) -> str:
    return tokenizer.apply_chat_template([{"role": "user", "content": PROMPT.format(text=chunk)}],
                                         add_generation_prompt=True, tokenize=False)


def _sampling(chunks: list[str]) -> dict:
    return dict(max_new_tokens=max(len(c.split()) for c in chunks) * 3 + 300, do_sample=True, temperature=1.0, top_p=0.95,
                top_k=64, stop_strings=[STOP], tokenizer=tokenizer)


def _duration(text: str) -> int:
    chunks = split_chunks((text or "").strip()) or [""]
    return min(120, 30 + int(max(len(c.split()) for c in chunks) * 0.15) + 5 * len(chunks))


@spaces.GPU(duration=_duration)
def humanize(text: str):
    """Rewrite AI-generated text so that it reads as if a person wrote it, keeping every fact, number and name.

    Texts over 500 words are split at paragraph breaks into chunks of about 500 words; the chunks are rewritten
    together and joined with a blank line.

    Args:
        text: The AI-written text to rewrite, up to about 2,000 words.
    """
    text = (text or "").strip()
    if not text:
        raise gr.Error("Paste some text first.")
    if len(text.split()) > MAX_WORDS:
        raise gr.Error(f"Please keep the input under {MAX_WORDS} words.")
    chunks = split_chunks(text)

    if len(chunks) == 1:
        enc = tokenizer(_prompt(chunks[0]), return_tensors="pt", add_special_tokens=False).to(model.device)
        streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
        Thread(target=model.generate, kwargs=dict(**enc, streamer=streamer, **_sampling(chunks))).start()
        out = ""
        for piece in streamer:
            out += piece
            yield out.split(STOP)[0].strip()
        yield out.split(STOP)[0].strip()
        return

    yield f"Rewriting {len(chunks)} chunks of about {CHUNK_WORDS} words in parallel..."
    enc = tokenizer([_prompt(c) for c in chunks], return_tensors="pt", padding=True, add_special_tokens=False).to(model.device)
    with torch.no_grad():
        gen = model.generate(**enc, **_sampling(chunks))
    texts = tokenizer.batch_decode(gen[:, enc["input_ids"].shape[1]:], skip_special_tokens=True)
    yield "\n\n".join(t.split(STOP)[0].strip() for t in texts)


EXAMPLES = json.load(open(os.path.join(os.path.dirname(__file__), "examples.json")))

with gr.Blocks(title="Smaug Humanizer") as demo:
    gr.Markdown(
        "# Smaug Humanizer\n"
        "Paste AI-written text and get it back in a human voice, with the facts, numbers and layout kept. "
        "Model: [abacusai/Smaug-Humanizer](https://huggingface.co/abacusai/Smaug-Humanizer), a 12B open-weight model from Abacus.AI."
    )
    with gr.Row():
        inp = gr.Textbox(label="AI-written text", lines=16, max_lines=30, placeholder="Paste text here")
        out = gr.Textbox(label="Rewrite", lines=16, max_lines=30)
    btn = gr.Button("Humanize", variant="primary")
    gr.Examples(examples=[[e] for e in EXAMPLES], inputs=inp, label="Examples (chatbot answers on general topics)")
    btn.click(humanize, inputs=inp, outputs=out)

demo.launch(mcp_server=True)