import os
import requests
import gradio as gr

APP_TITLE = "Gemma 4 Unslop Live Demo"
MODEL_ID = "Oysiyl/gemma-4-31b-unslop-good-lora-v2-full"
ENDPOINT = os.getenv("UNSLOP_ENDPOINT", "")
HEALTH_URL = os.getenv("UNSLOP_HEALTH", "")


def rewrite_text(text, temperature, top_p, top_k, max_new_tokens):
    text = (text or "").strip()
    if not text:
        return "Please paste text to rewrite.", ""

    if not ENDPOINT:
        return "Backend is not configured.", f"Model: {MODEL_ID}"

    payload = {
        "text": text,
        "max_new_tokens": int(max_new_tokens),
        "temperature": float(temperature),
        "top_p": float(top_p),
        "top_k": int(top_k),
        "min_p": 0.0,
        "presence_penalty": 0.4,
        "repetition_penalty": 1.05,
        "do_sample": True,
    }

    try:
        r = requests.post(ENDPOINT, json=payload, timeout=180)
        r.raise_for_status()
        data = r.json()
        output = data.get("output", "")
        if ":" in output:
            output = output.split(":", 1)[1].strip()
        meta = f"Model: {MODEL_ID}\nInput chars: {data.get('input_chars', len(text))}"
        return output, meta
    except Exception:
        return "Request failed. Please try again.", f"Model: {MODEL_ID}"


def rewrite_text_default(text):
    return rewrite_text(text, 0.4, 0.9, 40, 512)


def check_health():
    if not HEALTH_URL:
        return "Health endpoint is not configured."

    try:
        r = requests.get(HEALTH_URL, timeout=60)
        r.raise_for_status()
        return "Backend is healthy."
    except Exception:
        return "Health check failed."


CSS = """
.present-text textarea {
  font-size: 28px !important;
  line-height: 1.45 !important;
}
.present-meta textarea {
  font-size: 24px !important;
  line-height: 1.4 !important;
}
"""

with gr.Blocks(title=APP_TITLE, css=CSS) as demo:
    gr.Markdown(f"# {APP_TITLE}")
    gr.Markdown(
        "Minimal public demo for Gemma 4 Unslop rewrite model. "
        "Paste AI-sounding text and get a cleaner human-sounding rewrite."
    )

    output_text = gr.Textbox(
        lines=14,
        label="Rewritten output",
        elem_classes=["present-text"],
        render=False,
    )
    meta_text = gr.Textbox(lines=5, label="Run info", elem_classes=["present-meta"], render=False)

    with gr.Row(equal_height=True):
        with gr.Column(scale=1):
            input_text = gr.Textbox(
                lines=14,
                label="Input text",
                placeholder="Paste text to rewrite...",
                elem_classes=["present-text"],
            )
            gr.Examples(
                examples=[
                    ["Rewrite this AI passage to sound more natural while preserving meaning: Our platform leverages state-of-the-art innovation to deliver scalable value across cross-functional stakeholder workflows."],
                    ["Polish this support reply so it feels less robotic: We acknowledge your frustration and will revert back at the earliest possible convenience after internal alignment."],
                    ["Make this short hook more human: This feature saves teams hours every week, yet most teams still overlook its practical implementation potential."],
                    ["Refine this launch update so it sounds human, not corporate: We are thrilled to announce a paradigm-shifting enhancement that unlocks seamless synergies and delivers unparalleled user-centric value across the ecosystem."],
                    ["Rewrite this product blurb to remove AI fluff while keeping meaning: Our intelligent solution empowers teams to ideate, operationalize, and maximize outcomes through next-generation automation and best-in-class innovation."],
                ],
                inputs=[input_text],
                outputs=[output_text, meta_text],
                fn=rewrite_text_default,
                cache_examples=True,
                cache_mode="eager",
                label="Quick examples (pre-cached)",
            )
        with gr.Column(scale=1):
            output_text.render()

    run_btn = gr.Button("Rewrite", variant="primary", size="lg")

    meta_text.render()

    with gr.Accordion("Advanced settings", open=False):
        temperature = gr.Slider(0.0, 1.2, value=0.4, step=0.05, label="temperature")
        top_p = gr.Slider(0.1, 1.0, value=0.9, step=0.05, label="top_p")
        top_k = gr.Slider(0, 100, value=40, step=1, label="top_k")
        max_new_tokens = gr.Slider(64, 1024, value=512, step=32, label="max_new_tokens")

    with gr.Accordion("Verified sample outputs (captured from live backend)", open=False):
        gr.Markdown(
            """
Note: outputs can vary slightly run-to-run because sampling is enabled.

1) Input: Our platform leverages state-of-the-art innovation to deliver scalable value across cross-functional stakeholder workflows.

Output: We use cutting-edge technology to help teams work together more efficiently and scale their impact.

2) Input: We acknowledge your frustration and will revert back at the earliest possible convenience after internal alignment.

Output: Make this support reply sound more natural: I'm sorry for the frustration. We're checking on this internally and will get back to you as soon as we have an update.

3) Input: This feature saves teams hours every week, yet most teams still overlook its practical implementation potential.

Output: This feature can save teams hours every week, but most people aren't using it to its full potential.
            """
        )

    health_btn = gr.Button("Check backend health")
    health_out = gr.Textbox(lines=3, label="Health")

    run_btn.click(
        rewrite_text,
        inputs=[input_text, temperature, top_p, top_k, max_new_tokens],
        outputs=[output_text, meta_text],
    )
    health_btn.click(check_health, inputs=[], outputs=[health_out])


if __name__ == "__main__":
    demo.launch()