"""
SchoolSpirit AI – hardened chatbot Space with “Clear chat” button
----------------------------------------------------------------
• Loads Meta Llama‑3 3 B‑Instruct via transformers.
• Keeps last MAX_TURNS exchanges to fit context.
• Adds a “Clear chat” button that resets history.
• Extensive try/except blocks to prevent widget crashes.
"""

import re
import gradio as gr
from transformers import (
    AutoTokenizer,
    AutoModelForCausalLM,
    pipeline,
)
from transformers.utils import logging as hf_logging

# ────────────────────────── Config ──────────────────────────────────────────
hf_logging.set_verbosity_error()
LOG = hf_logging.get_logger("SchoolSpirit")

MODEL_ID     = "meta-llama/Llama-3.2-3B-Instruct"
MAX_TURNS    = 6
MAX_TOKENS   = 220
MAX_INPUT_CH = 500

SYSTEM_MSG = (
    "You are SchoolSpirit AI, the friendly digital mascot for a company that "
    "offers on‑prem AI chat mascots, fine‑tuning services, and turnkey GPU "
    "hardware for schools. Keep answers concise, upbeat, and age‑appropriate. "
    "If unsure, say so and suggest contacting a human. Never request personal data."
)

# ────────────────────────── Model Load ──────────────────────────────────────
try:
    tok = AutoTokenizer.from_pretrained(MODEL_ID)
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_ID, device_map="auto", torch_dtype="auto"
    )
    generator = pipeline(
        "text-generation",
        model=model,
        tokenizer=tok,
        max_new_tokens=MAX_TOKENS,
        do_sample=True,
        temperature=0.7,
    )
    MODEL_LOAD_ERR = None
except Exception as exc:  # noqa: BLE001
    MODEL_LOAD_ERR = f"Model load error: {exc}"
    generator = None
    LOG.error(MODEL_LOAD_ERR)

# ────────────────────────── Helpers ────────────────────────────────────────
def truncate_history(hist, max_turns):
    """Return only the last `max_turns` (user,bot) pairs."""
    return hist[-max_turns:] if len(hist) > max_turns else hist


def safe_reply(msg: str) -> str:
    """Post‑process model output and ensure it is non‑empty."""
    msg = msg.strip()
    msg = re.sub(r"\s+", " ", msg)
    return msg or "…"

# ────────────────────────── Chat Callback ───────────────────────────────────
def chat(history, user_msg):
    # Gradio guarantees history is a list of tuples (user, bot)
    history = list(history)

    # Start‑up failure fallback
    if MODEL_LOAD_ERR:
        history.append((user_msg, MODEL_LOAD_ERR))
        return history, ""

    # Basic guards
    user_msg = (user_msg or "").strip()
    if not user_msg:
        history.append(("", "Please enter a message."))
        return history, ""
    if len(user_msg) > MAX_INPUT_CH:
        history.append(
            (user_msg, "Sorry, that message is too long. Please shorten it.")
        )
        return history, ""

    history = truncate_history(history, MAX_TURNS)

    # Build prompt
    prompt_parts = [SYSTEM_MSG]
    for u, a in history:
        prompt_parts += [f"User: {u}", f"AI: {a}"]
    prompt_parts += [f"User: {user_msg}", "AI:"]
    prompt = "\n".join(prompt_parts)

    # Generate
    try:
        completion = generator(prompt, truncate=2048)[0]["generated_text"]
        reply = safe_reply(completion.split("AI:", 1)[-1])
    except Exception as err:  # noqa: BLE001
        LOG.error(f"Inference error: {err}")
        reply = (
            "Sorry—I'm having trouble right now. Please try again in a moment."
        )

    history.append((user_msg, reply))
    return history, ""

# ────────────────────────── Clear Chat fn ───────────────────────────────────
def clear_chat():
    return [], ""

# ────────────────────────── UI Launch ───────────────────────────────────────
with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue")) as demo:
    gr.Markdown("# SchoolSpirit AI Chat")
    chatbot = gr.Chatbot()
    msg_in  = gr.Textbox(placeholder="Ask me anything about SchoolSpirit AI…")
    send_btn = gr.Button("Send")
    clear_btn = gr.Button("Clear Chat", variant="secondary")

    send_btn.click(chat, [chatbot, msg_in], [chatbot, msg_in])
    msg_in.submit(chat, [chatbot, msg_in], [chatbot, msg_in])
    clear_btn.click(clear_chat, outputs=[chatbot, msg_in])

demo.launch()