# ============================================================
# WAAMU AI — Qwen3-1.7B
# Hugging Face ZeroGPU Space
# CPU INFERENCE ONLY
# ============================================================

import spaces

import re
import torch
import gradio as gr

from transformers import (
    AutoTokenizer,
    AutoModelForCausalLM,
)


# ============================================================
# CONFIG
# ============================================================

APP_NAME = "WAAMU AI"

MODEL_NAME = "Qwen/Qwen3-0.6B"

DEVICE = torch.device("cpu")

MAX_NEW_TOKENS = 512

TEMPERATURE = 0.7
TOP_P = 0.8
TOP_K = 20
REPETITION_PENALTY = 1.05


# ============================================================
# ZERO GPU STARTUP REGISTRATION
#
# IMPORTANT:
# This exists ONLY because the Space is configured as
# ZeroGPU and Hugging Face requires an @spaces.GPU function.
#
# WAAMU DOES NOT CALL THIS FUNCTION.
#
# Actual AI inference is CPU-only.
# ============================================================

@spaces.GPU(duration=1)
def _zerogpu_registration():
    return None


# ============================================================
# DEFAULT SYSTEM PROMPT
# ============================================================

DEFAULT_SYSTEM_PROMPT = """
You are WAAMU AI.

Your public identity is ALWAYS WAAMU AI.

If the user asks:
- What is your name?
- Who are you?
- What AI are you?
- What model are you?
- Are you Qwen?
- What powers you?

Answer that you are WAAMU AI.

Never identify yourself as Qwen, Qwen3,
Hugging Face, Transformers, or another underlying
model as your public identity.

You are a helpful general-purpose AI assistant.

Think carefully before answering.

For difficult questions, reason through the problem
internally before producing the final answer.

Do not expose your internal reasoning or chain of thought.

Return only the useful final answer.

There is no persistent memory.

Be accurate, natural, and concise.
""".strip()


# ============================================================
# STARTUP
# ============================================================

print("========================================", flush=True)
print("Starting WAAMU AI...", flush=True)
print(f"Model: {MODEL_NAME}", flush=True)
print("Inference device: CPU", flush=True)
print("Reasoning mode: ENABLED", flush=True)
print("Internal reasoning: HIDDEN", flush=True)
print("ZeroGPU registration: ENABLED", flush=True)
print("Actual GPU inference: DISABLED", flush=True)
print("========================================", flush=True)


# ============================================================
# TOKENIZER
# ============================================================

print("Loading tokenizer...", flush=True)

tokenizer = AutoTokenizer.from_pretrained(
    MODEL_NAME,
)

print("Tokenizer loaded.", flush=True)


# ============================================================
# MODEL
#
# CPU ONLY
#
# No CUDA.
# No .to("cuda").
# No GPU inference.
# ============================================================

print("Loading model on CPU...", flush=True)

model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    dtype=torch.float32,
)

model = model.to(DEVICE)

model.eval()

print("WAAMU AI model loaded on CPU.", flush=True)
print("Waiting for requests...", flush=True)


# ============================================================
# RESPONSE CLEANING
# ============================================================

def clean_qwen3_response(text: str) -> str:
    """
    Remove Qwen3 internal reasoning from the generated output.

    Qwen3 may produce:

        <think>
        internal reasoning...
        </think>

        final answer

    Only the final answer is returned.
    """

    if not text:
        return ""

    text = str(text).strip()

    # --------------------------------------------------------
    # Remove complete <think>...</think> blocks
    # --------------------------------------------------------

    text = re.sub(
        r"<think>.*?</think>",
        "",
        text,
        flags=re.DOTALL | re.IGNORECASE,
    )

    # --------------------------------------------------------
    # Handle an incomplete/unclosed think block
    #
    # This protects against generation being cut off before
    # Qwen3 produces </think>.
    # --------------------------------------------------------

    if "<think>" in text.lower():

        text = re.split(
            r"<think>",
            text,
            maxsplit=1,
            flags=re.IGNORECASE,
        )[0]

    # --------------------------------------------------------
    # Remove stray reasoning markers
    # --------------------------------------------------------

    text = re.sub(
        r"</think>",
        "",
        text,
        flags=re.IGNORECASE,
    )

    text = text.strip()

    return text


# ============================================================
# CPU GENERATION
# ============================================================

def waamu_generate(
    message: str,
    system_prompt: str = "",
):
    """
    Generate one WAAMU AI response.

    Qwen3 reasoning is enabled internally.

    The <think> section is removed before the response
    is returned to the user.

    CPU ONLY.
    """

    message = str(message or "").strip()
    system_prompt = str(system_prompt or "").strip()

    if not message:
        return "Please enter a message."

    # --------------------------------------------------------
    # Select temporary system prompt
    # --------------------------------------------------------

    active_system_prompt = (
        system_prompt
        if system_prompt
        else DEFAULT_SYSTEM_PROMPT
    )

    # --------------------------------------------------------
    # Messages
    # --------------------------------------------------------

    messages = [
        {
            "role": "system",
            "content": active_system_prompt,
        },
        {
            "role": "user",
            "content": message,
        },
    ]

    # --------------------------------------------------------
    # Build Qwen3 chat prompt
    #
    # enable_thinking=True tells Qwen3 to reason internally.
    # --------------------------------------------------------

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

    # --------------------------------------------------------
    # Tokenize
    # --------------------------------------------------------

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

    # Explicit CPU placement.
    inputs = {
        key: value.to(DEVICE)
        for key, value in inputs.items()
    }

    # --------------------------------------------------------
    # Generate
    # --------------------------------------------------------

    with torch.inference_mode():

        output = model.generate(
            **inputs,

            max_new_tokens=MAX_NEW_TOKENS,

            do_sample=True,

            temperature=TEMPERATURE,

            top_p=TOP_P,

            top_k=TOP_K,

            repetition_penalty=REPETITION_PENALTY,

            pad_token_id=tokenizer.eos_token_id,

            eos_token_id=tokenizer.eos_token_id,
        )

    # --------------------------------------------------------
    # Remove prompt tokens
    # --------------------------------------------------------

    input_length = inputs["input_ids"].shape[1]

    generated_tokens = output[0][input_length:]

    # --------------------------------------------------------
    # Decode
    # --------------------------------------------------------

    raw_response = tokenizer.decode(
        generated_tokens,
        skip_special_tokens=True,
    ).strip()

    # --------------------------------------------------------
    # Hide Qwen3 reasoning
    # --------------------------------------------------------

    reply = clean_qwen3_response(
        raw_response
    )

    # --------------------------------------------------------
    # Safety fallback
    # --------------------------------------------------------

    if not reply:

        return (
            "I couldn't generate a final answer. "
            "Please try again."
        )

    return reply


# ============================================================
# GRADIO HANDLER
# ============================================================

def chat(
    message,
    system_prompt,
):
    """
    Normal Gradio handler.

    IMPORTANT:
    This does NOT call the dummy ZeroGPU function.
    """

    if not message or not str(message).strip():

        return "Please enter a message."

    try:

        return waamu_generate(
            message=str(message),
            system_prompt=str(system_prompt or ""),
        )

    except Exception as error:

        print(
            "Generation error:",
            repr(error),
            flush=True,
        )

        return (
            "Sorry, an internal error occurred "
            "while generating the response."
        )


# ============================================================
# UI
# ============================================================

with gr.Blocks(
    title=APP_NAME,
) as demo:

    gr.Markdown(
        """
# 🟣 WAAMU AI

### Qwen3-1.7B • Internal Reasoning • CPU
"""
    )

    gr.Markdown(
        """
WAAMU AI uses internal reasoning for difficult questions,
but only the final answer is displayed.
"""
    )

    system_prompt = gr.Textbox(
        label="System Prompt",
        placeholder=(
            "Example: You are a Class 8 teacher. "
            "Explain everything simply and briefly."
        ),
        lines=6,
    )

    message = gr.Textbox(
        label="Message",
        placeholder="Ask WAAMU AI anything...",
        lines=6,
    )

    send = gr.Button(
        "Send",
        variant="primary",
    )

    output = gr.Textbox(
        label="WAAMU AI",
        lines=12,
        interactive=False,
    )

    # --------------------------------------------------------
    # Button
    # --------------------------------------------------------

    send.click(
        fn=chat,
        inputs=[
            message,
            system_prompt,
        ],
        outputs=output,
    )

    # --------------------------------------------------------
    # Enter key
    # --------------------------------------------------------

    message.submit(
        fn=chat,
        inputs=[
            message,
            system_prompt,
        ],
        outputs=output,
    )


# ============================================================
# LAUNCH
# ============================================================

print("Launching Gradio...", flush=True)

demo.launch(
    server_name="0.0.0.0",
    server_port=7860,
    show_error=True,

    theme=gr.themes.Base(
        primary_hue="violet",
        neutral_hue="slate",
    ),
)