pages/chat.py

"""Chat tab: message display with citations and mock RAG responses."""

import uuid
from datetime import datetime
from state import UserData, Message, get_active_notebook
from services.rag_engine import rag_answer


FILE_TYPE_ICONS = {
    "pdf": "📕", "pptx": "📊", "txt": "📝", "url": "🌐", "youtube": "🎬",
}


def format_chatbot_messages(state: UserData) -> list[dict]:
    """Convert notebook messages to gr.Chatbot format with embedded citations."""
    nb = get_active_notebook(state)
    if not nb or not nb.messages:
        return []

    formatted = []
    for msg in nb.messages:
        content = msg.content
        if msg.role == "assistant" and msg.citations:
            chips = ""
            for c in msg.citations:
                chips += f'<span class="citation-chip">📄 {c["source"]} · p.{c["page"]}</span>'
            content += f"\n\n<details><summary>View cited passages</summary>\n\n{chips}\n\n</details>"
        formatted.append({"role": msg.role, "content": content})
    return formatted


def render_no_sources_warning(state: UserData) -> str:
    nb = get_active_notebook(state)
    if not nb or len(nb.sources) == 0:
        return (
            '<div style="padding:14px 20px; background:rgba(234,179,8,0.08); '
            'border:1px solid rgba(234,179,8,0.2); border-radius:12px; color:#d4a017; '
            'font-size:0.9rem; margin-bottom:16px;">'
            'Upload sources in the <strong>Sources</strong> tab to start chatting with your documents.'
            '</div>'
        )
    return ""


def handle_chat_submit(message: str, state: UserData) -> tuple[UserData, list[dict], str, str]:
    """Handle user sending a chat message. Returns (state, chatbot_messages, textbox_value, warning_html)."""
    if not message or not message.strip():
        return state, format_chatbot_messages(state), "", render_no_sources_warning(state)

    nb = get_active_notebook(state)
    if not nb:
        return state, [], "", ""
    if len(nb.sources) == 0:
        return state, format_chatbot_messages(state), "", render_no_sources_warning(state)

    # Build chat history BEFORE appending the new message (to avoid double-counting)
    chat_history = [{"role": m.role, "content": m.content} for m in nb.messages]

    # Add user message
    user_msg = Message(
        id=str(uuid.uuid4()),
        role="user",
        content=message.strip(),
        citations=[],
        created_at=datetime.now().isoformat(),
    )
    nb.messages.append(user_msg)

    # Get actual response
    response = rag_answer(message.strip(), nb.id, chat_history=chat_history)

    # Add assistant message
    assistant_msg = Message(
        id=str(uuid.uuid4()),
        role="assistant",
        content=response["content"],
        citations=response["citations"],
        created_at=datetime.now().isoformat(),
    )
    nb.messages.append(assistant_msg)

    return state, format_chatbot_messages(state), "", render_no_sources_warning(state)


def handle_clear_chat(state: UserData) -> tuple[UserData, list[dict], str]:
    """Clear all messages from the active notebook."""
    nb = get_active_notebook(state)
    if nb:
        nb.messages = []
    return state, [], render_no_sources_warning(state)