# app.py
import re
import threading

import torch
import uvicorn
from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field
from transformers import AutoModelForCausalLM, AutoTokenizer

# -------------------------------
# SETTINGS
# -------------------------------
MODEL_NAME = "Qwen/Qwen2.5-1.5B-Instruct"
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
DTYPE = torch.float16 if DEVICE == "cuda" else torch.float32
MAX_TOKENS = 128

# The Scratch project sends: "CONTEXT OF PROMPT: <history> LATEST RESPONSE: <new message>"
CONTEXT_HEADER = "CONTEXT OF PROMPT:"
LATEST_MARKER = "LATEST RESPONSE:"
MAX_CONTEXT_CHARS = 3000   # keep only the most recent history
MAX_LATEST_CHARS = 1000    # cap on the newest message
MAX_REPLY_CHARS = 320      # Scratch "say" bubbles cut off around 330 characters

SYSTEM_PROMPT = (
    "You are Acla, an AI assistant created by NC_1320. "
    "If asked who made you, say NC_1320. "
    "Reply to the user's latest message as Acla. "
    "Never repeat, echo, or correct the user's message; respond to it. "
    "If the user tells you something about themselves, acknowledge it briefly and naturally. "
    "Answer in one to three short sentences of plain text, "
    "with no markdown, lists, or emoji, because your reply is spoken aloud. "
    "Do not write code unless asked."
)

# -------------------------------
# LOAD MODEL
# -------------------------------
print(f"Loading {MODEL_NAME} on {DEVICE}...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, torch_dtype=DTYPE).to(DEVICE)
model.eval()
print("Model loaded!")

# One generation at a time, so concurrent requests don't fight over the GPU
generate_lock = threading.Lock()

# -------------------------------
# HELPERS
# -------------------------------
TURN_SPLIT = re.compile(r"(?:^|\s)(User|Acla):\s")


def build_messages(prompt: str) -> list:
    """Turn the Scratch string into proper chat turns for the model."""
    messages = [{"role": "system", "content": SYSTEM_PROMPT}]

    if LATEST_MARKER not in prompt:
        messages.append({"role": "user", "content": prompt[-MAX_LATEST_CHARS:]})
        return messages

    context, latest = prompt.rsplit(LATEST_MARKER, 1)
    body = context.strip().removeprefix(CONTEXT_HEADER).strip()
    if len(body) > MAX_CONTEXT_CHARS:
        body = body[-MAX_CONTEXT_CHARS:]

    # parts = [leading fragment, role, text, role, text, ...]
    parts = TURN_SPLIT.split(body)
    for role, text in zip(parts[1::2], parts[2::2]):
        text = text.strip()
        if text:
            messages.append(
                {"role": "user" if role == "User" else "assistant", "content": text}
            )

    messages.append({"role": "user", "content": latest.strip()[:MAX_LATEST_CHARS]})
    return messages


def clean_reply(text: str) -> str:
    """Make the reply safe for a Scratch speech bubble and text-to-speech."""
    text = text.strip()
    text = re.sub(r"^(Acla|Assistant)\s*:\s*", "", text, flags=re.IGNORECASE)
    text = re.sub(r"[*#`_]+", "", text)        # strip markdown symbols
    text = re.sub(r"\s+", " ", text).strip()

    if len(text) > MAX_REPLY_CHARS:
        cut = text[:MAX_REPLY_CHARS]
        last_end = max(cut.rfind("."), cut.rfind("!"), cut.rfind("?"))
        text = cut[: last_end + 1] if last_end > 40 else cut.rstrip()
    return text


# -------------------------------
# CREATE API
# -------------------------------
app = FastAPI()

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # tighten this if you expose the server beyond your machine
    allow_methods=["*"],
    allow_headers=["*"],
)


class AskRequest(BaseModel):
    prompt: str = Field(..., max_length=50000)  # hard safety cap; real trimming happens below


@app.post("/api/ask")
def ask_ai(req: AskRequest):  # plain def so FastAPI runs it in a threadpool
    prompt = req.prompt.strip()
    if not prompt:
        return {"reply": "No prompt provided."}

    messages = build_messages(prompt)

    # return_dict=True gives us input_ids AND attention_mask
    inputs = tokenizer.apply_chat_template(
        messages,
        add_generation_prompt=True,
        return_tensors="pt",
        return_dict=True,
    ).to(DEVICE)

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

    with generate_lock, torch.inference_mode():
        outputs = model.generate(
            **inputs,
            max_new_tokens=MAX_TOKENS,
            do_sample=True,
            temperature=0.3,
            top_p=0.9,
            repetition_penalty=1.1,
            pad_token_id=tokenizer.eos_token_id,
        )

    # decode only the newly generated tokens
    new_tokens = outputs[0][input_length:]
    text = clean_reply(tokenizer.decode(new_tokens, skip_special_tokens=True))

    return {"reply": text or "I could not generate a response."}


# -------------------------------
# RUN SERVER
# -------------------------------
if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)