ui/callbacks.py

"""
Gradio callback functions.

These call chat_service / auth / history directly (in-process) rather than
going over HTTP, since the Gradio Blocks app is mounted into the same Flask
process (see app.py). Because Gradio callbacks run outside of a normal
Flask request context, every function that touches the database wraps its
body in `with _flask_app.app_context():`.

Auth state for the UI is NOT the Flask session cookie (that's only for the
REST API in auth/routes.py) — it's tracked per-browser-tab in a gr.State
dict: {"user_id": int, "email": str} or None when logged out.
"""

import time

from auth.db import db
from auth.models import User
from chat_service import handle_chat_message
from history import service as history_service
from logs.logger import get_logger

logger = get_logger(__name__)

_flask_app = None


def init_app(flask_app):
    """Called once from app.py so callbacks can open an app context."""
    global _flask_app
    _flask_app = flask_app


# ---------------------------------------------------------------------------
# Auth
# ---------------------------------------------------------------------------

def do_signup(email: str, password: str):
    """Returns (auth_state, status_message, is_error)."""
    email = (email or "").strip().lower()
    password = password or ""

    if not email or "@" not in email:
        return None, "Please enter a valid email address.", True
    if len(password) < 8:
        return None, "Password must be at least 8 characters.", True

    with _flask_app.app_context():
        if User.query.filter_by(email=email).first() is not None:
            return None, "An account with this email already exists.", True

        user = User(email=email)
        user.set_password(password)
        db.session.add(user)
        db.session.commit()
        auth_state = {"user_id": user.id, "email": user.email}

    return auth_state, f"Account created. Welcome, {email}!", False


def do_login(email: str, password: str):
    """Returns (auth_state, status_message, is_error)."""
    email = (email or "").strip().lower()
    password = password or ""

    with _flask_app.app_context():
        user = User.query.filter_by(email=email).first()
        if user is None or not user.check_password(password):
            return None, "Invalid email or password.", True
        auth_state = {"user_id": user.id, "email": user.email}

    return auth_state, f"Logged in as {email}.", False


def do_logout():
    return None, "Logged out.", False


# ---------------------------------------------------------------------------
# Chat
# ---------------------------------------------------------------------------

def send_message(user_message: str, chat_display: list, auth_state, session_state, anon_history: list):
    """
    Args:
        user_message: raw text from the input box.
        chat_display: current [{"role": ..., "content": ...}, ...] shown in
                       gr.Chatbot(type="messages").
        auth_state: {"user_id", "email"} or None.
        session_state: current chat_session id (logged-in only) or None.
        anon_history: client-side rolling history used only when logged out.

    Returns:
        (chat_display, session_state, anon_history, cleared_textbox)
    """
    user_message = (user_message or "").strip()
    if not user_message:
        return chat_display, session_state, anon_history, ""

    chat_display = chat_display + [{"role": "user", "content": user_message}]

    is_logged_in = bool(auth_state)

    with _flask_app.app_context():
        if is_logged_in:
            user = User.query.get(auth_state["user_id"])
            result = handle_chat_message(
                query=user_message,
                user=user,
                session_id=session_state,
            )
            session_state = result["session_id"]
        else:
            result = handle_chat_message(
                query=user_message,
                user=None,
                anonymous_history=anon_history,
            )
            anon_history = (anon_history or []) + [
                {"role": "user", "content": user_message},
                {"role": "assistant", "content": result["response"]},
            ]

    chat_display = chat_display + [{"role": "assistant", "content": result["response"]}]
    return chat_display, session_state, anon_history, ""


def stream_last_response(chat_display: list):
    """
    Fakes token-by-token streaming for the assistant's last message so the
    UI feels responsive even though llm.generate() currently returns the
    full string at once. Yields progressively longer chat_display copies.
    """
    if not chat_display or chat_display[-1]["role"] != "assistant":
        yield chat_display
        return

    full_text = chat_display[-1]["content"]
    prefix_display = chat_display[:-1]

    step = max(1, len(full_text) // 60)
    for i in range(0, len(full_text) + step, step):
        partial = full_text[:i]
        yield prefix_display + [{"role": "assistant", "content": partial}]
        time.sleep(0.012)

    yield chat_display


def new_chat():
    """Resets the chat window + clears the active session id."""
    return [], None, []


# ---------------------------------------------------------------------------
# History panel
# ---------------------------------------------------------------------------

def refresh_sessions(auth_state):
    """Returns gr.update(...) choices for the session list, newest first."""
    import gradio as gr

    if not auth_state:
        return gr.update(choices=[], value=None)

    with _flask_app.app_context():
        sessions = history_service.list_sessions(auth_state["user_id"])

    choices = [(s["title"] or f"Session {s['id']}", s["id"]) for s in sessions]
    return gr.update(choices=choices, value=None)


def open_session(session_id, auth_state):
    """Loads a past session's full transcript into the chat window."""
    if not auth_state or session_id is None:
        return [], None

    with _flask_app.app_context():
        full_session = history_service.get_full_session(session_id, auth_state["user_id"])

    if full_session is None:
        return [], None

    chat_display = [
        {"role": m["role"], "content": m["content"]} for m in full_session["messages"]
    ]
    return chat_display, session_id


def delete_session(session_id, auth_state):
    """Deletes a session and returns a cleared chat window if it was active."""
    if not auth_state or session_id is None:
        return [], None

    with _flask_app.app_context():
        history_service.delete_session(session_id, auth_state["user_id"])

    return [], None