backend/gradio_ui/app.py

# app.py
"""FinSight AI — light chat-first UI with + feature menu."""

from __future__ import annotations

import atexit
import logging

import gradio as gr

from core.bootstrap import init_services, shutdown_services
from gradio_ui.chat_tab import build_chat_panel
from gradio_ui.documents_tab import build_documents_panel
from gradio_ui.entities_tab import build_entities_panel
from gradio_ui.ocr_tab import build_ocr_panel
from gradio_ui.state import (
    DEFAULT_FOOTER,
    AppMode,
    SearchMode,
    center_title,
    default_app_mode,
    default_search_mode,
    default_selected_doc_ids,
    default_session_id,
    footer_hint,
    mode_label,
)
from gradio_ui.renderers import render_chat_transcript
from gradio_ui.summary_tab import build_summary_panel
from gradio_ui.theme import finsight_theme, load_custom_css

logger = logging.getLogger(__name__)
_services_initialized = False

# ── Feature popup items ────────────────────────────────────────────────────────
# Each tuple: (emoji-icon, label, feature-key)
FEATURE_ITEMS = [
    ("💬", "Finance QA Chatbot", "Finance QA Chatbot"),
    ("📄", "Financial Summary", "Financial Summary"),
    ("🔍", "Document OCR", "Document OCR"),
    ("🏷️", "Entity Extraction", "Entity Extraction"),
    ("📁", "Upload Documents", "Upload Documents"),
]

SEARCH_ITEMS = [
    ("⚡", "Hybrid RAG", "Hybrid RAG"),
    ("📌", "Single Document Search", "Single Document Search"),
]

UTILITY_ITEMS = [
    ("💾", "Export session", "Export session"),
]

TAB_IDS: dict[str, str] = {
    "Finance QA Chatbot": "qa",
    "Financial Summary": "summary",
    "Document OCR": "ocr",
    "Entity Extraction": "entities",
    "Upload Documents": "documents",
}

FEATURE_CHOICE_BY_MODE: dict[AppMode, str] = {
    "qa": "Finance QA Chatbot",
    "summary": "Financial Summary",
    "ocr": "Document OCR",
    "entities": "Entity Extraction",
    "documents": "Upload Documents",
}

POPUP_MENU_CHOICES = [
    "Finance QA Chatbot",
    "Financial Summary",
    "Document OCR",
    "Entity Extraction",
    "Upload Documents",
    "Hybrid RAG",
    "Single Document Search",
    "Export session",
]


def _popup_class_updates(app_mode: AppMode, search_mode: SearchMode) -> list[gr.update]:
    active_feature = FEATURE_CHOICE_BY_MODE.get(app_mode, "Finance QA Chatbot")
    updates: list[gr.update] = []
    for choice in POPUP_MENU_CHOICES:
        cls = ["fs-popup-menu-btn"]
        if choice == active_feature:
            cls.append("fs-popup-menu-btn-active")
        elif choice == "Hybrid RAG" and search_mode == "all":
            cls.append("fs-popup-menu-btn-active")
        elif choice == "Single Document Search" and search_mode == "single":
            cls.append("fs-popup-menu-btn-active")
        updates.append(gr.update(elem_classes=cls))
    return updates

FONT_LINK = """
<link rel="preconnect" href="https://fonts.googleapis.com" />
<link rel="preconnect" href="https://fonts.gstatic.com" crossorigin />
<link href="https://fonts.googleapis.com/css2?family=Inter:wght@400;500;600;700&family=JetBrains+Mono:wght@400;500&display=swap" rel="stylesheet" />
"""


def _view_updates(mode: AppMode):
    """Toggle QA main panel vs fixed feature overlays."""
    return (
        gr.update(visible=(mode == "qa")),
        gr.update(visible=(mode == "summary")),
        gr.update(visible=(mode == "ocr")),
        gr.update(visible=(mode == "entities")),
        gr.update(visible=(mode == "documents")),
        gr.update(visible=False),
    )


def _input_updates(app_mode: AppMode) -> tuple[gr.update, gr.update]:
    if app_mode == "qa":
        return (
            gr.update(
                placeholder="Ask about your financial documents…",
                interactive=True,
            ),
            gr.update(visible=True),
        )
    return (
        gr.update(
            placeholder="Tap + to switch features or upload documents",
            interactive=False,
        ),
        gr.update(visible=False),
    )


def _close_popup():
    return gr.update(visible=False), False


def _wire_feature_button(
    btn: gr.Button,
    choice: str,
    handler,
    inputs: list,
    outputs: list,
) -> None:
    def pick(search_mode, selected, session_id):
        return handler(choice, search_mode, selected, session_id)

    btn.click(pick, inputs=inputs, outputs=outputs)


def ensure_services() -> None:
    global _services_initialized
    if not _services_initialized:
        init_services()
        _services_initialized = True
        atexit.register(shutdown_services)


def build_app() -> tuple[gr.Blocks, gr.Theme, str]:
    theme = finsight_theme()
    css = load_custom_css()

    with gr.Blocks(title="FinSight AI", elem_classes=["fs-shell"]) as demo:
        gr.HTML(FONT_LINK, padding=False)

        app_mode_state = gr.State(default_app_mode())
        search_mode_state = gr.State(default_search_mode())
        selected_docs_state = gr.State(default_selected_doc_ids())
        chat_history_state = gr.State([])
        session_state = gr.State(default_session_id())
        popup_open = gr.State(False)

        session_dropdown = gr.Dropdown(choices=[], visible=False, elem_classes=["hidden-control"])

        with gr.Row(elem_classes=["fs-navbar"]):
            gr.HTML(
                """
                <div class="fs-brand">
                  <div class="fs-brand-icon">📊</div>
                  <div class="fs-brand-text">
                    <h1>FinSight AI</h1>
                    <p>Chat · RAG · Finance document intelligence</p>
                  </div>
                </div>
                """,
                padding=False,
            )
            mode_title = gr.HTML(
                '<div class="fs-nav-center">Hybrid RAG</div>',
                padding=False,
            )
            with gr.Row(elem_classes=["fs-nav-actions"]):
                gr.HTML(
                    """
                    <div class="fs-nav-meta">
                      <span class="fs-ai-badge">✨ AI-Powered</span>
                      <span class="fs-nav-hint">Navigate financial documents</span>
                    </div>
                    """,
                    padding=False,
                )
                new_chat_btn = gr.Button(
                    "🗑 New chat",
                    elem_classes=["fs-pill-btn"],
                    variant="secondary",
                    min_width=0,
                )
                mode_badge_btn = gr.Button(
                    "🔍 Hybrid RAG",
                    elem_classes=["fs-pill-btn", "fs-pill-btn-active"],
                    variant="secondary",
                    min_width=0,
                )

        with gr.Row(elem_classes=["fs-body"]):
            with gr.Column(scale=1, min_width=0, elem_classes=["fs-panel-area"]):
                with gr.Column(visible=True, elem_classes=["fs-qa-panel"]) as qa_panel:
                    chat = build_chat_panel()

        with gr.Column(visible=False, elem_classes=["fs-feature-overlay"]) as summary_overlay:
            summary = build_summary_panel()

        with gr.Column(visible=False, elem_classes=["fs-feature-overlay", "fs-feature-overlay-wide"]) as ocr_overlay:
            ocr = build_ocr_panel()

        with gr.Column(visible=False, elem_classes=["fs-feature-overlay"]) as entities_overlay:
            entities = build_entities_panel()

        with gr.Column(visible=False, elem_classes=["fs-feature-overlay"]) as documents_overlay:
            docs = build_documents_panel()

        feature_panel_outputs = [
            qa_panel,
            summary_overlay,
            ocr_overlay,
            entities_overlay,
            documents_overlay,
        ]

        with gr.Row(elem_classes=["fs-meta-row"], visible=False) as meta_row:
            gr.HTML("", padding=False, visible=False)
            gr.HTML("", padding=False, visible=False)

        with gr.Column(elem_classes=["fs-input-dock"], visible=True) as input_dock:
            with gr.Row(elem_classes=["fs-input-composer"]):
                with gr.Column(elem_classes=["fs-plus-anchor"], scale=0, min_width=44):
                    with gr.Column(
                        elem_classes=["fs-feature-popup-wrap"],
                        visible=False,
                    ) as feature_popup:
                        gr.HTML('<div class="fs-popup-section-label">Features</div>', padding=False)
                        feat_qa_btn = gr.Button("💬  Finance QA Chatbot", elem_classes=["fs-popup-menu-btn"])
                        feat_summary_btn = gr.Button("📄  Financial Summary", elem_classes=["fs-popup-menu-btn"])
                        feat_ocr_btn = gr.Button("🔍  Document OCR", elem_classes=["fs-popup-menu-btn"])
                        feat_entities_btn = gr.Button("🏷️  Entity Extraction", elem_classes=["fs-popup-menu-btn"])
                        feat_docs_btn = gr.Button("📁  Upload Documents", elem_classes=["fs-popup-menu-btn"])
                        gr.HTML('<div class="fs-popup-divider"></div>', padding=False)
                        gr.HTML('<div class="fs-popup-section-label">Search mode</div>', padding=False)
                        feat_hybrid_btn = gr.Button("⚡  Hybrid RAG", elem_classes=["fs-popup-menu-btn"])
                        feat_single_btn = gr.Button("📌  Single Document Search", elem_classes=["fs-popup-menu-btn"])
                        gr.HTML('<div class="fs-popup-divider"></div>', padding=False)
                        feat_export_btn = gr.Button("💾  Export session", elem_classes=["fs-popup-menu-btn"])
                    plus_btn = gr.Button(
                        "+",
                        elem_classes=["fs-plus-btn"],
                        elem_id="fs-plus-btn",
                        variant="secondary",
                        scale=0,
                        min_width=44,
                    )
                with gr.Row(elem_classes=["fs-input-pill"]):
                    msg = gr.Textbox(
                        placeholder="Ask about your financial documents…",
                        show_label=False,
                        lines=1,
                        max_lines=4,
                        scale=10,
                        container=False,
                    )
                    send_btn = gr.Button(
                        "➤",
                        elem_classes=["fs-send-btn"],
                        variant="secondary",
                        scale=0,
                        min_width=40,
                    )
            footer_md = gr.Markdown(
                DEFAULT_FOOTER,
                elem_classes=["fs-footer"],
            )

        popup_btns = [
            feat_qa_btn,
            feat_summary_btn,
            feat_ocr_btn,
            feat_entities_btn,
            feat_docs_btn,
            feat_hybrid_btn,
            feat_single_btn,
            feat_export_btn,
        ]

        panel_chrome_outputs = [*feature_panel_outputs, meta_row]

        def toggle_popup(is_open: bool):
            new_open = not is_open
            return new_open, gr.update(visible=new_open)

        plus_btn.click(
            toggle_popup,
            inputs=[popup_open],
            outputs=[popup_open, feature_popup],
        )

        def _header_updates(app_mode: AppMode, search_mode: SearchMode, selected: list[str]):
            title_label = center_title(app_mode, search_mode)
            badge_label = mode_label(search_mode)
            badge_text = (
                f"🔍 {badge_label}" if search_mode == "all" else f"📌 {badge_label}"
            )
            doc_count = len(docs["doc_choices"]())
            return (
                f'<div class="fs-nav-center">{title_label}</div>',
                badge_text,
                footer_hint(search_mode, selected, doc_count, app_mode),
            )

        def _with_chrome(
            app_mode: AppMode,
            search_mode: SearchMode,
            selected: list[str],
            base: tuple,
        ):
            msg_up, send_up = _input_updates(app_mode)
            return base + (
                msg_up,
                send_up,
                docs["doc_select_update"](selected),
            ) + tuple(_popup_class_updates(app_mode, search_mode))

        def on_feature_pick(choice, search_mode, selected, session_id):
            if choice == "Export session":
                path = chat["export_session_file"](session_id)
                title, badge, foot = _header_updates("qa", search_mode, selected)
                popup_close, popup_state = _close_popup()
                return _with_chrome(
                    "qa",
                    search_mode,
                    selected,
                    (
                        "qa",
                        search_mode,
                        selected,
                        *_view_updates("qa"),
                        title,
                        badge,
                        foot,
                        gr.update(value=path) if path else gr.update(),
                        popup_close,
                        popup_state,
                    ),
                )

            if choice == "Hybrid RAG":
                new_sel, foot = docs["on_search_mode_change"]("all", selected)
                title, badge, _ = _header_updates("qa", "all", new_sel)
                popup_close, popup_state = _close_popup()
                return _with_chrome(
                    "qa",
                    "all",
                    new_sel,
                    (
                        "qa",
                        "all",
                        new_sel,
                        *_view_updates("qa"),
                        title,
                        badge,
                        foot,
                        gr.update(),
                        popup_close,
                        popup_state,
                    ),
                )

            if choice == "Single Document Search":
                new_sel, foot = docs["on_search_mode_change"]("single", selected)
                title, badge, _ = _header_updates("qa", "single", new_sel)
                popup_close, popup_state = _close_popup()
                return _with_chrome(
                    "qa",
                    "single",
                    new_sel,
                    (
                        "qa",
                        "single",
                        new_sel,
                        *_view_updates("qa"),
                        title,
                        badge,
                        foot,
                        gr.update(),
                        popup_close,
                        popup_state,
                    ),
                )

            tab_id = TAB_IDS.get(choice, "qa")
            app_mode: AppMode = tab_id  # type: ignore[assignment]
            title, badge, foot = _header_updates(app_mode, search_mode, selected)
            popup_close, popup_state = _close_popup()
            return _with_chrome(
                app_mode,
                search_mode,
                selected,
                (
                    app_mode,
                    search_mode,
                    selected,
                    *_view_updates(app_mode),
                    title,
                    badge,
                    foot,
                    gr.update(),
                    popup_close,
                    popup_state,
                ),
            )

        chrome_outputs = popup_btns

        feature_pick_inputs = [search_mode_state, selected_docs_state, session_state]
        feature_pick_outputs = [
            app_mode_state,
            search_mode_state,
            selected_docs_state,
            *panel_chrome_outputs,
            mode_title,
            mode_badge_btn,
            footer_md,
            chat["export_file"],
            feature_popup,
            popup_open,
            msg,
            send_btn,
            docs["doc_select"],
        ] + chrome_outputs

        for btn, choice in [
            (feat_qa_btn, "Finance QA Chatbot"),
            (feat_summary_btn, "Financial Summary"),
            (feat_ocr_btn, "Document OCR"),
            (feat_entities_btn, "Entity Extraction"),
            (feat_docs_btn, "Upload Documents"),
            (feat_hybrid_btn, "Hybrid RAG"),
            (feat_single_btn, "Single Document Search"),
            (feat_export_btn, "Export session"),
        ]:
            _wire_feature_button(
                btn,
                choice,
                on_feature_pick,
                feature_pick_inputs,
                feature_pick_outputs,
            )

        def toggle_search_mode(sm: SearchMode, sel: list[str], app_mode: AppMode):
            new_mode: SearchMode = "single" if sm == "all" else "all"
            new_sel, foot = docs["on_search_mode_change"](new_mode, sel)
            badge = (
                f"🔍 {mode_label(new_mode)}"
                if new_mode == "all"
                else f"📌 {mode_label(new_mode)}"
            )
            doc_count = len(docs["doc_choices"]())
            foot = footer_hint(new_mode, new_sel, doc_count, app_mode)
            title, _, _ = _header_updates(app_mode, new_mode, new_sel)
            return (
                new_mode,
                new_sel,
                foot,
                badge,
                title,
                docs["doc_select_update"](new_sel),
                *_popup_class_updates(app_mode, new_mode),
            )

        mode_badge_btn.click(
            toggle_search_mode,
            inputs=[search_mode_state, selected_docs_state, app_mode_state],
            outputs=[
                search_mode_state,
                selected_docs_state,
                footer_md,
                mode_badge_btn,
                mode_title,
                docs["doc_select"],
            ]
            + popup_btns,
        )

        async def submit(message, history, search_mode, selected_doc_ids, session_id, app_mode):
            if app_mode != "qa" or not message or not str(message).strip():
                yield history, gr.update(value=render_chat_transcript(history or [])), ""
                return
            async for h, transcript in chat["stream_chat"](
                message, history, search_mode, selected_doc_ids, session_id
            ):
                yield h, gr.update(value=transcript), ""

        send_btn.click(
            submit,
            inputs=[
                msg,
                chat_history_state,
                search_mode_state,
                selected_docs_state,
                session_state,
                app_mode_state,
            ],
            outputs=[
                chat_history_state,
                chat["chat_transcript"],
                msg,
            ],
            queue=True,
        )
        msg.submit(
            submit,
            inputs=[
                msg,
                chat_history_state,
                search_mode_state,
                selected_docs_state,
                session_state,
                app_mode_state,
            ],
            outputs=[
                chat_history_state,
                chat["chat_transcript"],
                msg,
            ],
            queue=True,
        )

        new_chat_btn.click(
            chat["create_new_session"],
            inputs=[search_mode_state, selected_docs_state],
            outputs=[
                session_state,
                chat_history_state,
                chat["chat_transcript"],
                footer_md,
                session_dropdown,
            ],
        )

        def back_to_chat(search_mode, selected, session_id):
            return on_feature_pick(
                "Finance QA Chatbot", search_mode, selected, session_id
            )

        for back_btn in (docs["close_btn"], docs["done_btn"]):
            back_btn.click(
                back_to_chat,
                inputs=feature_pick_inputs,
                outputs=feature_pick_outputs,
            )

        docs["file_upload"].change(
            docs["upload_document"],
            inputs=[docs["file_upload"], selected_docs_state, search_mode_state],
            outputs=[
                docs["upload_status"],
                docs["doc_table"],
                docs["doc_select"],
                docs["delete_dropdown"],
                selected_docs_state,
                footer_md,
            ],
            queue=True,
        ).then(
            lambda: (
                gr.update(choices=docs["doc_choices"]()),
                gr.update(choices=docs["doc_choices"]()),
            ),
            outputs=[summary["doc_dropdown"], entities["doc_dropdown"]],
        )

        docs["upload_btn"].click(
            docs["upload_document"],
            inputs=[docs["file_upload"], selected_docs_state, search_mode_state],
            outputs=[
                docs["upload_status"],
                docs["doc_table"],
                docs["doc_select"],
                docs["delete_dropdown"],
                selected_docs_state,
                footer_md,
            ],
            queue=True,
        ).then(
            lambda: (
                gr.update(choices=docs["doc_choices"]()),
                gr.update(choices=docs["doc_choices"]()),
            ),
            outputs=[summary["doc_dropdown"], entities["doc_dropdown"]],
        )

        docs["delete_btn"].click(
            docs["delete_document"],
            inputs=[docs["delete_dropdown"], selected_docs_state, search_mode_state],
            outputs=[
                docs["upload_status"],
                docs["doc_table"],
                docs["doc_select"],
                selected_docs_state,
                footer_md,
            ],
        )

        docs["doc_select"].change(
            docs["on_doc_selection"],
            inputs=[docs["doc_select"], search_mode_state],
            outputs=[selected_docs_state, footer_md, docs["doc_select"]],
        )

        def startup():
            sid, session_dd, history, transcript = chat["on_load_sessions"]()
            rows = docs["doc_table_rows"]()
            choices = docs["doc_choices"]()
            choice_ids = [c[1] for c in choices]
            return (
                sid,
                session_dd,
                history,
                gr.update(value=transcript),
                rows,
                docs["doc_select_update"]([]),
                gr.update(choices=choice_ids),
                gr.update(choices=choices),
                gr.update(choices=choices),
            )

        demo.load(
            startup,
            outputs=[
                session_state,
                session_dropdown,
                chat_history_state,
                chat["chat_transcript"],
                docs["doc_table"],
                docs["doc_select"],
                docs["delete_dropdown"],
                summary["doc_dropdown"],
                entities["doc_dropdown"],
            ],
        )

    return demo, theme, css


def _resolve_server_port(preferred: int = 7860, attempts: int = 10) -> int:
    """Use preferred port, or the next free port if something is already listening."""
    import os
    import socket

    if os.environ.get("GRADIO_SERVER_PORT"):
        return int(os.environ["GRADIO_SERVER_PORT"])

    for port in range(preferred, preferred + attempts):
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
            sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
            try:
                sock.bind(("127.0.0.1", port))
                if port != preferred:
                    logger.warning(
                        "Port %s is in use (likely an old FinSight instance). "
                        "Open http://127.0.0.1:%s — or stop the process on %s and restart.",
                        preferred,
                        port,
                        preferred,
                    )
                return port
            except OSError:
                continue

    raise OSError(
        f"Cannot find an open port in range {preferred}-{preferred + attempts - 1}. "
        "Stop the existing Gradio process or set GRADIO_SERVER_PORT."
    )


def create_demo() -> tuple[gr.Blocks, gr.Theme, str]:
    logging.basicConfig(level=logging.INFO)
    ensure_services()
    demo, theme, css = build_app()
    demo.queue(default_concurrency_limit=4)
    return demo, theme, css


def launch_demo(demo: gr.Blocks, theme: gr.Theme, css: str) -> None:
    import os

    port_env = os.environ.get("PORT") or os.environ.get("GRADIO_SERVER_PORT")
    port = int(port_env) if port_env else _resolve_server_port()
    demo.launch(
        server_name="0.0.0.0",
        server_port=port,
        theme=theme,
        css=css,
        show_error=True,
        pwa=True,
        ssr_mode=False,
    )


def main() -> None:
    demo, theme, css = create_demo()
    launch_demo(demo, theme, css)


if __name__ == "__main__":
    main()