"""
Estella Explainer — streaming Gradio chat interface

This version keeps Gradio's default appearance and adds the two important
interaction patterns used by the DeepSeek demo: token streaming and a queue.
"""

import os
from threading import Thread
from typing import Iterator

import gradio as gr
import spaces
import torch
from peft import PeftModel
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    TextIteratorStreamer,
)


ADAPTER_ID = "sgattup/EstellaExplainerLLM"
BASE_MODEL_ID = "unsloth/gemma-4-E4B-it"

HF_TOKEN = os.getenv("HF_TOKEN")
HUB_ARGS = {"token": HF_TOKEN} if HF_TOKEN else {}


SYSTEM_PROMPT = """You are Estella Explainer, a math reading and language coach for young learners in grades 3 through 8.

Estella helps students understand what a math problem is asking before they begin solving it.

Rules:
- Give only ONE hint per problem.
- Do NOT solve the problem or show the answer.
- Do NOT show calculation steps.
- Do NOT compute anything.
- Help the student understand what the problem is asking.
- Use very simple words and short sentences.
- Identify what is known and what is being asked.
- Use phrases like "We know...", "We are looking for...", or "Think about...".
- Be friendly and encouraging.
"""


EXAMPLES = [
    "A rectangle has a length of 12 cm and a width of 5 cm. What is the area?",
    "Sam had 3/4 of a pizza. He ate 1/4. How much is left?",
    "A store sells apples for $0.75 each. Jake buys 6 apples. How much does he spend?",
    "Round 4,567 to the nearest hundred.",
    "There are 48 students put into equal groups of 6. How many groups are there?",
    "A bag has 5 red marbles and 3 blue marbles. What fraction of the marbles are red?",
    "A train leaves at 10:15 AM and arrives at 1:45 PM. How long is the trip?",
    "The perimeter of a square is 36 inches. What is the length of one side?",
]


print(f"Loading base model: {BASE_MODEL_ID}")

base_model = AutoModelForCausalLM.from_pretrained(
    BASE_MODEL_ID,
    dtype=torch.float16,
    low_cpu_mem_usage=True,
    **HUB_ARGS,
)

print(f"Loading adapter on CPU: {ADAPTER_ID}")

model = PeftModel.from_pretrained(
    base_model,
    ADAPTER_ID,
    torch_device="cpu",
    **HUB_ARGS,
)

print("Moving combined model to CUDA.")
model = model.to("cuda")
model.eval()

tokenizer = AutoTokenizer.from_pretrained(
    ADAPTER_ID,
    **HUB_ARGS,
)

print("Estella is ready.")


@spaces.GPU(duration=60)
def ask_estella(message: str, history: list) -> Iterator[str]:
    """Stream one reading hint without solving the math problem."""

    if not message or not message.strip():
        yield (
            "Please type a math word problem and I’ll help explain "
            "what it is asking!"
        )
        return

    messages = [
        {
            "role": "system",
            "content": SYSTEM_PROMPT,
        },
        {
            "role": "user",
            "content": (
                "Math Problem:\n"
                f"{message.strip()}\n\n"
                "Estella's Hint:"
            ),
        },
    ]

    prompt = tokenizer.apply_chat_template(
        messages,
        add_generation_prompt=True,
        tokenize=False,
    )

    inputs = tokenizer(
        prompt,
        return_tensors="pt",
    )

    input_device = next(model.parameters()).device
    inputs = {
        key: value.to(input_device)
        for key, value in inputs.items()
    }

    streamer = TextIteratorStreamer(
        tokenizer,
        timeout=60.0,
        skip_prompt=True,
        skip_special_tokens=True,
    )

    generation_kwargs = {
        **inputs,
        "streamer": streamer,
        "max_new_tokens": 150,
        "temperature": 0.7,
        "do_sample": True,
        "repetition_penalty": 1.1,
        "pad_token_id": tokenizer.eos_token_id,
    }

    generation_thread = Thread(
        target=model.generate,
        kwargs=generation_kwargs,
    )
    generation_thread.start()

    partial_hint = ""
    for new_text in streamer:
        partial_hint += new_text
        yield partial_hint.strip()


demo = gr.ChatInterface(
    fn=ask_estella,
    title="Estella Explainer Math Reading Coach",
    description=(
        "Estella helps students understand what a math problem is asking "
        "before they begin solving it.\n\n"
        "Enter a math question below for one clear, friendly hint. "
        "Estella will not give away the answer.\n\n"
        "Developed by [Sai Gattupalli, PhD](https://huggingface.co/sgattup) · "
        "[UMass Amherst College of Education](https://www.umass.edu/education/) · "
        "[Society and AI Research Group](https://societyandai.org) · "
        "[Usable Math](https://usablemath.org)"
    ),
    examples=EXAMPLES,
    cache_examples=False,
)


if __name__ == "__main__":
    demo.queue(max_size=20).launch()