"""Gradio web frontend for MiniRewrite.
Self-contained application ready for Hugging Face Spaces.
"""

import sys
import time
from pathlib import Path
import gradio as gr
import torch

# Ensure local src package is on path
ROOT_DIR = Path(__file__).resolve().parent
if str(ROOT_DIR / "src") not in sys.path:
    sys.path.insert(0, str(ROOT_DIR / "src"))

from minirewrite.pipeline import MiniRewrite

# Global model instance
print("Initializing MiniRewrite model...")
MODEL_DIR = ROOT_DIR if (ROOT_DIR / "best.pt").exists() else (ROOT_DIR / "runs" / "rtx4060")

# CPU inference is extremely fast (~50ms) for this 7.6M model and requires zero GPU cold-start
device = "cuda" if torch.cuda.is_available() else "cpu"
pipeline = MiniRewrite.from_pretrained(MODEL_DIR, checkpoint_name="best.pt", device=device)
print(f"Model loaded successfully on {pipeline.device}!")


def process_rewrite(text: str, task: str, tone: str, style: str):
    if not text or not text.strip():
        return "Please enter a sentence to rewrite.", "0 ms"

    start_time = time.perf_counter()
    result = pipeline.rewrite(
        text.strip(),
        task=task.lower(),
        tone=tone.lower(),
        style=style.lower(),
        max_source_length=128,
        max_new_tokens=128,
    )
    elapsed_ms = (time.perf_counter() - start_time) * 1000

    latency_str = f"Generated in {elapsed_ms:.1f} ms on {pipeline.device}"
    return result, latency_str


theme = gr.themes.Soft(
    primary_hue="indigo",
    secondary_hue="purple",
)

with gr.Blocks(theme=theme, title="MiniRewrite - Transformer Text Rewriter") as demo:
    with gr.Column(elem_id="title-banner"):
        gr.Markdown(
            """
            # ✍️ MiniRewrite
            **Instruction-Controlled Encoder-Decoder Transformer (7.6M Parameters)**
            
            Rewrites text by controlling **Task**, **Tone**, and **Style** in real time.
            """
        )

    with gr.Row():
        with gr.Column(scale=5):
            input_text = gr.Textbox(
                label="Input Sentence",
                placeholder="e.g. send me report asap",
                lines=3,
                max_lines=5,
            )

            with gr.Row():
                task_dropdown = gr.Dropdown(
                    label="Task",
                    choices=["rewrite", "simplify", "shorten"],
                    value="rewrite",
                )
                tone_dropdown = gr.Dropdown(
                    label="Tone",
                    choices=["professional", "friendly", "neutral", "casual"],
                    value="professional",
                )
                style_dropdown = gr.Dropdown(
                    label="Style",
                    choices=["polite", "direct", "concise", "plain"],
                    value="polite",
                )

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

        with gr.Column(scale=5):
            output_text = gr.Textbox(
                label="Rewritten Output",
                lines=3,
                max_lines=5,
                interactive=False,
            )
            latency_badge = gr.Markdown("Latency: ready")

    submit_btn.click(
        fn=process_rewrite,
        inputs=[input_text, task_dropdown, tone_dropdown, style_dropdown],
        outputs=[output_text, latency_badge],
    )

    gr.Markdown("### 💡 Click any preset example to test instantly:")
    gr.Examples(
        examples=[
            ["send me report asap", "rewrite", "professional", "polite"],
            ["send me report asap", "rewrite", "friendly", "direct"],
            ["cant come to meeting tmrw", "rewrite", "professional", "polite"],
            ["pls fix the bug today", "rewrite", "friendly", "direct"],
            ["In order to complete the project on time, we need more resources.", "shorten", "neutral", "concise"],
            ["There was a concert of Michael Jackson in my city last month .", "simplify", "neutral", "plain"],
        ],
        inputs=[input_text, task_dropdown, tone_dropdown, style_dropdown],
        outputs=[output_text, latency_badge],
        fn=process_rewrite,
        cache_examples=False,
    )

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