src/ui/_pages/main_interface.py

"""Main chat interface page for ChemGraph."""

import asyncio
import html
import json
import logging
import os
import pprint
import queue
import re
import threading
import uuid
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, Optional

import pandas as pd
import streamlit as st
from ase.io import read as ase_read

from chemgraph.agent.llm_agent import HumanInputRequired
from chemgraph.memory.schemas import SessionMessage
from chemgraph.memory.store import SessionStore
from chemgraph.models.supported_models import supported_argo_models
from chemgraph.schemas.ase_input import (
    get_available_calculator_names,
    get_default_calculator_name,
)
from chemgraph.utils.config_utils import (
    get_argo_user_from_nested_config,
    get_base_url_for_model_from_nested_config,
)
from chemgraph.utils.tool_mapping import (
    format_validation_failure_message,
    validate_completion,
)

from ui.agent_manager import initialize_agent
from ui.branding import LOGO_IMAGES, first_existing_asset
from ui.config import load_config
from ui.endpoint import check_local_model_endpoint
from ui.file_utils import (
    extract_log_dir_from_messages,
    find_latest_xyz_file_in_dir,
)
from ui.message_utils import (
    extract_messages_from_result,
    extract_molecular_structure,
    extract_xyz_from_report_html,
    find_html_filename,
    find_structure_in_messages,
    has_structure_signal,
    is_infrared_requested,
    normalize_latex_delimiters,
    normalize_message_content,
    split_markdown_latex_blocks,
    strip_viewer_from_report_html,
)
from ui.scientific_reminders import (
    QUERY_SPECIFICATION_HINT,
    render_chat_scientific_reminder,
)
from ui.session_utils import (
    conversation_entry_to_messages,
    generate_session_id,
    messages_from_result,
    session_to_conversation_history,
)
from ui.state import init_session_state
from ui.visualization import (
    STMOL_AVAILABLE,
    display_molecular_structure,
    visualize_trajectory,
)

# Re-use the constants from the configuration page
from ui._pages.configuration import normalize_workflow_name

logger = logging.getLogger(__name__)
_AGENT_RUN_LOCK = threading.Lock()


# ---------------------------------------------------------------------------
# Thin wrappers around config utilities
# ---------------------------------------------------------------------------


def _get_base_url_for_model(model_name: str, config: Dict[str, Any]) -> Optional[str]:
    """Resolve the configured base URL for a model.

    Parameters
    ----------
    model_name : str
        Selected model identifier.
    config : dict[str, Any]
        Nested UI configuration.

    Returns
    -------
    str or None
        Provider base URL, or ``None`` when not configured.
    """
    return get_base_url_for_model_from_nested_config(model_name, config)


def _initial_ui_log_root() -> str:
    """Return the root directory for per-chat UI artifacts."""
    env_log_dir = os.environ.get("CHEMGRAPH_LOG_DIR")
    if env_log_dir:
        return str(_ui_log_root_from_path(Path(env_log_dir).expanduser()))
    return str((Path.cwd() / "cg_logs").resolve())


def _ui_log_root_from_path(path: Path) -> Path:
    """Return the chat-log root even when given a stale session/query path."""
    current = path.resolve()
    while current.name.startswith(("query_", "session_", "ui_session_")):
        current = current.parent
    return current


def _ensure_chat_log_dir() -> str:
    """Create and activate a log directory owned by the current chat."""
    if st.session_state.get("ui_log_root"):
        st.session_state.ui_log_root = str(
            _ui_log_root_from_path(Path(st.session_state.ui_log_root).expanduser())
        )
    else:
        st.session_state.ui_log_root = _initial_ui_log_root()

    chat_log_dir = st.session_state.get("current_chat_log_dir")
    if chat_log_dir:
        chat_path = Path(chat_log_dir).expanduser().resolve()
        root_path = Path(st.session_state.ui_log_root).expanduser().resolve()
        if chat_path.parent != root_path:
            st.session_state.current_chat_log_dir = None
            chat_log_dir = None
    if not chat_log_dir:
        timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
        suffix = str(uuid.uuid4())[:8]
        chat_log_dir = str(
            Path(st.session_state.ui_log_root) / f"ui_session_{timestamp}_{suffix}"
        )
        st.session_state.current_chat_log_dir = chat_log_dir

    os.makedirs(chat_log_dir, exist_ok=True)
    os.environ["CHEMGRAPH_LOG_DIR"] = chat_log_dir
    return chat_log_dir


def _create_query_log_dir() -> str:
    """Create a per-query artifact directory under the active chat directory."""
    chat_log_dir = _ensure_chat_log_dir()
    query_idx = len(st.session_state.get("conversation_history", [])) + 1
    suffix = str(uuid.uuid4())[:8]
    query_log_dir = str(Path(chat_log_dir) / f"query_{query_idx:03d}_{suffix}")
    os.makedirs(query_log_dir, exist_ok=True)
    st.session_state.active_query_log_dir = query_log_dir
    return query_log_dir


def _begin_agent_run(*, from_pending: bool = False) -> bool:
    """Mark an agent run as active if no other run is in progress."""
    if st.session_state.get("agent_running") and not from_pending:
        return False
    if st.session_state.get("agent_run_lock_acquired"):
        return True
    if _AGENT_RUN_LOCK.locked():
        return False
    acquired = _AGENT_RUN_LOCK.acquire(blocking=False)
    if not acquired:
        return False
    st.session_state.agent_running = True
    st.session_state.agent_run_lock_acquired = True
    return True


def _agent_run_active() -> bool:
    """Return whether any Streamlit session has an active workflow run."""
    lock_active = _AGENT_RUN_LOCK.locked()
    has_pending_submission = st.session_state.get("pending_agent_submission") is not None
    if (
        st.session_state.get("agent_running")
        and not lock_active
        and not has_pending_submission
    ):
        # A browser refresh creates a fresh Streamlit session_state while the
        # process lock is the authoritative cross-session signal.
        st.session_state.agent_running = False
        st.session_state.agent_run_lock_acquired = False
    return bool(lock_active or st.session_state.get("agent_running"))


def _end_agent_run() -> None:
    """Clear the active-run marker and release the process-local run lock."""
    st.session_state.agent_running = False
    st.session_state.active_query_log_dir = None
    if st.session_state.get("agent_run_lock_acquired"):
        st.session_state.agent_run_lock_acquired = False
        try:
            _AGENT_RUN_LOCK.release()
        except RuntimeError:
            pass


def _session_id_from_url() -> Optional[str]:
    """Read the session id from the URL query string when present."""
    try:
        value = st.query_params.get("session")
    except Exception:
        return None
    if isinstance(value, list):
        return str(value[0]) if value else None
    return str(value) if value else None


def _set_session_id_in_url(session_id: Optional[str]) -> None:
    """Keep the active session id in the browser URL for refresh recovery."""
    try:
        if session_id:
            st.query_params["session"] = session_id
        else:
            st.query_params.pop("session", None)
    except Exception:
        pass


def _restore_agent_log_context(agent, old_agent_log_dir, old_env_log_dir) -> None:
    """Restore agent/env log state after a per-query run."""
    if agent is not None:
        agent.log_dir = old_agent_log_dir
    if st.session_state.get("current_chat_log_dir"):
        os.environ["CHEMGRAPH_LOG_DIR"] = st.session_state.current_chat_log_dir
    elif old_env_log_dir is not None:
        old_path = Path(old_env_log_dir).expanduser()
        if old_path.name.startswith("query_"):
            os.environ["CHEMGRAPH_LOG_DIR"] = str(_ui_log_root_from_path(old_path))
        else:
            os.environ["CHEMGRAPH_LOG_DIR"] = old_env_log_dir
    else:
        os.environ.pop("CHEMGRAPH_LOG_DIR", None)


def _resolve_structured_output_for_model(
    model_name: str, structured_output: bool
) -> tuple[bool, Optional[str]]:
    """Disable structured output for Argo models, including quick overrides.

    Parameters
    ----------
    model_name : str
        Selected model identifier.
    structured_output : bool
        Requested structured-output setting.

    Returns
    -------
    tuple[bool, str | None]
        Effective structured-output setting and optional warning message.
    """
    if model_name in supported_argo_models and structured_output:
        return (
            False,
            "Structured output is disabled for Argo models to avoid JSON parsing errors.",
        )
    return structured_output, None


# ---------------------------------------------------------------------------
# Page entry point
# ---------------------------------------------------------------------------


def render() -> None:
    """Render the Main Interface page."""
    init_session_state()
    _restore_session_from_url()

    config = st.session_state.config
    selected_model = config["general"]["model"]
    selected_workflow = normalize_workflow_name(config["general"]["workflow"])
    selected_output = config["general"]["output"]
    structured_output = config["general"]["structured"]
    generate_report = config["general"]["report"]
    human_supervised = config["general"].get("human_supervised", False)
    thread_id = config["general"]["thread"]

    # ----- Header -----
    logo_image = first_existing_asset(LOGO_IMAGES)
    if logo_image:
        st.image(logo_image, width=320)
    else:
        st.title("\U0001f9ea ChemGraph")

    st.markdown("""
    ChemGraph enables you to perform various **computational chemistry** tasks with
    natural-language queries using AI agents.
    """)

    # ----- Calculator availability sidebar -----
    _render_available_calculators_sidebar()
    _render_chat_controls()

    structured_output, ui_notice = _resolve_structured_output_for_model(
        selected_model, structured_output
    )
    st.session_state.ui_notice = ui_notice
    st.session_state.active_model = selected_model
    st.session_state.active_workflow = selected_workflow
    if ui_notice:
        st.info(ui_notice)
    render_chat_scientific_reminder()

    selected_base_url = _get_base_url_for_model(selected_model, config)
    endpoint_status = check_local_model_endpoint(selected_base_url)

    # ----- Session management sidebar -----
    _render_session_sidebar()

    # Reload config button
    if st.sidebar.button("\U0001f504 Reload Config", disabled=_agent_run_active()):
        st.session_state.config = load_config()
        st.success("\u2705 Configuration reloaded!")
        st.rerun()

    # ----- Auto-initialize agent -----
    _auto_initialize_agent(
        config,
        selected_model,
        selected_workflow,
        structured_output,
        selected_output,
        generate_report,
        human_supervised,
        selected_base_url,
    )

    # ----- Agent status sidebar -----
    _render_agent_status(selected_model, selected_workflow, thread_id, endpoint_status)

    # ----- Restore pending clarification from the visible history -----
    _restore_pending_interrupt_from_history()

    # ----- Conversation history -----
    _render_conversation_history(thread_id)

    # ----- Pending interrupt display -----
    _render_pending_interrupt()

    # ----- Example queries -----
    _render_example_queries(config, selected_model)

    # ----- Chat input (handles both normal queries and interrupt responses) -----
    is_interrupt = st.session_state.pending_human_question is not None
    agent_busy = _agent_run_active()
    if agent_busy:
        st.info("Agent is running. Please wait for the current task to finish.")

    prompt = st.chat_input(
        (
            "Type your response..."
            if is_interrupt
            else "Ask: molecule/reaction + property + calculator + conditions..."
        ),
        disabled=agent_busy,
    )

    # Check for example query submitted via button click
    if agent_busy:
        st.session_state.pop("_pending_example_query", None)
        prompt = None
    else:
        example_query = st.session_state.pop("_pending_example_query", None)
        if example_query:
            prompt = example_query

    if prompt:
        if is_interrupt:
            _handle_human_response(prompt, thread_id)
        else:
            _queue_query_submission(
                prompt, thread_id, endpoint_status, selected_base_url
            )

    _execute_pending_agent_submission(thread_id, endpoint_status, selected_base_url)


# ---------------------------------------------------------------------------
# Internal renderers
# ---------------------------------------------------------------------------


def _render_markdown_with_math(text: str) -> None:
    """Render Markdown text, sending display math blocks through ``st.latex``.

    Parameters
    ----------
    text : str
        Markdown text that may contain display math blocks.
    """
    for block_type, content in split_markdown_latex_blocks(text):
        if block_type == "latex":
            st.latex(_prepare_latex_block(content))
        else:
            st.markdown(content)


def _prepare_latex_block(content: str) -> str:
    """Clean display math for Streamlit's KaTeX renderer.

    Parameters
    ----------
    content : str
        Raw LaTeX block content.

    Returns
    -------
    str
        KaTeX-compatible display math.
    """
    lines = [line.strip() for line in content.splitlines() if line.strip()]
    if not lines:
        return ""
    if len(lines) == 1 or r"\begin{" in content:
        return " ".join(lines)
    return "\\begin{aligned}\n" + (r" \\" + "\n").join(lines) + "\n\\end{aligned}"


def _format_calculator_label(calculator_name: str) -> str:
    """Format calculator class names for display.

    Parameters
    ----------
    calculator_name : str
        Calculator class name or label.

    Returns
    -------
    str
        Human-readable calculator label.
    """
    label = calculator_name.removesuffix("Calc")
    if label == "TBLite":
        return "TBLite (xTB, GFN1-xTB, GFN2-xTB)"
    return label


def _render_available_calculators_sidebar() -> None:
    """Render the calculators detected during ChemGraph initialization."""
    available = get_available_calculator_names()
    default = get_default_calculator_name()

    with st.sidebar.expander("\U0001f9ee Available Calculators", expanded=True):
        st.caption("Detected during ChemGraph initialization.")

        for calculator_name in available:
            label = _format_calculator_label(calculator_name)
            if calculator_name == default:
                st.success(f"{label} (default)")
            else:
                st.markdown(f"- {label}")

        st.caption(
            "The agent uses this list when choosing calculators for ASE simulations. "
            "Availability does not imply suitability; verify calculator domain, "
            "charge/spin, and convergence settings for your system."
        )


def _start_new_chat() -> None:
    """Reset conversation state for a fresh chat session."""
    if st.session_state.get("agent_run_lock_acquired"):
        try:
            _AGENT_RUN_LOCK.release()
        except RuntimeError:
            pass
    st.session_state.conversation_history.clear()
    st.session_state.current_session_id = None
    st.session_state.session_created = False
    st.session_state.query_input = ""
    st.session_state.last_run_error = None
    st.session_state.last_run_result = None
    st.session_state.last_run_query = None
    st.session_state.pop("_pending_example_query", None)
    st.session_state.agent = None
    st.session_state.last_config = None
    st.session_state.current_chat_log_dir = None
    st.session_state.active_query_log_dir = None
    st.session_state.active_graph_thread_id = None
    st.session_state.agent_running = False
    st.session_state.agent_run_lock_acquired = False
    st.session_state.pending_agent_submission = None
    _set_session_id_in_url(None)
    os.environ.pop("CHEMGRAPH_LOG_DIR", None)
    _clear_interrupt_state()


def _ensure_graph_thread_id(thread_id: int) -> str:
    """Return a stable LangGraph thread ID for the current chat.

    The thread must persist across follow-up questions inside one chat so the
    agent can resolve references like "its dipole moment". It must also change
    when the user starts a new chat so old molecule/tool context does not leak
    into an unrelated task.
    """
    active_id = st.session_state.get("active_graph_thread_id")
    if not active_id:
        active_id = f"{thread_id}-{uuid.uuid4().hex[:8]}"
        st.session_state.active_graph_thread_id = active_id
    return active_id


def _render_chat_controls() -> None:
    """Render chat-level actions that must be available even without memory."""
    if st.sidebar.button(
        "New Chat",
        key="new_chat_btn",
        use_container_width=True,
        disabled=_agent_run_active(),
    ):
        _start_new_chat()
        st.rerun()


def _render_session_sidebar() -> None:
    """Render the session management panel in the sidebar."""
    store: Optional[SessionStore] = st.session_state.get("session_store")
    if store is None:
        return

    with st.sidebar.expander("\U0001f4c2 Sessions", expanded=False):
        # Show current session info
        current_sid = st.session_state.get("current_session_id")
        if current_sid:
            st.caption(f"Active session: `{current_sid}`")

        # List recent sessions
        try:
            sessions = store.list_sessions(limit=10)
        except Exception:
            sessions = []

        if not sessions:
            st.caption("No saved sessions yet.")
            return

        st.markdown("**Recent sessions:**")
        for s in sessions:
            # Highlight the active session
            is_active = current_sid and s.session_id == current_sid
            prefix = "\u25b6 " if is_active else ""
            label = s.title or "Untitled"
            if len(label) > 35:
                label = label[:32] + "..."

            col_load, col_del = st.columns([4, 1])
            with col_load:
                if st.button(
                    f"{prefix}{label}",
                    key=f"load_session_{s.session_id}",
                    use_container_width=True,
                    help=(
                        f"Model: {s.model_name} | "
                        f"Queries: {s.query_count} | "
                        f"{s.updated_at.strftime('%Y-%m-%d %H:%M')}"
                    ),
                ):
                    _load_session(s.session_id)
                    st.rerun()
            with col_del:
                if st.button(
                    "\U0001f5d1",
                    key=f"del_session_{s.session_id}",
                    help="Delete this session",
                ):
                    try:
                        store.delete_session(s.session_id)
                        # If we just deleted the active session, reset
                        if current_sid == s.session_id:
                            _start_new_chat()
                    except Exception as exc:
                        logger.warning(
                            "Failed to delete session %s: %s", s.session_id, exc
                        )
                    st.rerun()


def _load_session(session_id: str) -> None:
    """Load a stored session into the active conversation.

    Parameters
    ----------
    session_id : str
        Session ID or prefix selected in the sidebar.
    """
    store: Optional[SessionStore] = st.session_state.get("session_store")
    if store is None:
        return

    session = store.get_session(session_id)
    if session is None:
        st.sidebar.error(f"Session '{session_id}' not found.")
        return

    # Rebuild conversation_history from stored messages
    st.session_state.conversation_history = session_to_conversation_history(session)
    st.session_state.current_session_id = session.session_id
    st.session_state.session_created = True
    st.session_state.query_input = ""
    st.session_state.last_run_error = None
    st.session_state.last_run_result = None
    st.session_state.last_run_query = None
    st.session_state.current_chat_log_dir = session.log_dir
    st.session_state.active_graph_thread_id = f"session-{session.session_id}"
    _set_session_id_in_url(session.session_id)
    if session.log_dir:
        os.environ["CHEMGRAPH_LOG_DIR"] = session.log_dir


def _restore_session_from_url() -> None:
    """Restore the visible chat after a browser refresh using URL state."""
    if st.session_state.get("current_session_id"):
        return
    session_id = _session_id_from_url()
    if not session_id:
        return
    try:
        _load_session(session_id)
    except Exception as exc:
        logger.warning("Failed to restore session %s from URL: %s", session_id, exc)


def _active_session_metadata() -> tuple[str, str]:
    """Return model/workflow metadata matching the active UI run."""
    config = st.session_state.config
    model = (
        st.session_state.get("pending_interrupt_model")
        or st.session_state.get("active_model")
        or config["general"]["model"]
    )
    workflow = (
        st.session_state.get("pending_interrupt_workflow")
        or st.session_state.get("active_workflow")
        or config["general"]["workflow"]
    )
    return model, normalize_workflow_name(workflow)


def _ensure_ui_session(query: str) -> Optional[str]:
    """Create the UI session-store row if needed and return its id."""
    store: Optional[SessionStore] = st.session_state.get("session_store")
    if store is None:
        return None

    model, workflow = _active_session_metadata()
    if not st.session_state.session_created:
        sid = st.session_state.get("current_session_id") or generate_session_id()
        st.session_state.current_session_id = sid
        title = SessionStore.generate_title(query)
        store.create_session(
            session_id=sid,
            model_name=model,
            workflow_type=workflow,
            title=title,
            log_dir=st.session_state.get("current_chat_log_dir"),
        )
        st.session_state.session_created = True

    _set_session_id_in_url(st.session_state.current_session_id)
    return st.session_state.current_session_id


def _save_query_to_store(query: str) -> bool:
    """Persist the human query immediately so refresh does not hide it."""
    store: Optional[SessionStore] = st.session_state.get("session_store")
    if store is None:
        return False
    try:
        session_id = _ensure_ui_session(query)
        if session_id:
            store.save_messages(
                session_id,
                [SessionMessage(role="human", content=query)],
            )
            return True
    except Exception as exc:
        logger.warning("Failed to save pending query to session store: %s", exc)
    return False


def _save_exchange_to_store(
    query: str,
    result: Any,
    *,
    include_query: bool = True,
) -> None:
    """Persist a single query/result exchange to the SessionStore.

    Creates the session DB row on the first call, then appends messages.

    Parameters
    ----------
    query : str
        User query text.
    result : Any
        Agent result to persist as session messages.
    """
    store: Optional[SessionStore] = st.session_state.get("session_store")
    if store is None:
        return

    try:
        _ensure_ui_session(query)

        # Build SessionMessage objects for this exchange
        if include_query:
            entry = {"query": query, "result": result}
            messages = conversation_entry_to_messages(entry)
        else:
            messages = [
                msg for msg in messages_from_result(result) if msg.role != "human"
            ]
        if messages:
            store.save_messages(st.session_state.current_session_id, messages)
    except Exception as exc:
        # Best-effort persistence -- don't break the UI.
        logger.warning("Failed to save exchange to session store: %s", exc)


def _render_agent_status(
    selected_model: str,
    selected_workflow: str,
    thread_id: int,
    endpoint_status: dict,
) -> None:
    """Render sidebar status for the active agent.

    Parameters
    ----------
    selected_model : str
        Selected model name.
    selected_workflow : str
        Selected workflow name.
    thread_id : int
        Current LangGraph thread ID.
    endpoint_status : dict
        Local endpoint status dictionary.
    """
    st.sidebar.header("Agent Status")

    if st.session_state.agent:
        st.sidebar.success("\u2705 Agents Ready")
        st.sidebar.info(f"\U0001f9e0 Model: {selected_model}")
        st.sidebar.info(f"\u2699\ufe0f Workflow: {selected_workflow}")
        st.sidebar.info(f"\U0001f517 Thread ID: {thread_id}")
        st.sidebar.info(
            f"\U0001f4ac Messages: {len(st.session_state.conversation_history)}"
        )
        if endpoint_status["ok"]:
            st.sidebar.caption(f"LLM endpoint: {endpoint_status['message']}")
        else:
            st.sidebar.error(f"LLM endpoint issue: {endpoint_status['message']}")
        if st.session_state.pending_human_question is not None:
            st.sidebar.warning("Waiting for your input...")
        if st.session_state.last_run_error:
            st.sidebar.error("Last run error (see verbose info).")

        if st.sidebar.button("\U0001f504 Refresh Agents"):
            st.session_state.agent = None
            # Checkpoint is lost on re-init, so clear interrupt state
            _clear_interrupt_state()
            st.rerun()
    else:
        st.sidebar.error("\u274c Agents Not Ready")
        st.sidebar.info("Agents will initialize automatically...")
        if not endpoint_status["ok"]:
            st.sidebar.error(f"LLM endpoint issue: {endpoint_status['message']}")

    st.sidebar.markdown("---")
    st.sidebar.markdown("**\u2699\ufe0f Configuration**")
    st.sidebar.markdown(
        "Use the Configuration page to modify settings, API endpoints, and chemistry parameters."
    )
    st.sidebar.markdown("Current config loaded from: `config.toml`")


def _auto_initialize_agent(
    config: dict,
    selected_model: str,
    selected_workflow: str,
    structured_output: bool,
    selected_output: str,
    generate_report: bool,
    human_supervised: bool,
    selected_base_url: Optional[str],
) -> None:
    """Initialize or refresh the cached Streamlit agent when config changes.

    Parameters
    ----------
    config : dict
        Nested UI configuration.
    selected_model : str
        Selected model name.
    selected_workflow : str
        Selected workflow name.
    structured_output : bool
        Effective structured-output setting.
    selected_output : str
        Agent return mode.
    generate_report : bool
        Whether report generation is enabled.
    human_supervised : bool
        Whether human-supervision tools are enabled.
    selected_base_url : str, optional
        Model endpoint URL.
    """
    current_config = (
        selected_model,
        selected_workflow,
        structured_output,
        selected_output,
        generate_report,
        human_supervised,
        config["general"]["recursion_limit"],
        selected_base_url,
        get_argo_user_from_nested_config(config),
        st.session_state.get("current_chat_log_dir"),
    )

    if st.session_state.agent is None or st.session_state.last_config != current_config:
        with st.spinner("\U0001f680 Initializing ChemGraph agents..."):
            chat_log_dir = _ensure_chat_log_dir()
            agent = initialize_agent(
                selected_model,
                selected_workflow,
                structured_output,
                selected_output,
                generate_report,
                human_supervised,
                config["general"]["recursion_limit"],
                selected_base_url,
                get_argo_user_from_nested_config(config),
                log_dir=chat_log_dir,
            )
            st.session_state.agent = agent
            if agent is not None:
                st.session_state.last_config = (
                    selected_model,
                    selected_workflow,
                    structured_output,
                    selected_output,
                    generate_report,
                    human_supervised,
                    config["general"]["recursion_limit"],
                    selected_base_url,
                    get_argo_user_from_nested_config(config),
                    chat_log_dir,
                )
            else:
                st.session_state.last_config = None


def _render_conversation_history(thread_id: int) -> None:
    """Render all saved conversation exchanges.

    Parameters
    ----------
    thread_id : int
        Current LangGraph thread ID.
    """
    if not st.session_state.conversation_history:
        return

    pending = st.session_state.get("pending_agent_submission")
    pending_log_dir = pending.get("query_log_dir") if isinstance(pending, dict) else None
    pending_interrupt_log_dir = (
        st.session_state.get("pending_interrupt_log_dir")
        if st.session_state.get("pending_human_question") is not None
        else None
    )
    visible_idx = 1
    for entry in st.session_state.conversation_history:
        if (
            pending_log_dir
            and entry.get("status") == "running"
            and entry.get("log_dir") == pending_log_dir
        ):
            # The live streaming panel below renders this same submitted query.
            # Avoid showing a duplicate history bubble while the run is active.
            continue
        if (
            pending_interrupt_log_dir
            and entry.get("status") == "waiting_for_input"
            and entry.get("log_dir") == pending_interrupt_log_dir
        ):
            # The pending interrupt panel renders the original query, prior
            # clarification exchanges, and the current question.
            continue
        _render_single_exchange(visible_idx, entry, thread_id)
        visible_idx += 1


def _render_single_exchange(idx: int, entry: dict, thread_id: int) -> None:
    """Render one user-query / agent-response exchange.

    Parameters
    ----------
    idx : int
        One-based exchange index.
    entry : dict
        Conversation-history entry.
    thread_id : int
        Current LangGraph thread ID.
    """
    # User message
    with st.chat_message("user"):
        st.markdown(entry["query"])

    # Interrupt exchanges (if any occurred during this query)
    for exch in entry.get("interrupt_exchanges", []):
        with st.chat_message("assistant"):
            _render_markdown_with_math(exch["question"])
        with st.chat_message("user"):
            st.markdown(exch["answer"])

    messages = extract_messages_from_result(entry["result"])

    if not messages:
        with st.chat_message("assistant"):
            status = entry.get("status")
            if status == "error":
                st.error(f"Processing error: {entry.get('error', 'unknown error')}")
            elif status == "waiting_for_input":
                st.info("The agent is waiting for your input.")
            elif status == "running" or _agent_run_active():
                st.info("Agent is running. Please wait for the current task to finish.")
            else:
                st.info(
                    "No assistant response was saved for this run. "
                    "It may have been interrupted; please rerun the query."
                )
            _render_workflow_process_trace(messages, entry)
            _render_verbose_info(idx, messages, entry)
        return

    # Prefer validated structured output over intermediate LLM prose.
    final_answer = _final_answer_for_entry(entry, messages)

    # Display the AI response with visualizations
    with st.chat_message("assistant"):
        if final_answer:
            _render_markdown_with_math(final_answer)

        _render_workflow_process_trace(messages, entry)

        # Structure visualisation
        html_filename = find_html_filename(messages)
        _render_structure_section(idx, messages, final_answer, entry, html_filename)

        # HTML report
        if html_filename:
            _render_html_report(idx, html_filename, messages, entry)

        # IR spectrum
        if is_infrared_requested(messages):
            _render_ir_spectrum(idx, messages, entry)

        # Debug expander
        _render_verbose_info(idx, messages, entry)


def _extract_final_answer(messages: list) -> str:
    """Walk messages in reverse to find the last non-JSON AI message.

    Parameters
    ----------
    messages : list
        Message-like objects or dictionaries.

    Returns
    -------
    str
        Final displayable answer text, or an empty string.
    """
    final_answer = ""
    for message in reversed(messages):
        if hasattr(message, "content") and hasattr(message, "type"):
            content = normalize_message_content(message.content).strip()
            if message.type == "ai" and content:
                if not (
                    content.startswith("{")
                    and content.endswith("}")
                    and "numbers" in content
                ):
                    final_answer = content
                    break
        elif isinstance(message, dict):
            content = normalize_message_content(message.get("content", "")).strip()
            if message.get("type") == "ai" and content:
                if not (
                    content.startswith("{")
                    and content.endswith("}")
                    and "numbers" in content
                ):
                    final_answer = content
                    break
        elif hasattr(message, "content"):
            content = normalize_message_content(getattr(message, "content", "")).strip()
            if content and not (
                content.startswith("{")
                and content.endswith("}")
                and "numbers" in content
            ):
                final_answer = content
                break
    return final_answer


def _is_response_formatter_payload(payload: Any) -> bool:
    """Return whether *payload* looks like ChemGraph structured final output."""
    if not isinstance(payload, dict):
        return False
    formatter_fields = {
        "smiles",
        "scalar_answer",
        "scalar_answers",
        "dipole",
        "vibrational_answer",
        "ir_spectrum",
        "atoms_data",
        "_failure",
    }
    return any(field in payload for field in formatter_fields)


def _parse_structured_output_text(text: str) -> Optional[dict]:
    """Parse a structured-output JSON string if it matches ResponseFormatter."""
    stripped = normalize_message_content(text).strip()
    if not (stripped.startswith("{") and stripped.endswith("}")):
        return None
    try:
        payload = json.loads(stripped)
    except json.JSONDecodeError:
        return None
    if not _is_response_formatter_payload(payload):
        return None
    return payload


def _structured_output_from_entry(entry: dict, messages: list) -> Optional[dict]:
    """Return the validated final output saved by the graph, if available."""
    result = entry.get("result") if isinstance(entry, dict) else None
    if isinstance(result, dict):
        final_output = result.get("final_output")
        if _is_response_formatter_payload(final_output):
            return _augment_structured_output_from_messages(entry, final_output, messages)

    for message in reversed(messages):
        if hasattr(message, "content"):
            parsed = _parse_structured_output_text(getattr(message, "content", ""))
        elif isinstance(message, dict):
            parsed = _parse_structured_output_text(message.get("content", ""))
        else:
            parsed = None
        if parsed:
            return _augment_structured_output_from_messages(entry, parsed, messages)
    return None


def _augment_structured_output_from_messages(
    entry: dict,
    structured: dict,
    messages: list,
) -> dict:
    """Fill derived fields from saved tool messages when older output lacks them."""
    if structured.get("scalar_answers"):
        return structured
    query = str(entry.get("query") or "").strip()
    if not query or not messages:
        return structured
    try:
        validation = validate_completion(query, messages)
    except Exception:
        return structured
    candidate = (
        validation.get("structured_output") if isinstance(validation, dict) else None
    )
    if not isinstance(candidate, dict) or not candidate.get("scalar_answers"):
        return structured
    augmented = dict(structured)
    augmented["scalar_answers"] = candidate["scalar_answers"]
    return augmented


def _workflow_status_from_result(result: dict) -> str:
    """Return scientific workflow status, not just graph stream status."""
    if not isinstance(result, dict):
        return "complete"

    final_output = result.get("final_output")
    if isinstance(final_output, dict):
        failure = final_output.get("_failure")
        if isinstance(failure, dict):
            status = str(failure.get("status") or "").lower()
            if status in {"blocked", "failed", "error"}:
                return "failed" if status == "error" else status
            return "failed"

    ledger = result.get("scientific_ledger")
    if isinstance(ledger, dict):
        status = str(ledger.get("status") or "").lower()
        if status in {"complete", "blocked", "failed", "partial"}:
            return status

    validation = (
        result.get("completion_validation")
        if isinstance(result.get("completion_validation"), dict)
        else result.get("validation")
    )
    if isinstance(validation, dict) and "complete" in validation:
        return "complete" if validation.get("complete") else "blocked"

    return "complete"


def _workflow_status_label(workflow_status: str) -> str:
    return {
        "complete": "Complete",
        "blocked": "Blocked",
        "failed": "Failed",
        "partial": "Partial",
    }.get(str(workflow_status or "").lower(), "Complete")


def _streamlit_state_for_workflow_status(workflow_status: str) -> str:
    return "complete" if str(workflow_status or "").lower() == "complete" else "error"


def _format_number_for_display(value: Any) -> str:
    """Format numeric values without excessive trailing precision."""
    if isinstance(value, int):
        return str(value)
    if isinstance(value, float):
        return f"{value:.8g}"
    return str(value)


def _plain_scientific_unit(unit: Any) -> str:
    """Return a readable unit string for Markdown text."""
    text = normalize_latex_delimiters(str(unit or "")).strip()
    text = text.strip("$")
    text = re.sub(r"\\text\{([^{}]+)\}", r"\1", text)
    text = re.sub(r"\\mathrm\{([^{}]+)\}", r"\1", text)
    text = text.replace(r"\cdot", "·")
    text = text.replace(r"\AA", "Å")
    text = text.replace("Angstrom", "Å").replace("angstrom", "Å")
    text = re.sub(r"\s*\*\s*", "·", text)
    text = re.sub(r"\s*·\s*", "·", text)
    return text


def _format_structured_output_answer(structured: Optional[dict]) -> str:
    """Create user-facing Markdown from validated structured output.

    The UI should prefer this over intermediate LLM prose because these fields
    come from the graph's completion validator or ResponseFormatter output.
    """
    if not structured:
        return ""

    if isinstance(structured.get("_failure"), dict):
        return format_validation_failure_message(structured)

    lines: list[str] = []

    smiles = structured.get("smiles")
    if smiles:
        values = smiles if isinstance(smiles, list) else [smiles]
        lines.append("**SMILES:** " + ", ".join(f"`{value}`" for value in values))

    scalar_values = structured.get("scalar_answers")
    if isinstance(scalar_values, list):
        scalars = [
            item
            for item in scalar_values
            if isinstance(item, dict) and item.get("value") is not None
        ]
    else:
        scalar = structured.get("scalar_answer")
        scalars = (
            [scalar]
            if isinstance(scalar, dict) and scalar.get("value") is not None
            else []
        )

    for scalar in scalars:
        property_name = str(scalar.get("property") or "Scalar result").strip()
        value = _format_number_for_display(scalar.get("value"))
        unit = _plain_scientific_unit(scalar.get("unit"))
        suffix = f" {unit}" if unit else ""
        lines.append(f"**{property_name}:** {value}{suffix}")

    dipole = structured.get("dipole")
    if isinstance(dipole, dict) and dipole.get("value") is not None:
        values = dipole.get("value")
        if isinstance(values, list):
            vector = ", ".join(_format_number_for_display(value) for value in values)
            vector = f"[{vector}]"
        else:
            vector = _format_number_for_display(values)
        unit = _plain_scientific_unit(dipole.get("unit"))
        suffix = f" {unit}" if unit else ""
        lines.append(f"**Dipole moment:** {vector}{suffix}")

    vibrational = structured.get("vibrational_answer")
    if isinstance(vibrational, dict):
        frequencies = vibrational.get("frequency_cm1") or []
        if frequencies:
            sample = ", ".join(str(freq) for freq in frequencies[:10])
            extra = "" if len(frequencies) <= 10 else f", ... ({len(frequencies)} total)"
            lines.append(f"**Vibrational frequencies:** {sample}{extra} cm^-1")

    ir_spectrum = structured.get("ir_spectrum")
    if isinstance(ir_spectrum, dict):
        frequencies = ir_spectrum.get("frequency_cm1") or []
        intensities = ir_spectrum.get("intensity") or []
        if frequencies or intensities:
            lines.append(
                "**IR spectrum:** "
                f"{len(frequencies)} frequencies and {len(intensities)} intensities."
            )
        plot = ir_spectrum.get("plot")
        if plot:
            lines.append(f"**IR plot:** `{plot}`")

    if structured.get("atoms_data"):
        lines.append("**Structure:** atomic geometry data is available below.")

    return "\n\n".join(lines)


def _final_answer_for_entry(entry: dict, messages: list) -> str:
    """Return the final answer text for a UI exchange."""
    structured = _structured_output_from_entry(entry, messages)
    structured_answer = _format_structured_output_answer(structured)
    if structured_answer:
        return structured_answer
    return _extract_final_answer(messages)


def _render_structure_section(
    idx: int,
    messages: list,
    final_answer: str,
    entry: dict,
    html_filename: Optional[str],
) -> None:
    """Render molecular structure artifacts for an exchange.

    Parameters
    ----------
    idx : int
        One-based exchange index.
    messages : list
        Message-like objects from the exchange.
    final_answer : str
        Final assistant answer text.
    entry : dict
        Conversation-history entry.
    html_filename : str, optional
        HTML report path/filename, if detected.
    """
    structure = find_structure_in_messages(messages)
    if structure:
        display_molecular_structure(
            structure["atomic_numbers"],
            structure["positions"],
            title=f"Molecular Structure (Query {idx})",
        )
    else:
        structure_from_text = extract_molecular_structure(final_answer)
        if structure_from_text:
            display_molecular_structure(
                structure_from_text["atomic_numbers"],
                structure_from_text["positions"],
                title=f"Structure from Response {idx}",
            )
        elif not html_filename:
            if has_structure_signal(messages, entry.get("query", ""), final_answer):
                log_dir = _artifact_log_dir(messages, entry)
                if log_dir and os.path.isdir(log_dir):
                    latest_xyz = find_latest_xyz_file_in_dir(log_dir)
                    if latest_xyz:
                        try:
                            atoms = ase_read(latest_xyz)
                            display_molecular_structure(
                                atoms.get_atomic_numbers().tolist(),
                                atoms.get_positions().tolist(),
                                title=(
                                    f"Structure from {Path(latest_xyz).name} "
                                    f"(Query {idx})"
                                ),
                            )
                        except Exception as exc:
                            st.warning(f"Failed to load XYZ structure: {exc}")


def _render_html_report(
    idx: int, html_filename: str, messages: list, entry: dict
) -> None:
    """Render an HTML report expander and download button.

    Parameters
    ----------
    idx : int
        One-based exchange index.
    html_filename : str
        HTML report path or filename.
    messages : list
        Message-like objects from the exchange.
    entry : dict
        Conversation-history entry.
    """
    with st.expander("\U0001f4ca Report", expanded=False):
        try:
            resolved_html = _resolve_artifact_path(
                html_filename,
                _artifact_log_dir(messages, entry),
            )
            with open(resolved_html, "r", encoding="utf-8") as f:
                html_content = f.read()

            report_structure = extract_xyz_from_report_html(html_content)
            if report_structure:
                display_molecular_structure(
                    report_structure["atomic_numbers"],
                    report_structure["positions"],
                    title=f"Molecular Structure (Report {idx})",
                )

            cleaned_html = strip_viewer_from_report_html(html_content)
            st.download_button(
                "Download HTML Report",
                data=html_content,
                file_name=Path(resolved_html).name,
                mime="text/html",
                key=f"download_report_{idx}",
            )
            st.components.v1.html(cleaned_html, height=600, scrolling=True)
        except FileNotFoundError:
            st.warning(f"HTML file '{html_filename}' not found")
        except Exception as e:
            st.error(f"Error displaying HTML: {e}")


def _artifact_log_dir(messages: list, entry: dict) -> Optional[str]:
    """Return the log directory tied to a specific conversation entry.

    Parameters
    ----------
    messages : list
        Message-like objects from the exchange.
    entry : dict
        Conversation-history entry.

    Returns
    -------
    str or None
        Artifact/log directory, if found.
    """
    entry_log_dir = entry.get("log_dir")
    if entry_log_dir:
        return entry_log_dir
    return extract_log_dir_from_messages(messages)


def _latest_artifact_path(directory: Optional[str], pattern: str) -> Optional[str]:
    """Return the newest shallow match for an output artifact pattern.

    Parameters
    ----------
    directory : str, optional
        Directory to search.
    pattern : str
        Glob pattern to match.

    Returns
    -------
    str or None
        Newest matching file path, or ``None``.
    """
    if not directory or not os.path.isdir(directory):
        return None

    candidates: list[Path] = []
    try:
        candidates.extend(path for path in Path(directory).glob(pattern) if path.is_file())
    except OSError:
        return None

    if not candidates:
        return None
    return str(max(candidates, key=lambda path: path.stat().st_mtime))


def _resolve_artifact_path(filename: str, directory: Optional[str]) -> str:
    """Resolve an artifact path relative to its run directory when known.

    Parameters
    ----------
    filename : str
        Absolute or relative artifact path.
    directory : str, optional
        Run artifact directory.

    Returns
    -------
    str
        Resolved artifact path.
    """
    if os.path.isabs(filename):
        return filename
    if directory:
        return str(Path(directory) / filename)
    return filename


def _render_ir_spectrum(idx: int, messages: list, entry: dict) -> None:
    """Render IR spectrum plot, frequency table, and trajectory viewer.

    Parameters
    ----------
    idx : int
        One-based exchange index.
    messages : list
        Message-like objects from the exchange.
    entry : dict
        Conversation-history entry.
    """
    log_dir = _artifact_log_dir(messages, entry)
    ir_path = _latest_artifact_path(log_dir, "ir_spectrum*.png")
    freq_path = _latest_artifact_path(log_dir, "frequencies*.csv")

    if not ir_path and not freq_path:
        st.warning("IR spectrum not found.")
        return

    with st.expander("\U0001f50d IR Spectrum", expanded=True):
        col1, col2 = st.columns(2, border=True)

        with col1:
            if ir_path and os.path.exists(ir_path):
                st.image(ir_path)
            else:
                st.warning("IR spectrum plot not found.")

        with col2:
            if not freq_path or not os.path.exists(freq_path):
                st.warning("Frequencies file not found.")
                return

            df = pd.read_csv(
                freq_path,
                index_col=False,
                names=["filename", "frequency"],
            )
            modes = df.iloc[6:] if len(df) > 6 else df

            if modes.empty:
                st.warning("No vibrational frequencies found.")
                return

            st.write("**Select a frequency to visualize:**")
            freq_options = {}
            for mode_idx, row in modes.iterrows():
                freq_text = str(row["frequency"]).strip()
                suffix = "i" if freq_text.endswith("i") else ""
                try:
                    freq_value = float(freq_text.rstrip("i"))
                    label = f"Mode {mode_idx}: {freq_value:.2f}{suffix} cm\u207b\u00b9"
                except ValueError:
                    label = f"Mode {mode_idx}: {freq_text} cm\u207b\u00b9"
                freq_options[label] = mode_idx

            selected_freq = st.selectbox(
                "Frequency",
                list(freq_options.keys()),
                index=0,
                key=f"ir_frequency_select_{idx}",
            )
            traj_file = str(modes.loc[freq_options[selected_freq]]["filename"])
            traj_path = _resolve_artifact_path(traj_file, log_dir)
            if not os.path.exists(traj_path):
                st.warning(f"Trajectory file '{traj_file}' not found.")
            elif not STMOL_AVAILABLE:
                st.info("3D viewer not available; install stmol to animate trajectories.")
            else:
                import stmol
                from ase.io.trajectory import Trajectory

                traj = Trajectory(traj_path)
                view = visualize_trajectory(traj)
                view.zoomTo()
                stmol.showmol(view, height=400, width=500)


def _render_workflow_process_trace(messages: list, entry: dict) -> None:
    """Render the validated tool-path visualization for a saved exchange."""
    ledger = _scientific_ledger_from_entry(entry)
    if isinstance(ledger, dict):
        trace = _workflow_trace_from_ledger(ledger)
    else:
        trace = entry.get("workflow_trace")
        rebuilt_trace = _workflow_trace_from_messages(messages) if messages else None
        if not isinstance(trace, dict):
            trace = rebuilt_trace
        elif _trace_has_incomplete_run_ase(trace) and isinstance(rebuilt_trace, dict):
            trace = rebuilt_trace

    completed = trace.get("completed", []) if isinstance(trace, dict) else []
    active = trace.get("active", []) if isinstance(trace, dict) else []
    if not completed and not active:
        return

    st.html(_workflow_trace_diagram_html(completed, active))
    st.html(_workflow_tool_status_html(completed, active))
    identity_html = _identity_provenance_html(_identity_records_from_messages(messages))
    if identity_html:
        st.html(identity_html)


def _scientific_ledger_from_entry(entry: dict) -> Optional[dict]:
    """Return saved scientific ledger for a UI exchange, when available."""
    result = entry.get("result") if isinstance(entry, dict) else None
    if not isinstance(result, dict):
        return None
    ledger = result.get("scientific_ledger")
    return ledger if isinstance(ledger, dict) else None


def _workflow_trace_from_ledger(ledger: dict) -> dict[str, list]:
    """Build UI trace labels from canonical requirement ledger rows."""
    completed: list[str] = []
    active: list[dict] = []
    can_answer = bool(ledger.get("can_answer"))
    artifacts_by_id = {
        artifact.get("id"): artifact
        for artifact in ledger.get("artifacts", []) or []
        if isinstance(artifact, dict) and artifact.get("id")
    }
    for req in ledger.get("requirements", []) or []:
        if not isinstance(req, dict):
            continue
        label = _ledger_requirement_label(req, artifacts_by_id)
        status = str(req.get("status") or "").lower()
        if status == "done":
            completed.append(label)
        elif status in {"failed", "blocked"}:
            completed.append(f"{label} ({status})")
        elif can_answer:
            continue
        else:
            completed.append(f"{label} (incomplete)")
    return {"completed": completed, "active": active}


def _ledger_requirement_label(req: dict, artifacts_by_id: Optional[dict] = None) -> str:
    """Return a trace-compatible label for one ledger requirement."""
    kind = str(req.get("kind") or "requirement")
    species = req.get("species")
    target = f" on {species}" if species else ""
    if kind == "identity":
        return f"molecule_name_to_smiles:{target}".replace(": on", ":")
    if kind == "xyz":
        return f"smiles_to_coordinate_file:{target}".replace(": on", ":")
    if kind == "aggregation":
        return "calculator: aggregate reaction property"
    if kind in {"energy", "dipole", "thermo", "vib", "ir"}:
        calculator = _ledger_requirement_calculator(req, artifacts_by_id or {})
        calc_suffix = f"/{calculator}" if calculator else ""
        return f"run_ase: {kind}{calc_suffix}{target}"
    return f"{kind}:{target}".replace(": on", ":")


def _ledger_requirement_calculator(req: dict, artifacts_by_id: dict) -> str:
    """Return calculator family from a requirement's satisfied artifacts."""
    for artifact_id in req.get("satisfied_by", []) or []:
        artifact = artifacts_by_id.get(artifact_id)
        if not isinstance(artifact, dict):
            continue
        calculator = artifact.get("calculator")
        if calculator:
            return str(calculator)
    return ""


def _trace_has_incomplete_run_ase(trace: dict) -> bool:
    """Return whether a saved workflow trace contains stale incomplete run_ase rows."""
    labels = []
    if isinstance(trace, dict):
        labels.extend(str(label) for label in trace.get("completed", []) or [])
        labels.extend(_trace_entry_label(item) for item in trace.get("active", []) or [])
    return any("run_ase" in label and "(incomplete)" in label for label in labels)


def _render_verbose_info(idx: int, messages: list, entry: dict) -> None:
    """Render raw result/debug information for an exchange.

    Parameters
    ----------
    idx : int
        One-based exchange index.
    messages : list
        Message-like objects from the exchange.
    entry : dict
        Conversation-history entry.
    """
    structure = find_structure_in_messages(messages)
    with st.expander(f"\U0001f50d Verbose Info (Query {idx})", expanded=False):
        st.write(f"**Number of messages:** {len(messages)}")
        st.write(f"**Structure found:** {'Yes' if structure else 'No'}")
        raw_result = entry.get("result")
        if st.session_state.last_run_query == entry.get("query"):
            if st.session_state.last_run_error:
                st.write("**Last run error:**")
                st.code(str(st.session_state.last_run_error))
            if st.session_state.last_run_result is not None:
                raw_result = st.session_state.last_run_result

        st.write("**Raw result:**")
        st.code(pprint.pformat(raw_result, width=1, compact=False), language="text")


def _render_example_queries(config: dict, selected_model: str) -> None:
    """Show example queries that the user can click to submit directly.

    Parameters
    ----------
    config : dict
        Nested UI configuration.
    selected_model : str
        Selected model name.
    """
    # Hide after the first message or during an interrupt
    if (
        st.session_state.conversation_history
        or st.session_state.pending_human_question is not None
        or _agent_run_active()
    ):
        return

    with st.expander("Example Queries", expanded=False):
        st.markdown("**Based on your current configuration:**")
        st.markdown(f"- Model: {selected_model}")
        st.markdown(
            f"- Default Calculator: {config['chemistry']['calculators']['default']}"
        )
        st.caption(QUERY_SPECIFICATION_HINT)

        examples = [
            "What is the SMILES string for caffeine?",
            f"Optimize the geometry of water molecule using {config['chemistry']['calculators']['default']}",
            "Calculate the infrared spectrum of methanol with xtb calculator",
            "What is the reaction enthalpy of methane combustion using mace_mp",
        ]
        for ex in examples:
            if st.button(ex, key=f"ex_{ex}"):
                st.session_state._pending_example_query = ex
                st.rerun()


def _render_pending_interrupt() -> None:
    """Show the agent's pending question and any prior interrupt exchanges."""
    question = st.session_state.pending_human_question
    if question is None:
        return

    # Show the original user query that triggered the interrupt
    original_query = st.session_state.pending_interrupt_query
    if original_query:
        with st.chat_message("user"):
            _render_markdown_with_math(original_query)

    # Show any prior interrupt exchanges in this chain
    for exch in st.session_state.interrupt_exchanges:
        with st.chat_message("assistant"):
            _render_markdown_with_math(exch["question"])
        with st.chat_message("user"):
            _render_markdown_with_math(exch["answer"])

    # Show the current pending question
    with st.chat_message("assistant"):
        st.info("The agent needs your input to continue.", icon="\u2753")
        _render_markdown_with_math(question)

    # Cancel button
    if st.button("Cancel", key="cancel_interrupt"):
        _clear_interrupt_state()
        st.rerun()


def _restore_pending_interrupt_from_history() -> None:
    """Restore in-memory clarification state from a waiting history entry."""
    if st.session_state.pending_human_question is not None:
        return

    for entry in reversed(st.session_state.get("conversation_history", [])):
        if entry.get("status") != "waiting_for_input":
            continue

        question = entry.get("pending_human_question")
        resume_config = entry.get("pending_interrupt_config")
        if not question or not isinstance(resume_config, dict):
            return

        st.session_state.pending_human_question = question
        st.session_state.pending_interrupt_config = dict(resume_config)
        st.session_state.pending_interrupt_query = entry.get(
            "pending_interrupt_query", entry.get("query")
        )
        st.session_state.pending_interrupt_thread_id = entry.get(
            "pending_interrupt_thread_id", entry.get("thread_id")
        )
        st.session_state.pending_interrupt_prev_msg_count = int(
            entry.get("pending_interrupt_prev_msg_count", 0) or 0
        )
        st.session_state.pending_interrupt_model = entry.get("pending_interrupt_model")
        st.session_state.pending_interrupt_workflow = entry.get(
            "pending_interrupt_workflow"
        )
        st.session_state.pending_interrupt_log_dir = entry.get(
            "pending_interrupt_log_dir", entry.get("log_dir")
        )
        st.session_state.interrupt_count = int(entry.get("interrupt_count", 1) or 1)
        st.session_state.interrupt_exchanges = list(
            entry.get("interrupt_exchanges") or []
        )
        return


def _clear_interrupt_state() -> None:
    """Clear all interrupt-related session state."""
    st.session_state.pending_human_question = None
    st.session_state.pending_interrupt_config = None
    st.session_state.pending_interrupt_query = None
    st.session_state.pending_interrupt_thread_id = None
    st.session_state.pending_interrupt_prev_msg_count = 0
    st.session_state.pending_interrupt_model = None
    st.session_state.pending_interrupt_workflow = None
    st.session_state.pending_interrupt_log_dir = None
    st.session_state.interrupt_count = 0
    st.session_state.interrupt_exchanges = []


def _classify_message(msg):
    """Classify a LangGraph message for UI display.

    Parameters
    ----------
    msg : Any
        LangGraph/LangChain message to classify.

    Returns
    -------
    tuple or None
        ``("tool_call", [tool_items])``, ``("tool_result", tool_item)``, or
        ``None`` when not relevant for display.
    """
    if isinstance(msg, dict):
        tool_calls = msg.get("tool_calls")
        msg_type = msg.get("type") or msg.get("role")
        tool_name = msg.get("name")
        tool_call_id = msg.get("tool_call_id")
        content = msg.get("content")
    else:
        tool_calls = getattr(msg, "tool_calls", None)
        msg_type = getattr(msg, "type", None)
        tool_name = getattr(msg, "name", None)
        tool_call_id = getattr(msg, "tool_call_id", None)
        content = getattr(msg, "content", None)

    if tool_calls:
        items = [
            _tool_call_display_item(tc)
            for tc in tool_calls
            if isinstance(tc, dict)
        ]
        if items:
            return ("tool_call", items)
    if msg_type == "tool":
        if tool_name:
            return (
                "tool_result",
                {
                    "id": tool_call_id,
                    "name": tool_name,
                    "label": _tool_result_label(
                        tool_name,
                        content,
                    ),
                },
            )
    return None


def _tool_call_display_item(tool_call: dict) -> dict:
    """Return a compact UI label for an LLM tool call."""
    name = tool_call.get("name", "unknown")
    args = tool_call.get("args") or {}
    return {
        "id": tool_call.get("id"),
        "name": name,
        "label": _tool_call_label(name, args if isinstance(args, dict) else {}),
    }


def _tool_call_label(name: str, args: dict) -> str:
    """Format a short in-progress label for a tool call."""
    if name in {"molecule_name_to_smiles", "resolve_molecule_identity"}:
        molecule = args.get("name")
        return f"{name}: {molecule}" if molecule else name

    if name == "smiles_to_coordinate_file":
        smiles = args.get("smiles")
        output_file = args.get("output_file")
        if smiles and output_file:
            return f"{name}: {smiles} -> {Path(str(output_file)).name}"
        if smiles:
            return f"{name}: {smiles}"
        return name

    if name == "run_ase":
        params = args.get("params") if isinstance(args.get("params"), dict) else args
        driver = params.get("driver")
        structure = params.get("input_structure_file")
        calculator = _calculator_label(params.get("calculator"))
        parts = [part for part in [driver, calculator] if part]
        prefix = "/".join(parts) if parts else "ASE"
        if structure:
            return f"{name}: {prefix} on {Path(str(structure)).name}"
        return f"{name}: {prefix}"

    return name


def _tool_result_label(name: str, content: Any) -> str:
    """Format a completed-tool label using the actual tool output when possible."""
    data = _tool_result_content_to_data(content)
    if name in {"molecule_name_to_smiles", "resolve_molecule_identity"} and isinstance(data, dict):
        molecule = data.get("name") or data.get("input_name")
        smiles = data.get("smiles")
        details = []
        if data.get("source"):
            details.append(str(data["source"]))
        if data.get("cid"):
            details.append(f"CID {data['cid']}")
        if data.get("credibility_score") is not None:
            details.append(f"score {data['credibility_score']}")
        if data.get("requires_clarification") or data.get("needs_clarification"):
            details.append("needs clarification")
        suffix = f" ({', '.join(details)})" if details else ""
        if molecule and smiles:
            return f"{name}: {molecule} -> {smiles}{suffix}"
        if molecule:
            return f"{name}: {molecule}{suffix}"

    if name == "smiles_to_coordinate_file" and isinstance(data, dict):
        smiles = data.get("smiles")
        path = data.get("path") or data.get("output_file")
        if smiles and path:
            return f"{name}: {smiles} -> {Path(str(path)).name}"
        if path:
            return f"{name}: {Path(str(path)).name}"

    if name == "run_ase" and isinstance(data, dict):
        status = data.get("status")
        result = data.get("result") if isinstance(data.get("result"), dict) else {}
        details = []
        if "single_point_energy" in data:
            details.append("energy")
        if isinstance(result, dict) and result.get("thermochemistry"):
            details.append("thermo")
        if data.get("dipole_moment"):
            details.append("dipole")
        if isinstance(result, dict) and result.get("ir_spectrum"):
            details.append("ir")
        elif isinstance(result, dict) and result.get("vibrational_frequencies"):
            details.append("vib")
        detail = "/".join(details) if details else "ASE"
        structure_hint = _run_ase_result_structure_hint(data)
        if structure_hint:
            detail = f"{detail} on {structure_hint}"
        return f"{name}: {detail}{f' ({status})' if status else ''}"

    return name


def _run_ase_result_structure_hint(data: dict) -> str:
    """Return a compact structure filename from a run_ase result payload."""
    candidates: list[Any] = []
    for key in ("input_structure_file", "structure_file"):
        if data.get(key):
            candidates.append(data.get(key))
    simulation_input = data.get("simulation_input")
    if isinstance(simulation_input, dict) and simulation_input.get("input_structure_file"):
        candidates.append(simulation_input.get("input_structure_file"))

    message = data.get("message")
    if isinstance(message, str):
        match = re.search(r"saved to\s+(.+?\.json)\b", message)
        if match:
            candidates.append(match.group(1))

    for candidate in candidates:
        hint = _structure_hint_from_path(candidate)
        if hint:
            return hint
    return ""


def _structure_hint_from_path(path_value: Any) -> str:
    """Convert a result or structure path into a displayable XYZ-like basename."""
    if not path_value:
        return ""
    name = Path(str(path_value)).name
    if name.endswith(".xyz"):
        return name
    if name.startswith("output_") and name.endswith(".json"):
        species = name.removeprefix("output_").removesuffix(".json")
        if species:
            return f"{species}.xyz"
    return ""


def _tool_result_content_to_data(content: Any) -> Any:
    """Best-effort parse of a tool message payload for UI labels."""
    if isinstance(content, (dict, list)):
        return content
    text = str(content or "").strip()
    if not text:
        return text
    try:
        return json.loads(text)
    except json.JSONDecodeError:
        return text


def _workflow_trace_from_messages(messages: list) -> dict[str, list]:
    """Rebuild completed tool labels from saved LangGraph messages."""
    completed: list[str] = []
    active: list[dict] = []

    for message in messages or []:
        classified = _classify_message(message)
        if not classified:
            continue
        event_type, event_data = classified
        if event_type == "tool_call":
            for item in event_data:
                active.append(item)
        elif event_type == "tool_result":
            item = event_data
            _pop_active_trace_match(active, item)
            label = str(item.get("label") or item.get("name") or "")
            if label and label not in completed:
                completed.append(label)

    for item in active:
        label = _trace_entry_label(item)
        incomplete = f"{label} (incomplete)"
        if label and incomplete not in completed:
            completed.append(incomplete)

    return {"completed": completed, "active": []}


def _pop_active_trace_match(active: list[dict], result: Any) -> Optional[dict]:
    """Pop the active tool call matching a completed tool result."""
    result_id = result.get("id") if isinstance(result, dict) else None
    result_name = result.get("name") if isinstance(result, dict) else str(result)
    result_label = _trace_entry_label(result)
    result_hint = _trace_structure_hint(result_label)

    if result_id:
        for index, item in enumerate(active):
            if isinstance(item, dict) and item.get("id") == result_id:
                return active.pop(index)

    specific_matches: list[int] = []
    generic_matches: list[int] = []
    for index, item in enumerate(active):
        if not isinstance(item, dict):
            continue
        item_name = item.get("name")
        item_label = _trace_entry_label(item)
        if item_name != result_name and item_label != result_name:
            continue
        generic_matches.append(index)
        item_hint = _trace_structure_hint(item_label)
        if result_hint and item_hint and result_hint == item_hint:
            specific_matches.append(index)

    if specific_matches:
        return active.pop(specific_matches[0])
    if len(generic_matches) == 1:
        return active.pop(generic_matches[0])
    return None


def _identity_records_from_messages(messages: list) -> list[dict]:
    """Extract resolver provenance records from saved tool messages."""
    records: list[dict] = []
    for message in messages or []:
        classified = _classify_message(message)
        if not classified:
            continue
        event_type, event_data = classified
        if event_type != "tool_result":
            continue
        if event_data.get("name") not in {
            "molecule_name_to_smiles",
            "resolve_molecule_identity",
        }:
            continue
        content = message.get("content") if isinstance(message, dict) else getattr(message, "content", None)
        data = _tool_result_content_to_data(content)
        if isinstance(data, dict) and (data.get("smiles") or data.get("candidates")):
            records.append(data)
    return records


def _identity_provenance_html(records: list[dict]) -> str:
    """Render resolver provenance as auditable UI, not model reasoning."""
    if not records:
        return ""

    cards = []
    for record in records[-6:]:
        name = record.get("input_name") or record.get("name") or "chemical"
        smiles = record.get("smiles") or "n/a"
        source = record.get("source") or "unknown"
        cid = record.get("cid") or "n/a"
        score = record.get("credibility_score")
        score_text = "n/a" if score is None else str(score)
        flags = (
            record.get("identity_flags")
            if isinstance(record.get("identity_flags"), dict)
            else {}
        )
        active_flags = [
            key.replace("_", " ")
            for key, value in flags.items()
            if value is True
        ]
        if record.get("requires_clarification") or record.get("needs_clarification"):
            active_flags.append("needs clarification")
        warnings = record.get("warnings") or []
        if isinstance(warnings, str):
            warnings = [warnings]
        warning = record.get("warning")
        if warning and warning not in warnings:
            warnings.append(warning)
        flags_text = ", ".join(active_flags) if active_flags else "none"
        warning_text = " ".join(str(item) for item in warnings if item) or "none"
        candidate_count = len(record.get("candidates") or [])
        cards.append(
            f"""
            <div class="cg-identity-card">
              <div class="cg-identity-title">{html.escape(str(name))}</div>
              <div class="cg-identity-grid">
                <span>SMILES</span><code>{html.escape(str(smiles))}</code>
                <span>Source</span><strong>{html.escape(str(source))}</strong>
                <span>CID</span><strong>{html.escape(str(cid))}</strong>
                <span>Score</span><strong>{html.escape(score_text)}</strong>
                <span>Flags</span><em>{html.escape(flags_text)}</em>
                <span>Candidates</span><strong>{candidate_count}</strong>
              </div>
              <div class="cg-identity-warning">{html.escape(warning_text)}</div>
            </div>
            """
        )

    return f"""
    <style>
    .cg-identity-panel {{
      margin-top: 0.85rem;
      border-left: 3px solid #3b82f6;
      padding-left: 0.75rem;
    }}
    .cg-identity-heading {{
      font-weight: 700;
      margin-bottom: 0.55rem;
    }}
    .cg-identity-cards {{
      display: grid;
      grid-template-columns: repeat(auto-fit, minmax(260px, 1fr));
      gap: 0.65rem;
    }}
    .cg-identity-card {{
      border: 1px solid #334155;
      border-radius: 8px;
      padding: 0.75rem;
      background: rgba(15, 23, 42, 0.55);
    }}
    .cg-identity-title {{
      font-weight: 700;
      margin-bottom: 0.45rem;
    }}
    .cg-identity-grid {{
      display: grid;
      grid-template-columns: 90px minmax(0, 1fr);
      gap: 0.25rem 0.5rem;
      font-size: 0.78rem;
    }}
    .cg-identity-grid span {{
      color: #94a3b8;
    }}
    .cg-identity-grid code {{
      white-space: normal;
      overflow-wrap: anywhere;
    }}
    .cg-identity-warning {{
      margin-top: 0.55rem;
      color: #fbbf24;
      font-size: 0.78rem;
      line-height: 1.35;
    }}
    </style>
    <div class="cg-identity-panel">
      <div class="cg-identity-heading">Identity resolution</div>
      <div class="cg-identity-cards">{''.join(cards)}</div>
    </div>
    """


def _calculator_label(calculator: Any) -> Optional[str]:
    """Return a compact calculator name from string/dict calculator specs."""
    if isinstance(calculator, dict):
        return str(
            calculator.get("name")
            or calculator.get("type")
            or calculator.get("calculator")
            or ""
        ) or None
    if calculator:
        return str(calculator)
    return None


def _trace_entry_label(entry: Any) -> str:
    """Return the display label for a live trace entry."""
    if isinstance(entry, dict):
        return str(entry.get("label") or entry.get("name") or "unknown")
    return str(entry)


def _trace_structure_hint(label: str) -> str:
    """Extract a structure basename from a compact trace label."""
    text = str(label)
    match = re.search(r"\bon\s+([A-Za-z0-9_.-]+\.xyz)\b", text)
    if match:
        return match.group(1)
    match = re.search(r"\b(output_[A-Za-z0-9_.-]+\.json)\b", text)
    if match:
        return _structure_hint_from_path(match.group(1))
    return ""


def _is_failure_label(label: str) -> bool:
    """Return whether a tool/status label represents a failed result."""
    lowered = str(label).lower()
    return any(
        marker in lowered
        for marker in (
            "(failure)",
            "(failed)",
            "(error)",
            "(incomplete)",
            "(aborted)",
            "(blocked)",
            " failed",
            " failure",
            " error",
            " incomplete",
            " aborted",
            " blocked",
        )
    )


def _stage_status(
    completed_labels: list[str], active_labels: list[str], keys: list[str]
) -> tuple[str, str]:
    """Return CSS class and display text for a workflow stage."""
    if any(any(key in label for key in keys) for label in active_labels):
        return ("running", "running")
    matching_completed = [
        label
        for label in completed_labels
        if any(key in label for key in keys)
    ]
    if not matching_completed:
        return ("waiting", "waiting")
    last_matching = matching_completed[-1]
    if _is_failure_label(last_matching):
        return ("failed", "failed")
    if any(_is_failure_label(label) for label in matching_completed[:-1]):
        return ("recovered", "recovered")
    return ("done", "done")


def _simulation_stage(labels: list[str]) -> tuple[str, str, str]:
    """Infer the simulation stage from live tool labels."""
    joined = " ".join(labels).lower()
    if "run_ase: ir" in joined or "ir/" in joined or "/ir" in joined:
        return (
            "IR spectrum",
            "ASE",
            "Compute vibrational intensities and generate an infrared spectrum.",
        )
    if "thermo" in joined:
        return (
            "Thermochemistry",
            "ASE",
            "Compute enthalpy, entropy, and Gibbs corrections from the structure.",
        )
    if "dipole" in joined:
        return ("Dipole", "ASE", "Compute the molecular dipole moment from the geometry.")
    if "vib" in joined:
        return ("Vibrations", "ASE", "Compute vibrational frequencies from the structure.")
    if "opt" in joined:
        return ("Optimize", "ASE", "Relax the geometry before reporting properties.")
    return ("Energy / property", "ASE", "Run the requested calculator on the XYZ structure.")


def _calculator_route_info(labels: list[str], sim_title: str) -> dict[str, str]:
    """Infer the selected calculator and concise routing rationale from trace labels."""
    joined = " ".join(labels).lower()
    selected = "Waiting"
    reason = "Select an ASE-compatible calculator after the molecule and XYZ structure are available."

    if "tblite" in joined or "xtb" in joined or "gfn" in joined:
        selected = "TBLite / xTB"
        reason = (
            "Fast molecular backend; preferred here because IR, vibrations, "
            "thermochemistry, and dipole tasks need electronic/force information."
        )
    elif "mace" in joined:
        selected = "MACE / ML potential"
        reason = (
            "Machine-learned potential route; useful for fast geometry/energy-style "
            "work when the chemistry is inside the model domain."
        )
    elif "emt" in joined:
        selected = "EMT"
        reason = (
            "Lightweight ASE demo calculator; good for plumbing tests, not for "
            "serious organic electronic properties."
        )
    elif "orca" in joined:
        selected = "ORCA"
        reason = "External quantum-chemistry engine selected; requires a working ORCA executable/profile."
    elif "nwchem" in joined:
        selected = "NWChem"
        reason = "External quantum-chemistry engine selected; requires a working NWChem executable/profile."
    elif "run_ase:" in joined:
        selected = "ASE default"
        reason = (
            "The run reached ASE without an explicit calculator label; verify the "
            "configured default is suitable for this property."
        )

    if selected == "Waiting" and sim_title in {"IR spectrum", "Vibrations", "Thermochemistry", "Dipole"}:
        reason = (
            "For this molecular property, TBLite/xTB is the default low-cost option "
            "unless the user explicitly requests MACE, EMT, ORCA, or NWChem."
        )

    return {
        "selected": selected,
        "reason": reason,
        "options": "TBLite/xTB, MACE/ML, EMT, ORCA, NWChem",
    }


def _calculator_stage_status(
    completed_labels: list[str],
    active_labels: list[str],
) -> tuple[str, str]:
    """Return workflow status for the calculator selection stage."""
    labels = completed_labels + active_labels
    if not any("run_ase" in label for label in labels):
        return ("waiting", "waiting")
    if any("run_ase" in label for label in active_labels):
        return ("done", "selected")
    matching_completed = [label for label in completed_labels if "run_ase" in label]
    if not matching_completed:
        return ("waiting", "waiting")
    last_matching = matching_completed[-1]
    if _is_failure_label(last_matching):
        return ("failed", "check")
    if any(_is_failure_label(label) for label in matching_completed[:-1]):
        return ("recovered", "recovered")
    return ("done", "selected")


def _workflow_trace_diagram_html(completed: list[str], active: list[dict]) -> str:
    """Build a diagram-style workflow explanation without raw model thoughts."""
    completed_labels = [str(label) for label in completed]
    active_labels = [_trace_entry_label(item) for item in active]
    all_labels = completed_labels + active_labels
    sim_title, sim_token, sim_why = _simulation_stage(all_labels)
    calculator_info = _calculator_route_info(all_labels, sim_title)
    species_state, species_text = _stage_status(
        completed_labels,
        active_labels,
        ["molecule_name_to_smiles", "resolve_molecule_identity"],
    )
    coords_state, coords_text = _stage_status(
        completed_labels, active_labels, ["smiles_to_coordinate_file"]
    )
    calculator_state, calculator_text = _calculator_stage_status(
        completed_labels, active_labels
    )
    sim_state, sim_text = _stage_status(completed_labels, active_labels, ["run_ase"])
    composer_ready = sim_state in {"done", "recovered"} and not active_labels
    composer_state = "done" if composer_ready else "waiting"
    composer_text = "ready" if composer_state == "done" else "waiting"

    stages = [
        (
            "done",
            "planned",
            "Q",
            "Intent",
            "Identify the requested property and required chemistry workflow.",
        ),
        (
            species_state,
            species_text,
            "ID",
            "Species",
            "Resolve molecule names to explicit SMILES identities.",
        ),
        (
            coords_state,
            coords_text,
            "XYZ",
            "3D structure",
            "Generate coordinates because ASE calculators need atoms in space.",
        ),
        (
            calculator_state,
            calculator_text,
            "CALC",
            "Calculator",
            "Select the simulation backend and check property/backend fit.",
        ),
        (sim_state, sim_text, sim_token, sim_title, sim_why),
        (
            composer_state,
            composer_text,
            "OUT",
            "Composer",
            "Validate tool outputs and turn them into the final answer.",
        ),
    ]

    stage_cards = []
    for state, status, token, title, why in stages:
        stage_cards.append(
            f"""
            <div class="cg-stage cg-stage-{state}">
              <div class="cg-stage-token">{html.escape(token)}</div>
              <div class="cg-stage-title">{html.escape(title)}</div>
              <div class="cg-stage-status">{html.escape(status)}</div>
              <div class="cg-stage-why">{html.escape(why)}</div>
            </div>
            """
        )

    return f"""
    <style>
    .cg-workflow-panel {{
      display: grid;
      gap: 0.85rem;
      padding: 0.2rem 0 0.1rem;
    }}
    .cg-workflow-heading {{
      display: flex;
      align-items: center;
      justify-content: space-between;
      gap: 1rem;
      font-weight: 700;
      color: #f5f7fb;
    }}
    .cg-workflow-heading small {{
      color: #9aa4b2;
      font-weight: 500;
    }}
    .cg-workflow-diagram {{
      display: grid;
      grid-template-columns: repeat(6, minmax(120px, 1fr));
      gap: 0.75rem;
      align-items: stretch;
    }}
    .cg-stage {{
      position: relative;
      min-height: 130px;
      border: 1px solid #303743;
      border-radius: 8px;
      background: #111820;
      padding: 0.85rem 0.75rem;
      display: grid;
      grid-template-rows: auto auto auto 1fr;
      gap: 0.35rem;
    }}
    .cg-stage:not(:last-child)::after {{
      content: "";
      position: absolute;
      top: 50%;
      right: -0.78rem;
      width: 0.78rem;
      height: 2px;
      background: #3a4350;
      transform: translateY(-50%);
    }}
    .cg-stage-token {{
      width: 2.25rem;
      height: 2.25rem;
      border-radius: 999px;
      display: inline-flex;
      align-items: center;
      justify-content: center;
      background: #1d2632;
      color: #dce6f2;
      font-size: 0.78rem;
      font-weight: 800;
      letter-spacing: 0;
    }}
    .cg-stage-title {{
      color: #f8fafc;
      font-weight: 800;
      line-height: 1.2;
    }}
    .cg-stage-status {{
      justify-self: start;
      border-radius: 999px;
      padding: 0.12rem 0.5rem;
      font-size: 0.72rem;
      font-weight: 700;
      text-transform: uppercase;
      letter-spacing: 0;
    }}
    .cg-stage-why {{
      color: #aab4c1;
      font-size: 0.78rem;
      line-height: 1.35;
    }}
    .cg-stage-done {{
      border-color: #2f8f5b;
      background: #0d1d18;
    }}
    .cg-stage-done .cg-stage-token,
    .cg-stage-done .cg-stage-status {{
      background: #1d6f46;
      color: #eafff3;
    }}
    .cg-stage-recovered {{
      border-color: #2f8f5b;
      background: #0d1d18;
    }}
    .cg-stage-recovered .cg-stage-token,
    .cg-stage-recovered .cg-stage-status {{
      background: #1d6f46;
      color: #eafff3;
    }}
    .cg-stage-running {{
      border-color: #c8932f;
      background: #21190b;
    }}
    .cg-stage-running .cg-stage-token,
    .cg-stage-running .cg-stage-status {{
      background: #8b641d;
      color: #fff5d8;
    }}
    .cg-stage-failed {{
      border-color: #9f3a3a;
      background: #211010;
    }}
    .cg-stage-failed .cg-stage-token,
    .cg-stage-failed .cg-stage-status {{
      background: #8f2f2f;
      color: #fff0f0;
    }}
    .cg-stage-failed .cg-stage-why {{
      color: #f0b8b8;
    }}
    .cg-stage-waiting {{
      border-color: #394150;
      background: #0f141b;
    }}
    .cg-stage-waiting .cg-stage-status {{
      background: #2a3140;
      color: #b8c1cf;
    }}
    .cg-workflow-caveat {{
      color: #9aa4b2;
      border-left: 3px solid #3a4350;
      padding-left: 0.65rem;
      font-size: 0.78rem;
      line-height: 1.35;
    }}
    .cg-calculator-route {{
      display: grid;
      grid-template-columns: minmax(140px, 0.7fr) minmax(240px, 1.5fr) minmax(220px, 1fr);
      gap: 0.7rem;
      border: 1px solid #303743;
      border-radius: 8px;
      background: #0f141b;
      padding: 0.7rem;
    }}
    .cg-route-block {{
      min-width: 0;
    }}
    .cg-route-label {{
      color: #94a3b8;
      font-size: 0.72rem;
      font-weight: 800;
      text-transform: uppercase;
      letter-spacing: 0;
      margin-bottom: 0.25rem;
    }}
    .cg-route-value {{
      color: #f8fafc;
      font-size: 0.84rem;
      line-height: 1.35;
      overflow-wrap: anywhere;
    }}
    .cg-route-options {{
      display: flex;
      flex-wrap: wrap;
      gap: 0.35rem;
    }}
    .cg-option-chip {{
      border: 1px solid #334155;
      border-radius: 999px;
      color: #dbeafe;
      background: #172033;
      padding: 0.16rem 0.45rem;
      font-size: 0.75rem;
      line-height: 1.2;
    }}
    @media (max-width: 1100px) {{
      .cg-workflow-diagram {{
        grid-template-columns: repeat(2, minmax(150px, 1fr));
      }}
      .cg-stage:not(:last-child)::after {{
        display: none;
      }}
      .cg-calculator-route {{
        grid-template-columns: 1fr;
      }}
    }}
    </style>
    <div class="cg-workflow-panel">
      <div class="cg-workflow-heading">
        <span>Workflow diagram</span>
        <small>validated tool path, not raw model thoughts</small>
      </div>
      <div class="cg-workflow-diagram">
        {''.join(stage_cards)}
      </div>
      <div class="cg-calculator-route">
        <div class="cg-route-block">
          <div class="cg-route-label">Selected calculator</div>
          <div class="cg-route-value">{html.escape(calculator_info["selected"])}</div>
        </div>
        <div class="cg-route-block">
          <div class="cg-route-label">Reason</div>
          <div class="cg-route-value">{html.escape(calculator_info["reason"])}</div>
        </div>
        <div class="cg-route-block">
          <div class="cg-route-label">Options</div>
          <div class="cg-route-options">
            {''.join(f'<span class="cg-option-chip">{html.escape(option.strip())}</span>' for option in calculator_info["options"].split(","))}
          </div>
        </div>
      </div>
      <div class="cg-workflow-caveat">
        Check identity, charge/spin, calculator domain, convergence, and units before using numeric outputs.
      </div>
    </div>
    """


def _tool_status_parts(label: str) -> tuple[str, str]:
    """Split a compact tool label into tool name and operation detail."""
    if ": " not in label:
        return (label, "")
    tool_name, detail = label.split(": ", 1)
    return (tool_name, detail)


def _workflow_tool_status_html(completed: list[str], active: list[dict]) -> str:
    """Build an illustrative tool status board with escaped dynamic labels."""
    rows = []
    entries = [
        ("failed" if _is_failure_label(str(label)) else "done", str(label), "failed" if _is_failure_label(str(label)) else "done")
        for label in completed
    ] + [
        ("running", _trace_entry_label(item), "running")
        for item in active
    ]

    for state, label, status in entries:
        tool_name, detail = _tool_status_parts(label)
        rows.append(
            f"""
            <div class="cg-tool-row cg-tool-row-{state}">
              <div class="cg-tool-dot"></div>
              <div class="cg-tool-body">
                <div class="cg-tool-name">{html.escape(tool_name)}</div>
                <div class="cg-tool-detail">{html.escape(detail or label)}</div>
              </div>
              <div class="cg-tool-badge">{html.escape(status)}</div>
            </div>
            """
        )

    if not rows:
        rows.append(
            """
            <div class="cg-tool-row cg-tool-row-waiting">
              <div class="cg-tool-dot"></div>
              <div class="cg-tool-body">
                <div class="cg-tool-name">Waiting</div>
                <div class="cg-tool-detail">No tool call has started yet.</div>
              </div>
              <div class="cg-tool-badge">waiting</div>
            </div>
            """
        )

    return f"""
    <style>
    .cg-tool-board {{
      display: grid;
      gap: 0.45rem;
      margin-top: 0.7rem;
    }}
    .cg-tool-board-title {{
      color: #f5f7fb;
      font-weight: 800;
      margin-bottom: 0.15rem;
    }}
    .cg-tool-row {{
      display: grid;
      grid-template-columns: 0.8rem minmax(0, 1fr) auto;
      align-items: center;
      gap: 0.7rem;
      border: 1px solid #303743;
      border-radius: 8px;
      padding: 0.55rem 0.65rem;
      background: #10161d;
    }}
    .cg-tool-dot {{
      width: 0.62rem;
      height: 0.62rem;
      border-radius: 999px;
      background: #6b7280;
    }}
    .cg-tool-name {{
      color: #dbe5f0;
      font-size: 0.74rem;
      font-weight: 800;
      text-transform: uppercase;
      letter-spacing: 0;
      line-height: 1.2;
    }}
    .cg-tool-detail {{
      color: #f8fafc;
      font-family: ui-monospace, SFMono-Regular, Menlo, Monaco, Consolas, monospace;
      font-size: 0.82rem;
      line-height: 1.25;
      overflow-wrap: anywhere;
    }}
    .cg-tool-badge {{
      border-radius: 999px;
      padding: 0.16rem 0.55rem;
      font-size: 0.72rem;
      font-weight: 800;
      text-transform: uppercase;
      color: #cbd5e1;
      background: #273142;
    }}
    .cg-tool-row-done {{
      border-color: #246d49;
      background: #0c1b16;
    }}
    .cg-tool-row-done .cg-tool-dot,
    .cg-tool-row-done .cg-tool-badge {{
      background: #1d6f46;
      color: #eafff3;
    }}
    .cg-tool-row-done .cg-tool-detail {{
      color: #25d66f;
      font-weight: 800;
    }}
    .cg-tool-row-running {{
      border-color: #9a7124;
      background: #20180c;
    }}
    .cg-tool-row-running .cg-tool-dot,
    .cg-tool-row-running .cg-tool-badge {{
      background: #8b641d;
      color: #fff5d8;
    }}
    .cg-tool-row-failed {{
      border-color: #9f3a3a;
      background: #211010;
    }}
    .cg-tool-row-failed .cg-tool-dot,
    .cg-tool-row-failed .cg-tool-badge {{
      background: #8f2f2f;
      color: #fff0f0;
    }}
    .cg-tool-row-failed .cg-tool-detail {{
      color: #ffb4b4;
      font-weight: 800;
    }}
    .cg-tool-row-waiting {{
      border-color: #394150;
      background: #0f141b;
    }}
    </style>
    <div class="cg-tool-board">
      <div class="cg-tool-board-title">Tool status</div>
      {''.join(rows)}
    </div>
    """


def _stream_workflow(stream_input, config, agent, msg_queue):
    """Run the agent workflow in a background thread, pushing events to a queue.

    Events pushed:
        ("tool_call", [tool_names])   — agent is calling tool(s)
        ("tool_result", tool_name)    — a tool finished
        ("interrupt", question_str)
        ("done", last_state)
        ("error", exception)

    Parameters
    ----------
    stream_input : dict or Command
        Initial workflow input or resume command.
    config : dict
        LangGraph run configuration.
    agent : ChemGraph
        Active ChemGraph agent.
    msg_queue : queue.Queue
        Queue receiving stream events for the UI thread.
    """
    from langgraph.errors import GraphInterrupt

    async def _run():
        """Stream the workflow and enqueue UI events."""
        prev_msgs: list = []
        last_st = None
        interrupt_val = None

        try:
            async for s in agent.workflow.astream(
                stream_input, stream_mode="values", config=config
            ):
                if "__interrupt__" in s:
                    int_data = s["__interrupt__"]
                    if isinstance(int_data, (list, tuple)) and int_data:
                        interrupt_val = int_data[0].value
                    elif hasattr(int_data, "value"):
                        interrupt_val = int_data.value
                    else:
                        interrupt_val = {"question": "The workflow needs your input."}

                if "messages" in s and s["messages"] != prev_msgs:
                    new_message = s["messages"][-1]
                    classified = _classify_message(new_message)
                    if classified:
                        msg_queue.put(classified)
                    prev_msgs = s["messages"]
                last_st = s
        except GraphInterrupt as gi:
            interrupts = gi.args[0] if gi.args else []
            if interrupts:
                interrupt_val = interrupts[0].value
            else:
                interrupt_val = {"question": "The workflow needs your input."}

        # Check checkpoint for pending interrupts
        if interrupt_val is None:
            try:
                snapshot = agent.workflow.get_state(config)
                if snapshot and snapshot.tasks:
                    for t in snapshot.tasks:
                        t_interrupts = getattr(t, "interrupts", None)
                        if t_interrupts:
                            interrupt_val = t_interrupts[0].value
                            break
            except Exception:
                pass

        if interrupt_val is not None:
            if isinstance(interrupt_val, dict):
                q = interrupt_val.get(
                    "question",
                    interrupt_val.get("message", str(interrupt_val)),
                )
            else:
                q = str(interrupt_val)
            msg_queue.put(("interrupt", q))
        else:
            msg_queue.put(("done", last_st))

    try:
        asyncio.run(_run())
    except HumanInputRequired as hir:
        msg_queue.put(("interrupt", hir.question))
    except Exception as exc:
        msg_queue.put(("error", exc))


def _poll_and_display(msg_queue, status_container, placeholder, thread):
    """Poll the message queue and render a compact tool-call log.

    Uses a single ``st.empty()`` placeholder to re-render the full list
    each time, so completed tools get a checkmark and only the active
    tool shows a spinner indicator.

    Returns:
        ("done", last_state) | ("interrupt", question) | ("error", exception)

    Parameters
    ----------
    msg_queue : queue.Queue
        Queue receiving stream events.
    status_container : DeltaGenerator
        Streamlit status container.
    placeholder : DeltaGenerator
        Placeholder used for the tool-call log.
    thread : threading.Thread
        Background stream thread.

    Returns
    -------
    tuple
        ``("done", state)``, ``("interrupt", question)``, or
        ``("error", exception)``.
    """
    completed: list[str] = []  # tools that finished
    active: list[dict] = []  # tools currently running

    def _snapshot() -> dict[str, list]:
        return {
            "completed": list(completed),
            "active": [dict(item) for item in active if isinstance(item, dict)],
        }

    def _render():
        """Render the current tool-call status list."""
        with placeholder.container():
            st.html(_workflow_trace_diagram_html(completed, active))
            st.html(_workflow_tool_status_html(completed, active))

    def _active_labels() -> list[str]:
        return [_trace_entry_label(item) for item in active]

    def _pop_active_for_result(result: Any) -> Optional[dict]:
        return _pop_active_trace_match(active, result)

    while True:
        try:
            event_type, event_data = msg_queue.get(timeout=0.1)
        except queue.Empty:
            if not thread.is_alive():
                try:
                    event_type, event_data = msg_queue.get_nowait()
                except queue.Empty:
                    return ("error", RuntimeError("Stream ended without result."))
            else:
                continue

        if event_type == "tool_call":
            active.clear()
            active.extend(event_data)
            label = ", ".join(_active_labels())
            status_container.update(label=f"Running {label}", state="running")
            _render()
        elif event_type == "tool_result":
            # Move this specific tool from active to completed
            active_item = _pop_active_for_result(event_data)
            result_label = (
                str(event_data.get("label"))
                if isinstance(event_data, dict) and event_data.get("label")
                else _trace_entry_label(active_item or event_data)
            )
            if result_label not in completed:
                completed.append(result_label)
            if active:
                status_container.update(
                    label=f"Running {', '.join(_active_labels())}", state="running"
                )
            else:
                status_container.update(label="Thinking...", state="running")
            _render()
        elif event_type in ("done", "interrupt", "error"):
            # Only tool_result events are completed. Unmatched tool calls mean the
            # graph stopped before the tool returned and should not render green.
            for label in _active_labels():
                incomplete = f"{label} (incomplete)"
                if incomplete not in completed:
                    completed.append(incomplete)
            active.clear()
            _render()
            return (event_type, event_data, _snapshot())


def _find_conversation_entry_by_log_dir(log_dir: Optional[str]) -> Optional[dict]:
    """Return the visible exchange tied to a query artifact directory."""
    if not log_dir:
        return None
    for entry in reversed(st.session_state.get("conversation_history", [])):
        if entry.get("log_dir") == log_dir:
            return entry
    return None


def _error_result(error: Any) -> dict:
    """Return a persistable assistant message for workflow-level errors."""
    return {
        "messages": [
            {
                "type": "ai",
                "content": f"Processing error: {error}",
            }
        ],
        "final_output": None,
        "scientific_ledger": None,
        "validation": None,
        "completion_validation": None,
        "run_context": None,
    }


def _queue_query_submission(
    query: str,
    thread_id: int,
    endpoint_status: dict,
    selected_base_url: Optional[str],
) -> None:
    """Reserve a workflow run and rerun so the input renders disabled first."""
    if not endpoint_status["ok"]:
        msg = (
            f"Cannot reach local model endpoint `{selected_base_url}`. "
            f"{endpoint_status['message']}"
        )
        st.session_state.last_run_error = RuntimeError(msg)
        st.error(msg)
        return
    if not st.session_state.agent:
        st.error("Agent not ready. Please check configuration and try again.")
        return

    trimmed_query = query.strip()
    if not trimmed_query:
        return

    if not _begin_agent_run():
        st.info("A ChemGraph task is already running. Please wait for it to finish.")
        return

    agent = st.session_state.agent
    _ensure_chat_log_dir()

    # Create the persistent sessions while the agent still points to the
    # chat-level directory; calculation artifacts go into query dirs below.
    try:
        agent._ensure_session(trimmed_query)
    except Exception:
        pass
    query_saved = _save_query_to_store(trimmed_query)

    query_log_dir = _create_query_log_dir()
    run_thread_id = _ensure_graph_thread_id(thread_id)
    st.session_state.last_run_query = trimmed_query
    st.session_state.last_run_error = None
    st.session_state.last_run_result = None

    st.session_state.conversation_history.append(
        {
            "query": trimmed_query,
            "result": {"messages": []},
            "thread_id": run_thread_id,
            "log_dir": query_log_dir,
            "status": "running",
        }
    )
    st.session_state.pending_agent_submission = {
        "kind": "query",
        "query": trimmed_query,
        "thread_id": thread_id,
        "run_thread_id": run_thread_id,
        "query_log_dir": query_log_dir,
        "query_saved": query_saved,
    }
    st.rerun()


def _execute_pending_agent_submission(
    thread_id: int,
    endpoint_status: dict,
    selected_base_url: Optional[str],
) -> None:
    """Run a reserved pending submission after the disabled input is rendered."""
    pending = st.session_state.get("pending_agent_submission")
    if not pending or pending.get("kind") != "query":
        return

    _handle_query_submission(
        pending["query"],
        pending.get("thread_id", thread_id),
        endpoint_status,
        selected_base_url,
        run_reserved=True,
        query_log_dir=pending.get("query_log_dir"),
        run_thread_id=pending.get("run_thread_id"),
        query_saved=bool(pending.get("query_saved")),
    )


def _handle_query_submission(
    query: str,
    thread_id: int,
    endpoint_status: dict,
    selected_base_url: Optional[str],
    *,
    run_reserved: bool = False,
    query_log_dir: Optional[str] = None,
    run_thread_id: Optional[str] = None,
    query_saved: bool = False,
) -> None:
    """Handle a submitted user query and stream the workflow response.

    Parameters
    ----------
    query : str
        User query text.
    thread_id : int
        Current LangGraph thread ID.
    endpoint_status : dict
        Local endpoint status dictionary.
    selected_base_url : str, optional
        Model endpoint URL used in error messages.
    """
    if not endpoint_status["ok"]:
        msg = (
            f"Cannot reach local model endpoint `{selected_base_url}`. "
            f"{endpoint_status['message']}"
        )
        st.session_state.last_run_error = RuntimeError(msg)
        st.error(msg)
        return
    if not st.session_state.agent:
        st.error("Agent not ready. Please check configuration and try again.")
        return
    if not query.strip():
        return

    trimmed_query = query.strip()
    if not _begin_agent_run(from_pending=run_reserved):
        st.info("A ChemGraph task is already running. Please wait for it to finish.")
        return

    agent = st.session_state.agent
    old_agent_log_dir = getattr(agent, "log_dir", None)
    old_env_log_dir = None
    try:
        _ensure_chat_log_dir()
        old_env_log_dir = os.environ.get("CHEMGRAPH_LOG_DIR")

        # Create the persistent session while the agent still points to the
        # chat-level directory; only calculation artifacts go into query dirs.
        try:
            agent._ensure_session(trimmed_query)
        except Exception:
            pass

        if not query_saved:
            query_saved = _save_query_to_store(trimmed_query)

        query_log_dir = query_log_dir or _create_query_log_dir()
        agent.log_dir = query_log_dir
        os.environ["CHEMGRAPH_LOG_DIR"] = query_log_dir

        run_thread_id = run_thread_id or _ensure_graph_thread_id(thread_id)
        cfg = {"configurable": {"thread_id": run_thread_id}}
        cfg["recursion_limit"] = agent.recursion_limit
        st.session_state.last_run_query = trimmed_query
        st.session_state.last_run_error = None
        st.session_state.last_run_result = None
        pending_entry = _find_conversation_entry_by_log_dir(query_log_dir)
        if pending_entry is None:
            pending_entry = {
                "query": trimmed_query,
                "result": {"messages": []},
                "thread_id": run_thread_id,
                "log_dir": query_log_dir,
                "status": "running",
            }
            st.session_state.conversation_history.append(pending_entry)
        else:
            pending_entry["status"] = "running"

        # Snapshot message count before streaming so we can isolate new messages
        prev_msg_count = 0
        try:
            snapshot = agent.workflow.get_state(cfg)
            if snapshot and snapshot.values:
                prev_msg_count = len(snapshot.values.get("messages", []))
        except Exception:
            pass

        # Show the user's message immediately
        with st.chat_message("user"):
            st.markdown(trimmed_query)

        # Stream agent response with live tool-call display
        with st.chat_message("assistant"):
            msg_q: queue.Queue = queue.Queue()
            inputs = {"messages": trimmed_query}

            stream_thread = threading.Thread(
                target=_stream_workflow,
                args=(inputs, cfg, agent, msg_q),
                daemon=True,
            )

            status = st.status("Thinking...", expanded=True)
            with status:
                tool_log = st.empty()
            stream_thread.start()
            event_type, event_data, workflow_trace = _poll_and_display(
                msg_q, status, tool_log, stream_thread
            )
            stream_thread.join(timeout=5)

            if event_type == "done":
                last_state = event_data
                if last_state is None:
                    msg = "Workflow produced no output."
                    result = _error_result(msg)
                    pending_entry.update(
                        {
                            "result": result,
                            "thread_id": run_thread_id,
                            "log_dir": query_log_dir,
                            "workflow_trace": workflow_trace,
                            "status": "error",
                            "error": msg,
                        }
                    )
                    _save_exchange_to_store(
                        trimmed_query,
                        result,
                        include_query=not query_saved,
                    )
                    st.session_state.pending_agent_submission = None
                    st.error(msg)
                    return

                # Only keep messages from this query (not prior thread history)
                all_msgs = last_state.get("messages", [])
                new_msgs = all_msgs[prev_msg_count:]
                result = {
                    "messages": new_msgs,
                    "final_output": last_state.get("final_output"),
                    "scientific_ledger": last_state.get("scientific_ledger"),
                    "validation": last_state.get("validation"),
                    "completion_validation": last_state.get("completion_validation"),
                    "run_context": last_state.get("run_context"),
                }
                workflow_status = _workflow_status_from_result(result)
                status.update(
                    label=_workflow_status_label(workflow_status),
                    state=_streamlit_state_for_workflow_status(workflow_status),
                    expanded=False,
                )

                # Save messages to persistent session store (best-effort)
                try:
                    agent._save_messages_to_store(last_state, trimmed_query)
                except Exception:
                    pass

                st.session_state.last_run_result = result
                pending_entry.update(
                    {
                        "query": trimmed_query,
                        "result": result,
                        "thread_id": run_thread_id,
                        "log_dir": query_log_dir,
                        "workflow_trace": workflow_trace,
                        "status": workflow_status,
                    }
                )
                _save_exchange_to_store(
                    trimmed_query,
                    result,
                    include_query=not query_saved,
                )
                st.session_state.query_input = ""
                st.session_state.pending_agent_submission = None
                st.rerun()

            elif event_type == "interrupt":
                status.update(
                    label="Waiting for input", state="complete", expanded=False
                )
                cfg_for_resume = dict(cfg)
                st.session_state.pending_human_question = event_data
                st.session_state.pending_interrupt_config = cfg_for_resume
                st.session_state.pending_interrupt_query = trimmed_query
                st.session_state.pending_interrupt_thread_id = thread_id
                st.session_state.pending_interrupt_prev_msg_count = prev_msg_count
                st.session_state.pending_interrupt_model = st.session_state.get(
                    "active_model"
                )
                st.session_state.pending_interrupt_workflow = st.session_state.get(
                    "active_workflow"
                )
                st.session_state.pending_interrupt_log_dir = query_log_dir
                st.session_state.interrupt_count = 1
                st.session_state.interrupt_exchanges = []
                pending_entry.update(
                    {
                        "status": "waiting_for_input",
                        "workflow_trace": workflow_trace,
                        "pending_human_question": event_data,
                        "pending_interrupt_config": cfg_for_resume,
                        "pending_interrupt_query": trimmed_query,
                        "pending_interrupt_thread_id": thread_id,
                        "pending_interrupt_prev_msg_count": prev_msg_count,
                        "pending_interrupt_model": st.session_state.get(
                            "active_model"
                        ),
                        "pending_interrupt_workflow": st.session_state.get(
                            "active_workflow"
                        ),
                        "pending_interrupt_log_dir": query_log_dir,
                        "interrupt_count": 1,
                        "interrupt_exchanges": [],
                    }
                )
                st.session_state.pending_agent_submission = None
                st.rerun()

            else:  # error
                status.update(label="Error", state="error", expanded=False)
                st.session_state.last_run_error = event_data
                pending_entry["status"] = "error"
                pending_entry["error"] = str(event_data)
                pending_entry["workflow_trace"] = workflow_trace
                result = _error_result(event_data)
                pending_entry["result"] = result
                _save_exchange_to_store(
                    trimmed_query,
                    result,
                    include_query=not query_saved,
                )
                st.session_state.pending_agent_submission = None
                st.error(f"Processing error: {event_data}")
                st.rerun()
    finally:
        pending_submission = st.session_state.get("pending_agent_submission")
        if (
            isinstance(pending_submission, dict)
            and pending_submission.get("query_log_dir") == query_log_dir
        ):
            st.session_state.pending_agent_submission = None
        _restore_agent_log_context(agent, old_agent_log_dir, old_env_log_dir)
        _end_agent_run()


def _handle_human_response(answer: str, thread_id: int) -> None:
    """Resume the agent workflow with the human's answer.

    Parameters
    ----------
    answer : str
        Human response to the pending interrupt question.
    thread_id : int
        Current LangGraph thread ID.
    """
    from langgraph.types import Command

    agent = st.session_state.agent
    resume_config = st.session_state.pending_interrupt_config
    original_query = st.session_state.pending_interrupt_query
    current_question = st.session_state.pending_human_question
    interrupt_count = st.session_state.interrupt_count

    if agent is None or resume_config is None:
        st.error("Agent was re-initialized. Please submit your query again.")
        _clear_interrupt_state()
        return
    if not _begin_agent_run():
        st.info("A ChemGraph task is already running. Please wait for it to finish.")
        return

    MAX_INTERRUPTS = 10
    old_agent_log_dir = getattr(agent, "log_dir", None)
    old_env_log_dir = os.environ.get("CHEMGRAPH_LOG_DIR")
    resume_log_dir = (
        st.session_state.get("pending_interrupt_log_dir")
        or st.session_state.get("active_query_log_dir")
        or _create_query_log_dir()
    )

    try:
        agent.log_dir = resume_log_dir
        os.environ["CHEMGRAPH_LOG_DIR"] = resume_log_dir
        pending_entry = _find_conversation_entry_by_log_dir(resume_log_dir)
        if pending_entry is not None:
            pending_entry["status"] = "running"

        # Record this exchange
        st.session_state.interrupt_exchanges.append(
            {"question": current_question, "answer": answer}
        )

        # Show the user's reply immediately
        with st.chat_message("user"):
            st.markdown(answer)

        # Stream resumed agent response
        with st.chat_message("assistant"):
            msg_q: queue.Queue = queue.Queue()
            resume_cmd = Command(resume=answer)

            stream_thread = threading.Thread(
                target=_stream_workflow,
                args=(resume_cmd, resume_config, agent, msg_q),
                daemon=True,
            )

            status = st.status("Processing your response...", expanded=True)
            with status:
                tool_log = st.empty()
            stream_thread.start()
            event_type, event_data, workflow_trace = _poll_and_display(
                msg_q, status, tool_log, stream_thread
            )
            stream_thread.join(timeout=5)

            if event_type == "done":
                result_state = event_data

                if result_state is None:
                    msg = "Resume produced no output."
                    final_result = _error_result(msg)
                    if pending_entry is not None:
                        pending_entry.update(
                            {
                                "query": original_query,
                                "result": final_result,
                                "thread_id": thread_id,
                                "log_dir": resume_log_dir,
                                "workflow_trace": workflow_trace,
                                "status": "error",
                                "error": msg,
                            }
                        )
                    _save_exchange_to_store(
                        original_query,
                        final_result,
                        include_query=pending_entry is None,
                    )
                    st.error(msg)
                    _clear_interrupt_state()
                    return

                # Only keep messages from this query (not prior thread history)
                prev_msg_count = st.session_state.get(
                    "pending_interrupt_prev_msg_count", 0
                )
                all_msgs = result_state.get("messages", [])
                new_msgs = all_msgs[prev_msg_count:]
                final_result = {
                    "messages": new_msgs,
                    "final_output": result_state.get("final_output"),
                    "scientific_ledger": result_state.get("scientific_ledger"),
                    "validation": result_state.get("validation"),
                    "completion_validation": result_state.get("completion_validation"),
                    "run_context": result_state.get("run_context"),
                }
                workflow_status = _workflow_status_from_result(final_result)
                status.update(
                    label=_workflow_status_label(workflow_status),
                    state=_streamlit_state_for_workflow_status(workflow_status),
                    expanded=False,
                )

                exchanges = list(st.session_state.interrupt_exchanges)
                st.session_state.last_run_result = final_result
                entry_payload = {
                    "query": original_query,
                    "result": final_result,
                    "thread_id": thread_id,
                    "log_dir": resume_log_dir,
                    "interrupt_exchanges": exchanges,
                    "workflow_trace": workflow_trace,
                    "status": workflow_status,
                }
                if pending_entry is not None:
                    pending_entry.update(entry_payload)
                else:
                    st.session_state.conversation_history.append(entry_payload)
                _save_exchange_to_store(
                    original_query,
                    final_result,
                    include_query=pending_entry is None,
                )
                st.session_state.query_input = ""
                _clear_interrupt_state()
                st.rerun()

            elif event_type == "interrupt":
                status.update(
                    label="Waiting for input", state="complete", expanded=False
                )
                new_count = interrupt_count + 1
                if new_count > MAX_INTERRUPTS:
                    st.error(
                        "Agent exceeded maximum number of follow-up questions. Aborting."
                    )
                    if pending_entry is not None:
                        pending_entry["status"] = "error"
                        pending_entry["error"] = "Exceeded maximum interrupts."
                    _clear_interrupt_state()
                    return
                st.session_state.pending_human_question = event_data
                st.session_state.interrupt_count = new_count
                if pending_entry is not None:
                    pending_entry.update(
                        {
                            "status": "waiting_for_input",
                            "workflow_trace": workflow_trace,
                            "pending_human_question": event_data,
                            "pending_interrupt_config": resume_config,
                            "pending_interrupt_query": original_query,
                            "pending_interrupt_thread_id": thread_id,
                            "pending_interrupt_prev_msg_count": st.session_state.get(
                                "pending_interrupt_prev_msg_count", 0
                            ),
                            "pending_interrupt_model": st.session_state.get(
                                "pending_interrupt_model"
                            ),
                            "pending_interrupt_workflow": st.session_state.get(
                                "pending_interrupt_workflow"
                            ),
                            "pending_interrupt_log_dir": resume_log_dir,
                            "interrupt_count": new_count,
                            "interrupt_exchanges": list(
                                st.session_state.interrupt_exchanges
                            ),
                        }
                    )
                st.rerun()

            else:  # error
                status.update(label="Error", state="error", expanded=False)
                st.session_state.last_run_error = event_data
                final_result = _error_result(event_data)
                if pending_entry is not None:
                    pending_entry["status"] = "error"
                    pending_entry["error"] = str(event_data)
                    pending_entry["workflow_trace"] = workflow_trace
                    pending_entry["result"] = final_result
                _save_exchange_to_store(
                    original_query,
                    final_result,
                    include_query=pending_entry is None,
                )
                st.error(f"Error during resume: {event_data}")
                _clear_interrupt_state()
                st.rerun()
    finally:
        _restore_agent_log_context(agent, old_agent_log_dir, old_env_log_dir)
        _end_agent_run()