import gradio as gr
from functools import lru_cache

@lru_cache(maxsize=1)
def get_rewriter():
    """
    Use a lightweight instruction-following model that can run on CPU.
    flan-t5-small is generally more reliable for "rewrite" than GPT2-style models.
    """
    from transformers import pipeline
    return pipeline(
        task="text2text-generation",
        model="google/flan-t5-small",
        device=-1,  # CPU
    )

def build_prompt(text: str, style: str) -> str:
    text = (text or "").strip()
    if not text:
        return ""

    if style == "More formal":
        return (
            "Rewrite the text in a more formal tone. "
            "Keep the original meaning. Output only the rewritten text.\n\n"
            f"Text: {text}"
        )
    if style == "More friendly":
        return (
            "Rewrite the text in a friendly, warm tone. "
            "Keep the original meaning. Output only the rewritten text.\n\n"
            f"Text: {text}"
        )
    return (
        "Rewrite the text to be shorter and clearer. "
        "Keep the original meaning. Output only the rewritten text.\n\n"
        f"Text: {text}"
    )

def rewrite(text: str, style: str) -> str:
    text = (text or "").strip()
    if not text:
        return "Please enter some text."

    prompt = build_prompt(text, style)

    try:
        rewriter = get_rewriter()
        out = rewriter(
            prompt,
            max_new_tokens=128,
            do_sample=False,
        )
        result = (out[0].get("generated_text") or "").strip()
        return result if result else "No output. Try a shorter input."
    except Exception as e:
        return (
            "Error: failed to run the model.\n"
            "If this is the first run, the Space may still be downloading the model.\n\n"
            f"Details: {type(e).__name__}: {e}"
        )

with gr.Blocks(title="AI Text Rewriter") as demo:
    gr.Markdown(
        "AI Text Rewriter\n"
        "Paste a sentence or short paragraph, choose a style, then rewrite with AI."
    )

    with gr.Row():
        style = gr.Radio(
            ["More formal", "More friendly", "Shorter"],
            value="More friendly",
            label="Rewrite style",
        )

    text_input = gr.Textbox(
        label="Your text",
        placeholder="Type or paste text here...",
        lines=5,
    )

    with gr.Row():
        btn = gr.Button("Rewrite with AI")
        clear = gr.Button("Clear")

    output = gr.Textbox(label="Result", lines=6)

    btn.click(fn=rewrite, inputs=[text_input, style], outputs=output)
    clear.click(fn=lambda: ("", ""), inputs=None, outputs=[text_input, output])

demo.launch()