ui/layout.py

"""
Gradio Blocks layout for the DSA Tutor Chatbot.

Custom layout only (per spec) — no gr.ChatInterface. Structure:

    Top bar: title, auth status, Login/Logout buttons
    Sign-up dialog (modal-ish accordion, opened via a "Sign up" link)
    Main row:
        [ Toggleable history panel ]  [ Chat window + input + send/new chat ]

Streaming: the assistant's reply is revealed progressively via
ui.callbacks.stream_last_response (see that module's docstring for why
it's simulated rather than a true token stream).
"""

import gradio as gr

from ui import callbacks
from ui.theme import CUSTOM_CSS, dsa_theme


def build_ui(flask_app) -> gr.Blocks:
    callbacks.init_app(flask_app)

    with gr.Blocks(theme=dsa_theme, css=CUSTOM_CSS, title="DSA Tutor Chatbot") as demo:
        auth_state = gr.State(None)          # {"user_id":.., "email":..} | None
        session_state = gr.State(None)       # active ChatSession.id | None
        anon_history = gr.State([])          # rolling history for anonymous users
        history_visible_state = gr.State(False)  # tracks History panel visibility

        # ---------------------------------------------------------------
        # Top bar
        # ---------------------------------------------------------------
        with gr.Row(elem_id="topbar"):
            gr.Markdown("### 🧮 DSA Tutor Chatbot", elem_id="topbar-title")
            with gr.Row(elem_id="topbar-actions"):
                auth_status_md = gr.Markdown("Not logged in", elem_id="auth-status")
                history_toggle_btn = gr.Button("📜 History", size="sm", scale=0, elem_classes=["small-btn"])
                login_btn = gr.Button("Log in", size="sm", scale=0, elem_classes=["small-btn"])
                signup_btn = gr.Button("Sign up", size="sm", scale=0, elem_classes=["small-btn"])
                logout_btn = gr.Button("Log out", size="sm", scale=0, visible=False, elem_classes=["small-btn"])

        # ---------------------------------------------------------------
        # Login popover
        # ---------------------------------------------------------------
        with gr.Group(visible=False) as login_panel:
            gr.Markdown("#### Log in")
            login_email = gr.Textbox(label="Email")
            login_password = gr.Textbox(label="Password", type="password")
            with gr.Row():
                login_submit = gr.Button("Log in", variant="primary", size="sm", elem_classes=["small-btn"])
                login_cancel = gr.Button("Cancel", size="sm", elem_classes=["small-btn"])
            login_status = gr.Markdown("")

        # ---------------------------------------------------------------
        # Sign-up dialog
        # ---------------------------------------------------------------
        with gr.Group(visible=False) as signup_panel:
            gr.Markdown("#### Create an account")
            signup_email = gr.Textbox(label="Email")
            signup_password = gr.Textbox(label="Password (min 8 characters)", type="password")
            with gr.Row():
                signup_submit = gr.Button("Sign up", variant="primary", size="sm", elem_classes=["small-btn"])
                signup_cancel = gr.Button("Cancel", size="sm", elem_classes=["small-btn"])
            signup_status = gr.Markdown("")

        # ---------------------------------------------------------------
        # Main area: history panel (toggleable) + chat
        # ---------------------------------------------------------------
        with gr.Row(elem_id="main-row"):
            with gr.Column(scale=1, visible=False, elem_id="history-panel") as history_panel:
                gr.Markdown("**Previous sessions**")
                session_list = gr.Radio(choices=[], label="", elem_classes=["session-item"])
                with gr.Row():
                    refresh_sessions_btn = gr.Button("Refresh", size="sm", elem_classes=["small-btn"])
                    delete_session_btn = gr.Button("Delete selected", size="sm", elem_classes=["small-btn"])

            with gr.Column(scale=4, elem_id="chat-column"):
                chatbot = gr.Chatbot(
                    label="DSA Tutor",
                    type="messages",
                    height=360,
                    elem_id="chat-window",
                    avatar_images=(None, None),
                )
                with gr.Row(elem_id="input-row"):
                    msg_box = gr.Textbox(
                        placeholder="Ask about any DSA topic (e.g. 'Explain merge sort')...",
                        show_label=False,
                        scale=8,
                    )
                    send_btn = gr.Button("Send", variant="primary", scale=1, size="sm")
                new_chat_btn = gr.Button("🆕 New chat", size="sm", elem_classes=["small-btn"])

        # ---------------------------------------------------------------
        # Wiring: top bar toggles
        # ---------------------------------------------------------------
        def _toggle_history(is_visible):
            new_visible = not is_visible
            return new_visible, gr.update(visible=new_visible)

        history_toggle_btn.click(
            _toggle_history,
            inputs=[history_visible_state],
            outputs=[history_visible_state, history_panel],
        )

        login_btn.click(lambda: gr.update(visible=True), outputs=login_panel)
        login_cancel.click(lambda: gr.update(visible=False), outputs=login_panel)

        signup_btn.click(lambda: gr.update(visible=True), outputs=signup_panel)
        signup_cancel.click(lambda: gr.update(visible=False), outputs=signup_panel)

        def _auth_ui_update(auth_state_value):
            if auth_state_value:
                status = f"Logged in as **{auth_state_value['email']}**"
                return (
                    status,
                    gr.update(visible=False),  # login_btn
                    gr.update(visible=False),  # signup_btn
                    gr.update(visible=True),   # logout_btn
                )
            return (
                "Not logged in",
                gr.update(visible=True),
                gr.update(visible=True),
                gr.update(visible=False),
            )

        def _on_login(email, password):
            auth_state_value, message, is_error = callbacks.do_login(email, password)
            status_md, login_show, signup_show, logout_show = _auth_ui_update(auth_state_value)
            panel_visible = gr.update(visible=is_error)
            return auth_state_value, status_md, login_show, signup_show, logout_show, panel_visible, message

        login_submit.click(
            _on_login,
            inputs=[login_email, login_password],
            outputs=[auth_state, auth_status_md, login_btn, signup_btn, logout_btn, login_panel, login_status],
        )

        def _on_signup(email, password):
            auth_state_value, message, is_error = callbacks.do_signup(email, password)
            status_md, login_show, signup_show, logout_show = _auth_ui_update(auth_state_value)
            panel_visible = gr.update(visible=is_error)
            return auth_state_value, status_md, login_show, signup_show, logout_show, panel_visible, message

        signup_submit.click(
            _on_signup,
            inputs=[signup_email, signup_password],
            outputs=[auth_state, auth_status_md, login_btn, signup_btn, logout_btn, signup_panel, signup_status],
        )

        def _on_logout():
            auth_state_value, message, _ = callbacks.do_logout()
            status_md, login_show, signup_show, logout_show = _auth_ui_update(auth_state_value)
            return auth_state_value, status_md, login_show, signup_show, logout_show, [], None, []

        logout_btn.click(
            _on_logout,
            outputs=[
                auth_state, auth_status_md, login_btn, signup_btn, logout_btn,
                chatbot, session_state, anon_history,
            ],
        )

        # ---------------------------------------------------------------
        # Wiring: chat
        # ---------------------------------------------------------------
        def _send_and_stream(user_message, chat_display, auth_state_value, session_state_value, anon_hist):
            chat_display, session_state_value, anon_hist, cleared = callbacks.send_message(
                user_message, chat_display, auth_state_value, session_state_value, anon_hist
            )
            for partial in callbacks.stream_last_response(chat_display):
                yield partial, session_state_value, anon_hist, cleared

        send_btn.click(
            _send_and_stream,
            inputs=[msg_box, chatbot, auth_state, session_state, anon_history],
            outputs=[chatbot, session_state, anon_history, msg_box],
        )
        msg_box.submit(
            _send_and_stream,
            inputs=[msg_box, chatbot, auth_state, session_state, anon_history],
            outputs=[chatbot, session_state, anon_history, msg_box],
        )

        new_chat_btn.click(
            callbacks.new_chat,
            outputs=[chatbot, session_state, anon_history],
        )

        # ---------------------------------------------------------------
        # Wiring: history panel
        # ---------------------------------------------------------------
        refresh_sessions_btn.click(
            callbacks.refresh_sessions,
            inputs=[auth_state],
            outputs=[session_list],
        )

        session_list.change(
            callbacks.open_session,
            inputs=[session_list, auth_state],
            outputs=[chatbot, session_state],
        )

        delete_session_btn.click(
            callbacks.delete_session,
            inputs=[session_list, auth_state],
            outputs=[chatbot, session_state],
        ).then(
            callbacks.refresh_sessions,
            inputs=[auth_state],
            outputs=[session_list],
        )

        # Refresh the session list whenever a login/signup succeeds.
        auth_state.change(
            callbacks.refresh_sessions,
            inputs=[auth_state],
            outputs=[session_list],
        )

    return demo