app.py

"""
Ms. Q backend - FastAPI app.
Run locally:   uvicorn app:app --reload --port 7860
Deployed on:   Hugging Face Spaces (Docker SDK), port 7860
"""
import uuid
from datetime import datetime, timezone

from fastapi import FastAPI, HTTPException, Depends, status
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import Response
from pydantic import BaseModel, EmailStr, Field

from config import settings
from database import users_collection, conversations_collection, documents_collection
from auth import hash_password, verify_password, create_access_token, get_current_user
from persona import MS_Q_SYSTEM_PROMPT
from llm import generate_reply
from tts import synthesize_speech
from embeddings import embed_texts
from rag import chunk_text, retrieve_relevant_chunks

app = FastAPI(title="Ms. Q API")

origins = [o.strip() for o in settings.ALLOWED_ORIGINS.split(",") if o.strip()]
app.add_middleware(
    CORSMiddleware,
    allow_origins=origins,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


# ---------- Schemas ----------
class RegisterRequest(BaseModel):
    name: str = Field(min_length=1, max_length=80)
    email: EmailStr
    password: str = Field(min_length=8)


class LoginRequest(BaseModel):
    email: EmailStr
    password: str


class TokenResponse(BaseModel):
    access_token: str
    token_type: str = "bearer"
    name: str


class ChatMessage(BaseModel):
    role: str  # "user" | "assistant"
    content: str


class ChatRequest(BaseModel):
    message: str = Field(min_length=1, max_length=4000)


class ChatResponse(BaseModel):
    reply: str
    used_web_search: bool = False
    used_documents: bool = False


class TTSRequest(BaseModel):
    text: str = Field(min_length=1, max_length=4000)


class DocumentUploadRequest(BaseModel):
    title: str = Field(min_length=1, max_length=200)
    content: str = Field(min_length=1, max_length=50000)


# ---------- Health ----------
@app.get("/")
async def root():
    return {"status": "Ms. Q is online"}


# ---------- Auth ----------
@app.post("/api/auth/register", response_model=TokenResponse)
async def register(body: RegisterRequest):
    existing = await users_collection().find_one({"email": body.email})
    if existing:
        raise HTTPException(status_code=400, detail="An account with this email already exists.")

    doc = {
        "name": body.name,
        "email": body.email,
        "password_hash": hash_password(body.password),
        "created_at": datetime.now(timezone.utc),
    }
    await users_collection().insert_one(doc)
    token = create_access_token(subject=body.email)
    return TokenResponse(access_token=token, name=body.name)


@app.post("/api/auth/login", response_model=TokenResponse)
async def login(body: LoginRequest):
    user = await users_collection().find_one({"email": body.email})
    if not user or not verify_password(body.password, user["password_hash"]):
        raise HTTPException(status_code=401, detail="Incorrect email or password.")
    token = create_access_token(subject=body.email)
    return TokenResponse(access_token=token, name=user["name"])


# ---------- Chat ----------
@app.post("/api/chat", response_model=ChatResponse)
async def chat(body: ChatRequest, current_user: dict = Depends(get_current_user)):
    email = current_user["email"]
    convo = await conversations_collection().find_one({"email": email})
    history = convo["messages"] if convo else []

    # Keep the last 20 turns so context doesn't grow unbounded
    trimmed_history = history[-20:]

    api_messages = [{"role": m["role"], "content": m["content"]} for m in trimmed_history]

    # Document RAG: pull in anything relevant the person has shared with her.
    # Fails soft (returns []) if no documents exist or the vector index isn't set up.
    relevant_chunks = await retrieve_relevant_chunks(email, body.message)
    if relevant_chunks:
        context_block = "\n---\n".join(relevant_chunks)
        llm_user_message = (
            "[Background from documents this person has shared with you - "
            f"use naturally if relevant]\n{context_block}\n\n"
            f"[Their actual message]\n{body.message}"
        )
    else:
        llm_user_message = body.message

    api_messages.append({"role": "user", "content": llm_user_message})

    try:
        reply_text, used_web_search = await generate_reply(MS_Q_SYSTEM_PROMPT, api_messages)
    except RuntimeError as e:
        raise HTTPException(status_code=500, detail=str(e))

    # Store the ORIGINAL message (not the RAG-augmented version) so chat
    # history stays clean and isn't re-injected with duplicate context later.
    new_messages = trimmed_history + [
        {"role": "user", "content": body.message, "at": datetime.now(timezone.utc)},
        {"role": "assistant", "content": reply_text, "at": datetime.now(timezone.utc)},
    ]

    await conversations_collection().update_one(
        {"email": email},
        {"$set": {"email": email, "messages": new_messages, "updated_at": datetime.now(timezone.utc)}},
        upsert=True,
    )

    return ChatResponse(
        reply=reply_text,
        used_web_search=used_web_search,
        used_documents=bool(relevant_chunks),
    )


@app.get("/api/chat/history")
async def chat_history(current_user: dict = Depends(get_current_user)):
    convo = await conversations_collection().find_one({"email": current_user["email"]})
    if not convo:
        return {"messages": []}
    return {"messages": [{"role": m["role"], "content": m["content"]} for m in convo["messages"]]}


# ---------- Text-to-speech (Ms. Q's voice) ----------
@app.post("/api/tts")
async def text_to_speech(body: TTSRequest, current_user: dict = Depends(get_current_user)):
    try:
        audio_bytes, mime_type = await synthesize_speech(body.text)
    except RuntimeError as e:
        raise HTTPException(status_code=500, detail=str(e))

    return Response(content=audio_bytes, media_type=mime_type)


# ---------- Document memory (RAG) ----------
@app.post("/api/documents")
async def upload_document(body: DocumentUploadRequest, current_user: dict = Depends(get_current_user)):
    chunks = chunk_text(body.content)
    if not chunks:
        raise HTTPException(status_code=400, detail="No content to store.")

    vectors = await embed_texts(chunks)

    doc_id = str(uuid.uuid4())
    now = datetime.now(timezone.utc)
    rows = [
        {
            "doc_id": doc_id,
            "email": current_user["email"],
            "doc_title": body.title,
            "chunk_index": i,
            "text": chunk,
            "embedding": vector,
            "created_at": now,
        }
        for i, (chunk, vector) in enumerate(zip(chunks, vectors))
    ]
    await documents_collection().insert_many(rows)

    return {"doc_id": doc_id, "title": body.title, "chunk_count": len(chunks)}


@app.get("/api/documents")
async def list_documents(current_user: dict = Depends(get_current_user)):
    pipeline = [
        {"$match": {"email": current_user["email"]}},
        {
            "$group": {
                "_id": "$doc_id",
                "title": {"$first": "$doc_title"},
                "chunk_count": {"$sum": 1},
                "created_at": {"$min": "$created_at"},
            }
        },
        {"$sort": {"created_at": -1}},
    ]
    cursor = documents_collection().aggregate(pipeline)
    results = await cursor.to_list(length=200)
    return {
        "documents": [
            {
                "doc_id": r["_id"],
                "title": r["title"],
                "chunk_count": r["chunk_count"],
                "created_at": r["created_at"],
            }
            for r in results
        ]
    }


@app.delete("/api/documents/{doc_id}")
async def delete_document(doc_id: str, current_user: dict = Depends(get_current_user)):
    result = await documents_collection().delete_many({"doc_id": doc_id, "email": current_user["email"]})
    if result.deleted_count == 0:
        raise HTTPException(status_code=404, detail="Document not found.")
    return {"deleted": True}