import os

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import spaces  # must come before torch
import torch
import gradio as gr
from threading import Thread
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer

MODEL_ID = "Altworld/Hemmingway-1"

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID, dtype=torch.bfloat16, attn_implementation="sdpa"
).to("cuda")
model.eval()

EXAMPLES = [
    "Write the text I send my landlord about the broken boiler.",
    "Tell my manager I need Friday off for a family thing, without oversharing.",
    "Decline a friend's wedding invite nicely. It's in Portugal and I can't afford it.",
    "Chase a client who's 30 days late on a $2,400 invoice. Firm but friendly.",
    "Write a short note to a neighbour asking them to keep the dog quiet at night.",
    "Talk my brother into coming to Mum's 60th. He's been avoiding family stuff.",
    "Pitch me an original TV series: logline, the main characters and what they want, "
    "and a season one arc episode by episode. No sci-fi clichés, and make the ending hurt a little.",
    "Draft the plot of a small, quiet indie film about two estranged sisters clearing out "
    "their late father's house. Three acts, and give me one scene written out in full.",
]


def _to_messages(message: str, history: list, system_prompt: str) -> list[dict]:
    msgs = []
    if system_prompt.strip():
        msgs.append({"role": "system", "content": system_prompt.strip()})
    for turn in history:
        # Skip the collapsed "Thinking" bubbles; only real replies go back in.
        if turn.get("metadata"):
            continue
        content = turn["content"]
        if isinstance(content, list):
            content = " ".join(c.get("text", "") for c in content if isinstance(c, dict))
        if content:
            msgs.append({"role": turn["role"], "content": content})
    msgs.append({"role": "user", "content": message})
    return msgs


def _duration(message, history, system_prompt, thinking, max_new_tokens, *args, **kwargs):
    return min(240, 30 + int(max_new_tokens) // 12 + (60 if thinking else 0))


@spaces.GPU(duration=_duration, size="xlarge")
def respond(
    message: str,
    history: list,
    system_prompt: str = "",
    thinking: bool = False,
    max_new_tokens: int = 1280,
    temperature: float = 0.7,
    top_p: float = 0.95,
):
    """Ask Hemmingway-1 to write something: a text, email, note, or reply.

    Args:
        message: What you want written.
        history: Previous chat turns.
        system_prompt: Optional instructions that apply to the whole chat.
        thinking: Let the model reason before answering (slower).
        max_new_tokens: Maximum length of the reply.
        temperature: Sampling temperature; higher is more varied.
        top_p: Nucleus sampling cutoff.
    """
    msgs = _to_messages(message, history, system_prompt)
    template_kwargs = {"enable_thinking": bool(thinking)}
    if thinking:
        template_kwargs["reasoning_effort"] = "low"
    input_ids = tokenizer.apply_chat_template(
        msgs, add_generation_prompt=True, return_tensors="pt", return_dict=True, **template_kwargs
    ).to("cuda")

    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
    Thread(
        target=model.generate,
        kwargs=dict(
            **input_ids,
            streamer=streamer,
            max_new_tokens=int(max_new_tokens) + (1024 if thinking else 0),
            do_sample=True,
            temperature=float(temperature),
            top_p=float(top_p),
            top_k=20,
        ),
        daemon=True,
    ).start()

    raw = ""
    for piece in streamer:
        raw += piece
        if not thinking:
            yield raw.strip()
            continue
        # The template already opened <think>, so text before </think> is reasoning.
        if "</think>" in raw:
            reasoning, answer = raw.split("</think>", 1)
            yield [
                gr.ChatMessage(content=reasoning.strip(), metadata={"title": "Thinking", "status": "done"}),
                gr.ChatMessage(content=answer.strip()),
            ]
        else:
            yield [gr.ChatMessage(content=raw.strip(), metadata={"title": "Thinking…", "status": "pending"})]


CSS = """
#col-container { max-width: 860px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
#hero h1 { margin-bottom: 0; }
#hero p { margin-top: 4px; opacity: .75; }
"""

with gr.Blocks(title="Hemmingway-1") as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            "# ✍️ Hemmingway-1\n"
            "The AI that writes like a person. Ask for the text, the email, the awkward note — "
            "you get the message, not a memo.  \n"
            "[Model card](https://huggingface.co/Altworld/Hemmingway-1) · 27B · bf16 · Apache-2.0",
            elem_id="hero",
        )
        gr.ChatInterface(
            fn=respond,
            chatbot=gr.Chatbot(height=520, show_label=False, layout="bubble"),
            textbox=gr.Textbox(placeholder="What should Hemmingway write for you?", submit_btn=True),
            examples=[[e] for e in EXAMPLES],
            cache_examples=False,
            additional_inputs=[
                gr.Textbox(label="System prompt (optional)", placeholder="e.g. I'm British, keep it casual.", lines=2),
                gr.Checkbox(label="Thinking (slower, shows reasoning)", value=False),
                gr.Slider(64, 2048, value=1280, step=64, label="Max new tokens"),
                gr.Slider(0.1, 1.5, value=0.7, step=0.05, label="Temperature"),
                gr.Slider(0.5, 1.0, value=0.95, step=0.01, label="Top-p"),
            ],
            additional_inputs_accordion=gr.Accordion("Settings", open=False),
        )
        gr.Markdown(
            "<small>English-first. It can be wrong and still sound certain — don't use it for medical, legal or financial decisions.</small>"
        )

if __name__ == "__main__":
    demo.launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True)