backend/app/memory/router.py

"""
Memory router — persistent AI chat history, user preferences (theme + language).
All endpoints require JWT authentication.
"""
from typing import Optional
from fastapi import APIRouter, Depends, Query
from sqlalchemy.orm import Session
from pydantic import BaseModel

from app.database.database import get_db
from app.database.models import (
    User, ChatSession, ChatMessage, UserPreference, generate_uuid
)
from app.auth.router import get_current_user

router = APIRouter(prefix="/api/memory", tags=["Memory"])


# ─── Schemas ──────────────────────────────────────────────────────────────────
class SaveMessageRequest(BaseModel):
    session_id: Optional[str] = None
    role: str                          # user | assistant
    content: str
    session_title: Optional[str] = None


class PreferenceUpdate(BaseModel):
    theme: Optional[str] = None        # dark | light
    language: Optional[str] = None     # en | hi | mr


# ─── Helpers ──────────────────────────────────────────────────────────────────
def _get_or_create_prefs(db: Session, user_id: str) -> UserPreference:
    prefs = db.query(UserPreference).filter(UserPreference.user_id == user_id).first()
    if not prefs:
        prefs = UserPreference(id=generate_uuid(), user_id=user_id)
        db.add(prefs)
        db.commit()
        db.refresh(prefs)
    return prefs


# ─── Chat history ─────────────────────────────────────────────────────────────
@router.get("/history")
def get_history(
    session_id: Optional[str] = None,
    limit: int = Query(default=50, ge=1, le=200),
    current_user: User = Depends(get_current_user),
    db: Session = Depends(get_db),
):
    """Return persisted chat messages for the current user."""
    query = db.query(ChatMessage).filter(ChatMessage.user_id == current_user.id)
    if session_id:
        query = query.filter(ChatMessage.session_id == session_id)
    messages = query.order_by(ChatMessage.created_at.asc()).limit(limit).all()

    # Also return sessions list
    sessions = (
        db.query(ChatSession)
        .filter(ChatSession.user_id == current_user.id)
        .order_by(ChatSession.updated_at.desc())
        .limit(20)
        .all()
    )

    return {
        "messages": [
            {
                "id": m.id,
                "session_id": m.session_id,
                "role": m.role,
                "content": m.content,
                "created_at": m.created_at.isoformat() if m.created_at else None,
            }
            for m in messages
        ],
        "sessions": [
            {
                "id": s.id,
                "title": s.title,
                "created_at": s.created_at.isoformat() if s.created_at else None,
                "updated_at": s.updated_at.isoformat() if s.updated_at else None,
            }
            for s in sessions
        ],
        "total": len(messages),
    }


@router.post("/save")
def save_message(
    req: SaveMessageRequest,
    current_user: User = Depends(get_current_user),
    db: Session = Depends(get_db),
):
    """Persist a chat message to the database."""
    # Get or create session
    session_id = req.session_id
    if not session_id:
        # Create a new session
        session = ChatSession(
            id=generate_uuid(),
            user_id=current_user.id,
            title=req.session_title or req.content[:40] + ("..." if len(req.content) > 40 else ""),
        )
        db.add(session)
        db.flush()
        session_id = session.id
    else:
        session = db.query(ChatSession).filter(
            ChatSession.id == session_id,
            ChatSession.user_id == current_user.id,
        ).first()
        if not session:
            session = ChatSession(
                id=session_id,
                user_id=current_user.id,
                title=req.session_title or "Chat",
            )
            db.add(session)
            db.flush()

    msg = ChatMessage(
        id=generate_uuid(),
        user_id=current_user.id,
        session_id=session_id,
        role=req.role,
        content=req.content,
    )
    db.add(msg)
    db.commit()
    db.refresh(msg)

    return {
        "id": msg.id,
        "session_id": session_id,
        "role": msg.role,
        "content": msg.content,
        "created_at": msg.created_at.isoformat() if msg.created_at else None,
    }


@router.delete("/clear")
def clear_history(
    session_id: Optional[str] = None,
    current_user: User = Depends(get_current_user),
    db: Session = Depends(get_db),
):
    """Clear chat history — optionally scoped to a session."""
    query = db.query(ChatMessage).filter(ChatMessage.user_id == current_user.id)
    if session_id:
        query = query.filter(ChatMessage.session_id == session_id)
        db.query(ChatSession).filter(
            ChatSession.id == session_id,
            ChatSession.user_id == current_user.id,
        ).delete()
    deleted = query.delete()
    db.commit()
    return {"deleted": deleted, "message": "History cleared."}


# ─── Preferences ──────────────────────────────────────────────────────────────
@router.get("/preferences")
def get_preferences(
    current_user: User = Depends(get_current_user),
    db: Session = Depends(get_db),
):
    prefs = _get_or_create_prefs(db, current_user.id)
    return {"theme": prefs.theme, "language": prefs.language}


@router.patch("/preferences")
def update_preferences(
    req: PreferenceUpdate,
    current_user: User = Depends(get_current_user),
    db: Session = Depends(get_db),
):
    prefs = _get_or_create_prefs(db, current_user.id)
    if req.theme in ("dark", "light"):
        prefs.theme = req.theme
    if req.language in ("en", "hi", "mr"):
        prefs.language = req.language
    db.commit()
    db.refresh(prefs)
    return {"theme": prefs.theme, "language": prefs.language}