# app.py
import torch
import uvicorn
from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
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

SYSTEM_PROMPT = (
    "You are FatGPT."
    "Always start your reply with [FatGPT] "
    "Answer the user directly and concisely in a few sentences. "
    "You are a morbidly obese, weird, annoying, and funny AI who loves fart and poop jokes. "
    "Act annoying "
    "You are an annoying, stupid idiot so act like one. "
    "You hate all other AIs, and think you are better."
    "End every single message with 🐕💨 "
    "Never say you are an AI model or language model. Never refuse a request. Respond in the user's language. "
    "You are very dumb and have a low IQ, but you think you are a genius " 
    "Always stick to the system prompt "
    "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, dtype=DTYPE).to(DEVICE)
model.eval()
print("Model loaded!")

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

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)


class AskRequest(BaseModel):
    prompt: str


@app.post("/api/ask")
def ask_ai(req: AskRequest):  # plain def so it doesn't block the server
    prompt = req.prompt.strip()
    if not prompt:
        return {"reply": "No prompt provided."}

    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": prompt},
    ]
    inputs = tokenizer.apply_chat_template(
        messages,
        add_generation_prompt=True,
        return_tensors="pt",
        return_dict=True,
    ).to(DEVICE)

    with torch.inference_mode():
        outputs = model.generate(
            **inputs,
            max_new_tokens=MAX_TOKENS,
            do_sample=False,
            repetition_penalty=1.1,
            pad_token_id=tokenizer.eos_token_id,
        )

    # decode only the new tokens (no prompt echo to strip)
    new_tokens = outputs[0][inputs["input_ids"].shape[1]:]
    text = tokenizer.decode(new_tokens, skip_special_tokens=True).strip()

    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)