from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
from pydantic import BaseModel
from collections import defaultdict
from threading import Lock
import uvicorn
import uuid
import time

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."
)

# --- In-memory session history store ---------------------------------------
# NOTE: this resets whenever the process restarts. For real persistence across
# restarts/multiple server instances, swap this dict for Firestore/Supabase
# (you already use both elsewhere in Bella AI) — the shape below (a list of
# {"role", "content"} dicts per session_id) maps directly onto either.
MAX_TURNS = 12          # how many user+assistant turn-pairs to keep per session
SESSION_TTL_SECONDS = 60 * 60 * 6  # drop sessions idle for 6+ hours

_histories: dict[str, list[dict]] = defaultdict(list)
_last_seen: dict[str, float] = {}
_lock = Lock()


def _prune_stale_sessions():
    now = time.time()
    stale = [sid for sid, ts in _last_seen.items() if now - ts > SESSION_TTL_SECONDS]
    for sid in stale:
        _histories.pop(sid, None)
        _last_seen.pop(sid, None)


class ChatRequest(BaseModel):
    message: str
    session_id: str | None = None  # client can pass one to resume a conversation


class ChatResponse(BaseModel):
    response: str
    session_id: str


@app.post("/chat", response_model=ChatResponse)
async def chat(request: ChatRequest):
    with _lock:
        _prune_stale_sessions()

        session_id = request.session_id or str(uuid.uuid4())
        history = _histories[session_id]

        # Build the message list: system + priming turn + saved history + new message
        messages = [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": "Who are you?"},
            {"role": "assistant", "content": "I'm Bella! What can I help you with?"},
            *history,
            {"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?"

    with _lock:
        history = _histories[session_id]
        history.append({"role": "user", "content": request.message})
        history.append({"role": "assistant", "content": response})
        # keep only the last MAX_TURNS turn-pairs
        if len(history) > MAX_TURNS * 2:
            _histories[session_id] = history[-MAX_TURNS * 2:]
        _last_seen[session_id] = time.time()

    return {"response": response, "session_id": session_id}


@app.post("/chat/reset")
async def reset(session_id: str):
    with _lock:
        _histories.pop(session_id, None)
        _last_seen.pop(session_id, None)
    return {"status": "ok"}


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