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

app = FastAPI()

# Allow requests from any origin (so your HTML page can call this API)
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

# Load model + tokenizer
model_id = "Qwen/Qwen2.5-0.5B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id)

chat_pipeline = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
)

SYSTEM_PROMPT = (
    "You are Bella, a witty and friendly female AI assistant. "
    "Always respond as Bella, never use any other name. "
    "Keep your replies short, casual, and to the point."
)


class ChatRequest(BaseModel):
    message: str


@app.post("/chat")
async def chat(request: ChatRequest):
    # Priming turn anchors the persona much more strongly for small models
    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": "Who are you?"},
        {"role": "assistant", "content": "I'm Bella! What can I help you with?"},
        {"role": "user", "content": request.message},
    ]

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

    outputs = chat_pipeline(
        prompt,
        max_new_tokens=256,
        do_sample=True,
        temperature=0.2,
        top_p=0.9,
        repetition_penalty=1.2,
        pad_token_id=tokenizer.eos_token_id,
    )

    generated_text = outputs[0]["generated_text"]

    # Strip off the prompt, leaving only the new reply
    response = generated_text[len(prompt):].strip()

    # Safety net: cut off if the model starts hallucinating a new turn
    for stop_token in ["<|im_start|>", "<|im_end|>", "User:", "system"]:
        if stop_token in response:
            response = response.split(stop_token)[0].strip()

    if not response:
        response = "Hmm, I didn't quite catch that — can you rephrase?"

    return {"response": response}


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