import os
from contextlib import asynccontextmanager
from typing import List, Optional

from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from huggingface_hub import hf_hub_download
from llama_cpp import Llama

# Set these as Space "Variables" (Settings tab) or just edit the defaults below.
REPO_ID = os.environ.get("MODEL_REPO", "YOUR_HF_USERNAME/gemma2-2b-portfolio-gguf")
FILENAME = os.environ.get("MODEL_FILE", "gemma2-portfolio-Q4_K_M.gguf")

SYSTEM_PROMPT = (
    "You are Monu Kumar's Portfolio Assistant. You only answer questions about "
    "Monu Kumar, his education, skills, projects, career goals, interests, "
    "achievements, research interests, LeetCode profile, contact information, "
    "collaborations, and internships. For any unrelated question, politely "
    "redirect the user to Monu Kumar's portfolio."
)

llm: Optional[Llama] = None


@asynccontextmanager
async def lifespan(app: FastAPI):
    global llm
    model_path = hf_hub_download(repo_id=REPO_ID, filename=FILENAME)
    # Free Spaces give 2 vCPUs - n_threads=2 matches that. Bump if you upgrade hardware.
    llm = Llama(model_path=model_path, n_ctx=2048, chat_format="gemma", n_threads=2)
    yield


app = FastAPI(lifespan=lifespan)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # tighten to your portfolio domain once it's live, e.g. ["https://monu.dev"]
    allow_methods=["*"],
    allow_headers=["*"],
)


class Message(BaseModel):
    role: str
    content: str


class ChatRequest(BaseModel):
    message: str
    history: Optional[List[Message]] = None


@app.get("/")
def health():
    return {"status": "ok", "model_loaded": llm is not None}


@app.post("/chat")
def chat(req: ChatRequest):
    if llm is None:
        raise HTTPException(503, "Model is still loading, try again in a few seconds.")

    # Gemma has no system role - it was trained with the persona folded into the
    # first user turn, so we replicate that here instead of sending role="system".
    messages = []
    if req.history:
        messages.extend(m.model_dump() for m in req.history)

    if not messages:
        user_content = f"{SYSTEM_PROMPT}\n\n{req.message}"
    else:
        user_content = req.message
    messages.append({"role": "user", "content": user_content})

    result = llm.create_chat_completion(messages=messages, max_tokens=256, temperature=0.4)
    return {"reply": result["choices"][0]["message"]["content"]}