ui/chat.py

"""Streamlit UI: chat streams; audits stream in background; button reveals."""

from __future__ import annotations

import uuid
from datetime import timedelta

import streamlit as st

from agent.background import (
    schedule_audit,
    schedule_missing_audits,
    sync_audits,
)
from agent.chat import chat_reply_stream
from agent.state import ChatTurn, MessageAudit
from config import (
    AUDIT_MODEL,
    AUDIT_USE_WEB_SEARCH,
    CHAT_MODEL,
    SEARCH_MODEL,
    get_api_key,
)
from ui.render import conversation_table_html

APP_TITLE = "Anti-Hallucination Chat"

_COMPOSER_CSS = """
<style>
/* Tighten the pinned bottom composer row */
[data-testid="stBottomBlockContainer"] [data-testid="stHorizontalBlock"] {
  align-items: center;
  gap: 0.55rem;
}
[data-testid="stBottomBlockContainer"] [data-testid="column"]:last-child button {
  width: 100%;
  white-space: nowrap;
  border-radius: 0.75rem;
  min-height: 2.75rem;
}
</style>
"""


def _ensure_state() -> None:
    if "turns" not in st.session_state:
        st.session_state.turns: list[ChatTurn] = []
    if "audits_by_id" not in st.session_state:
        st.session_state.audits_by_id: dict[str, MessageAudit] = {}
    if "show_audit_view" not in st.session_state:
        st.session_state.show_audit_view = False
    if "audit_session_id" not in st.session_state:
        st.session_state.audit_session_id = str(uuid.uuid4())
    if "streaming_now" not in st.session_state:
        st.session_state.streaming_now = False
    if "partial_assistant" not in st.session_state:
        st.session_state.partial_assistant = ""


def _sync_background() -> bool:
    return sync_audits(
        st.session_state.audit_session_id,
        st.session_state.audits_by_id,
    )


def _audits_active() -> bool:
    return any(
        audit.get("status") in ("running", "pending")
        for audit in st.session_state.audits_by_id.values()
    )


def _render_turn_table(*, streaming_assistant: str | None = None) -> None:
    st.markdown(
        conversation_table_html(
            st.session_state.turns,
            st.session_state.audits_by_id,
            show_audit=bool(st.session_state.show_audit_view),
            streaming_assistant=streaming_assistant,
        ),
        unsafe_allow_html=True,
    )


def _toggle_audit_view() -> None:
    """Toggle columns 2–3 (claims / reasoning) on or off."""
    opening = not st.session_state.show_audit_view
    if opening:
        schedule_missing_audits(
            st.session_state.audit_session_id,
            st.session_state.turns,
        )
        _sync_background()
    st.session_state.show_audit_view = opening
    st.rerun()


def _check_button_label() -> str:
    if st.session_state.show_audit_view:
        if _audits_active():
            return "Hide audit (streaming…)"
        return "Hide audit"
    if _audits_active():
        return "Check (ready)"
    return "Hallucination check"


def _finalize_partial_assistant() -> None:
    """If a prior run was interrupted mid-stream, keep whatever text we got."""
    if not st.session_state.get("streaming_now"):
        return
    partial = (st.session_state.get("partial_assistant") or "").strip()
    st.session_state.streaming_now = False
    st.session_state.partial_assistant = ""
    if not partial:
        return
    if st.session_state.turns and st.session_state.turns[-1]["role"] == "assistant":
        return
    assistant_turn: ChatTurn = {
        "id": str(uuid.uuid4()),
        "role": "assistant",
        "content": partial,
    }
    st.session_state.turns.append(assistant_turn)
    schedule_audit(
        st.session_state.audit_session_id,
        assistant_turn,
        list(st.session_state.turns),
    )


def _render_composer() -> str | None:
    """Message input + check toggle, pinned to the viewport bottom.

    While a new prompt is being submitted (this run will stream), the check
    button is omitted so it cannot interrupt the assistant reply.
    """
    st.markdown(_COMPOSER_CSS, unsafe_allow_html=True)
    with st.bottom:
        input_col, check_col = st.columns([5.2, 1.55], vertical_alignment="center")
        with input_col:
            prompt = st.chat_input("Message…", key="ah_chat_input")
        with check_col:
            if prompt:
                # Hide check for the streaming run — keeps layout slot empty.
                st.empty()
            elif st.button(
                _check_button_label(),
                type="primary",
                key="ah_hallucination_check",
                use_container_width=True,
            ):
                _toggle_audit_view()
    return prompt


@st.fragment(run_every=timedelta(seconds=1))
def _audit_refresh_fragment() -> None:
    """Always-mounted fragment so run_every never targets a missing id.

    Polls background audit snapshots and streams bullets into the reveal panel.
    """
    active = _sync_background()
    if not st.session_state.show_audit_view:
        return
    if st.session_state.get("streaming_now"):
        return

    _render_turn_table()
    if active or _audits_active():
        st.caption("⏳ Audit stream updating…")


def _stream_assistant_reply() -> str | None:
    """Stream the next assistant reply into the table (no chat-bubble UI)."""
    table_slot = st.empty()
    err_slot = st.empty()
    chunks: list[str] = []

    def _render_partial(text: str) -> None:
        st.session_state.partial_assistant = text
        table_slot.markdown(
            conversation_table_html(
                st.session_state.turns,
                st.session_state.audits_by_id,
                show_audit=bool(st.session_state.show_audit_view),
                streaming_assistant=text or "…",
            ),
            unsafe_allow_html=True,
        )

    _render_partial("")
    try:
        for piece in chat_reply_stream(list(st.session_state.turns)):
            chunks.append(str(piece))
            _render_partial("".join(chunks))
    except Exception as exc:  # noqa: BLE001
        err_slot.error(f"Chat failed: {exc}")
        return None

    text = "".join(chunks).strip()
    if not text:
        err_slot.error("Chat failed: empty model response")
        return None
    return text


def run_chat() -> None:
    st.set_page_config(
        page_title=APP_TITLE,
        page_icon="💬",
        layout="wide",
    )
    st.title(APP_TITLE)
    search_note = (
        f"search {SEARCH_MODEL}" if AUDIT_USE_WEB_SEARCH else "web search off"
    )
    st.caption(
        f"v5 · chat {CHAT_MODEL} · audit {AUDIT_MODEL} · {search_note}"
    )

    if not get_api_key():
        st.error(
            "GOOGLE_API_KEY is missing. "
            "Set it in a local `.env` file or as a Hugging Face Space secret."
        )
        st.stop()

    _ensure_state()
    _sync_background()
    # If a previous run was cut off mid-stream, keep any partial assistant text.
    _finalize_partial_assistant()

    prompt = _render_composer()
    if prompt:
        st.session_state.streaming_now = True
        st.session_state.partial_assistant = ""

    # Keep this fragment mounted on every run — conditional fragments break the
    # Cursor Cloud preview when run_every targets a removed fragment id.
    _audit_refresh_fragment()

    if prompt:
        user_turn: ChatTurn = {
            "id": str(uuid.uuid4()),
            "role": "user",
            "content": prompt,
        }
        st.session_state.turns.append(user_turn)
        schedule_audit(
            st.session_state.audit_session_id,
            user_turn,
            list(st.session_state.turns),
        )

        text = _stream_assistant_reply()
        st.session_state.streaming_now = False
        st.session_state.partial_assistant = ""
        if text is None:
            st.session_state.turns.pop()
            return

        assistant_turn: ChatTurn = {
            "id": str(uuid.uuid4()),
            "role": "assistant",
            "content": text,
        }
        st.session_state.turns.append(assistant_turn)
        schedule_audit(
            st.session_state.audit_session_id,
            assistant_turn,
            list(st.session_state.turns),
        )
        st.rerun()
        return

    if not st.session_state.show_audit_view:
        _render_turn_table()