ui_chat.py

import html
import gradio as gr

from logic import (
    send_message_public,
    maybe_generate_image,
    initial_greeting,
    DEFAULT_AVATAR_ID,
)

# -------------------------------------------------------------
# CSS + JS (modern UI, responsive, smooth scroll)
# -------------------------------------------------------------
UI_STYLE = """
<style>
:root {
    --user-bg: #e8f1ff;
    --user-border: #8cb4ff;

    --avatar-bg: #dfffe0;
    --avatar-border: #8cd190;

    --panel-border: #d4d4d4;
    --panel-radius: 16px;
    color-scheme: light;
}

/* Chat panel container */
#chat-panel {
    width: 100%;
}

/* Scrollable chat area */
.chat-scroll {
    max-height: 800px;
    min-height: 340px;
    overflow-y: auto;
    padding: 20px;
    background: white;
    border: 3px solid var(--panel-border);
    border-radius: var(--panel-radius);
    display: flex;
    flex-direction: column;
    gap: 16px;
    scroll-behavior: smooth;
}

.chat-empty {
    text-align: center;
    color: #6b7280;
    font-style: italic;
}

/* Chat row alignment */
.chat-row {
    display: flex;
    width: 100%;
}

.message-row.user {
    justify-content: flex-end;
}
.message-row.avatar {
    justify-content: flex-start;
}

/* Chat bubble */
.chat-bubble {
    width: 75%;
    padding: 12px 14px;
    border-radius: 14px;
    border: 2px solid transparent;
    display: flex;
    flex-direction: column;
    gap: 8px;
    box-shadow: 0 2px 6px rgba(0,0,0,0.05);
}

.chat-bubble.user {
    background: var(--user-bg);
    border-color: var(--user-border);
    margin-left: auto;
}
.chat-bubble.avatar {
    background: var(--avatar-bg);
    border-color: var(--avatar-border);
    margin-right: auto;
}

/* Inline images */
.chat-bubble img {
    width: 70%;
    border-radius: 12px;
    border: 1px solid #d0d7e2;
}

.avatar-chat-text {
    font-size: 1rem;
    font-weight: 600;
    color: #0f172a !important;
    line-height: 1.45;
}

/* Mobile layout */
@media (max-width: 768px) {
    .chat-bubble {
        width: 100%;
    }
    .chat-bubble img {
        width: 100%;
    }
}

/* Dark theme overrides for device-level dark mode */
@media (prefers-color-scheme: dark) {
    :root {
        --user-bg: #1f2937;
        --user-border: #60a5fa;

        --avatar-bg: #0f172a;
        --avatar-border: #34d399;

        --panel-border: #1f2937;
    }

    body,
    .gradio-container {
        background: #020617 !important;
        color: #f1f5f9 !important;
    }

    .chat-scroll {
        background: #0b1120;
        border-color: var(--panel-border);
    }

    .chat-empty {
        color: #94a3b8;
    }

    .avatar-chat-text {
        color: #f8fafc !important;
    }
}

/* Keep chat readable while backend events run */
.gradio-container .block.loading,
.gradio-container .block.pending,
.gradio-container .loading,
.gradio-container .pending {
    filter: none !important;
    opacity: 1 !important;
}
.gradio-container .block.loading *,
.gradio-container .block.pending *,
.gradio-container .loading *,
.gradio-container .pending * {
    filter: none !important;
    opacity: 1 !important;
    color: inherit !important;
}
.gradio-container .block.loading::after,
.gradio-container .block.loading::before,
.gradio-container .block.pending::after,
.gradio-container .block.pending::before,
.gradio-container .loading::after,
.gradio-container .loading::before,
.gradio-container .pending::after,
.gradio-container .pending::before {
    display: none !important;
}
</style>
"""

SCROLL_JS_BODY = """
const app = document.querySelector("gradio-app");
const root = app && app.shadowRoot ? app.shadowRoot : document;
const host = root.querySelector("#chat-panel");
if (!host) {
    console.warn("[scroll-js] host not found");
} else {
    const panel = host.querySelector(".chat-scroll") || host.querySelector(".svelte-1n1h5do") || host;
    if (!panel) {
        console.warn("[scroll-js] scrollable panel not found");
    } else {
        const run = () => {
            panel.scrollTop = panel.scrollHeight;
        };
        run();
        setTimeout(run, 32);
        setTimeout(run, 160);
        panel.querySelectorAll("img").forEach((img) => {
            if (img.dataset.scrollWatcher === "1") return;
            img.dataset.scrollWatcher = "1";
            const fire = () => {
                panel.scrollTop = panel.scrollHeight;
            };
            if (img.complete) {
                setTimeout(fire, 32);
            } else {
                img.addEventListener("load", fire, { once: true });
                img.addEventListener("error", fire, { once: true });
            }
        });
        if (!panel.dataset.observeScroll) {
            panel.dataset.observeScroll = "1";
            const observer = new MutationObserver(() => {
                panel.scrollTop = panel.scrollHeight;
            });
            observer.observe(panel, { childList: true, subtree: true });
        }
    }
}
"""

SCROLL_SNIPPET = f"<script>(function(){{{SCROLL_JS_BODY}}})()</script>"
SCROLL_BUTTON_JS = f"() => {{{SCROLL_JS_BODY}}}"


def render_chat(history):
    rows = ["<div class='chat-scroll'>"]

    if not history:
        rows.append("<div class='chat-empty'>Start the conversation...</div>")
    else:
        for entry in history:
            speaker = entry.get("speaker", "avatar")
            text = entry.get("text", "") or ""
            safe_text = html.escape(text) if isinstance(text, str) else str(text)

            row_class = "chat-row message-row user" if speaker == "user" else "chat-row message-row avatar"
            bubble_class = "chat-bubble user" if speaker == "user" else "chat-bubble avatar"

            html_block = [f"<div class='{row_class}'>"]
            html_block.append(f"<div class='{bubble_class}'>")

            img = entry.get("image_data")
            if img and speaker != "user":
                html_block.append(f"<img src='{img}' />")

            html_block.append(f"<div class='avatar-chat-text'>{safe_text}</div>")

            if img and speaker == "user":
                html_block.append(f"<img src='{img}' />")

            html_block.append("</div></div>")
            rows.append("".join(html_block))

    rows.append("</div>")
    rows.append(SCROLL_SNIPPET)
    return "\n".join(rows)


def build_chat_tab(blocks: gr.Blocks):
    with gr.Tab("Chat"):
        avatar_id_input = gr.Textbox(
            label="Avatar ID",
            value=DEFAULT_AVATAR_ID,
            interactive=True,
            scale=2
        )

        chat_display = gr.HTML(render_chat([]), elem_id="chat-panel", sanitize_html=False)
        msg_input = gr.Textbox(label="Message", placeholder="Say something...")
        state = gr.State([])
        pending_message = gr.State("")

        with gr.Row():
            with gr.Column(scale=1):
                send_btn = gr.Button("Send", variant="primary")
                tool_status = gr.Markdown("Tool used: n/a")
            generate_toggle = gr.Checkbox(label="Generate Images", value=True)
            clear_btn = gr.Button("Clear")

        def stash_message(msg):
            return msg, ""

        def mark_processing():
            return gr.update(value="Generating...", interactive=False)

        def mark_ready():
            return gr.update(value="Send", interactive=True)

        def on_send(aid, msg, history, gen_flag):
            history, tool = send_message_public(aid, msg, history, gen_flag)
            return render_chat(history), "", history, f"Tool used: {tool}", gen_flag

        def preview_user_message(msg, history):
            history = list(history or [])
            msg = (msg or "").strip()
            if not msg:
                return render_chat(history)
            history.append({"speaker": "user", "text": msg})
            return render_chat(history)

        def on_image_gen(aid, history, gen_flag):
            history = maybe_generate_image(aid, history, gen_flag)
            return render_chat(history), history, gen_flag

        def mark_image_processing(gen_flag):
            if not gen_flag:
                return gr.update()
            return gr.update(label="Generate Images (processing...)", interactive=False)

        def mark_image_ready(gen_flag):
            if not gen_flag:
                return gr.update()
            return gr.update(label="Generate Images", interactive=True)

        send_btn.click(
            stash_message,
            inputs=[msg_input],
            outputs=[pending_message, msg_input],
            queue=False
        ).then(
            preview_user_message,
            inputs=[pending_message, state],
            outputs=[chat_display],
            queue=False,
            show_progress="hidden"
        ).then(
            mark_processing,
            outputs=[send_btn],
            show_progress="hidden"
        ).then(
            on_send,
            inputs=[avatar_id_input, pending_message, state, generate_toggle],
            outputs=[chat_display, msg_input, state, tool_status, generate_toggle],
            show_progress="hidden"
        ).then(
            mark_image_processing,
            inputs=[generate_toggle],
            outputs=[generate_toggle],
            show_progress="hidden"
        ).then(
            on_image_gen,
            inputs=[avatar_id_input, state, generate_toggle],
            outputs=[chat_display, state, generate_toggle],
            show_progress="hidden"
        ).then(
            mark_image_ready,
            inputs=[generate_toggle],
            outputs=[generate_toggle],
            show_progress="hidden"
        ).then(
            mark_ready,
            outputs=[send_btn],
            show_progress="hidden"
        ).then(
            None,
            js=SCROLL_BUTTON_JS,
            show_progress="hidden"
        )


        msg_input.submit(
            stash_message,
            inputs=[msg_input],
            outputs=[pending_message, msg_input],
            queue=False
        ).then(
            preview_user_message,
            inputs=[pending_message, state],
            outputs=[chat_display],
            queue=False,
            show_progress="hidden"
        ).then(
            mark_processing,
            outputs=[send_btn],
            show_progress="hidden"
        ).then(
            on_send,
            inputs=[avatar_id_input, pending_message, state, generate_toggle],
            outputs=[chat_display, msg_input, state, tool_status, generate_toggle],
            show_progress="hidden"
        ).then(
            mark_image_processing,
            inputs=[generate_toggle],
            outputs=[generate_toggle],
            show_progress="hidden"
        ).then(
            on_image_gen,
            inputs=[avatar_id_input, state, generate_toggle],
            outputs=[chat_display, state, generate_toggle],
            show_progress="hidden"
        ).then(
            mark_image_ready,
            inputs=[generate_toggle],
            outputs=[generate_toggle],
            show_progress="hidden"
        ).then(
            mark_ready,
            outputs=[send_btn],
            show_progress="hidden"
        ).then(
            None,
            js=SCROLL_BUTTON_JS,
            show_progress="hidden"
        )

        clear_btn.click(
            lambda: (render_chat([]), "", [], "Tool used: n/a", True),
            outputs=[chat_display, msg_input, state, tool_status, generate_toggle],
            show_progress="hidden"
        ).then(
            None,
            js=SCROLL_BUTTON_JS,
            show_progress="hidden"
        )

        def load_greeting(aid):
            hist = initial_greeting(aid)
            return render_chat(hist), "", hist, "Tool used: greet", True

        blocks.load(
            load_greeting,
            inputs=[avatar_id_input],
            outputs=[chat_display, msg_input, state, tool_status, generate_toggle],
            show_progress="hidden"
        ).then(
            mark_image_processing,
            inputs=[generate_toggle],
            outputs=[generate_toggle],
            show_progress="hidden"
        ).then(
            on_image_gen,
            inputs=[avatar_id_input, state, generate_toggle],
            outputs=[chat_display, state, generate_toggle],
            show_progress="hidden"
        ).then(
            mark_image_ready,
            inputs=[generate_toggle],
            outputs=[generate_toggle],
            show_progress="hidden"
        ).then(
            None,
            js=SCROLL_BUTTON_JS,
            show_progress="hidden"
        )