compliance-scanner-agent/compliance-dashboard/src/pages/chat.rs

use dioxus::prelude::*;

use crate::components::page_header::PageHeader;
use crate::infrastructure::chat::{
    fetch_embedding_status, send_chat_message, trigger_embedding_build, ChatHistoryMessage,
    SourceRef,
};

/// A UI-level chat message
#[derive(Clone, Debug)]
struct UiChatMessage {
    role: String,
    content: String,
    sources: Vec<SourceRef>,
}

#[component]
pub fn ChatPage(repo_id: String) -> Element {
    let mut messages: Signal<Vec<UiChatMessage>> = use_signal(Vec::new);
    let mut input_text = use_signal(String::new);
    let mut loading = use_signal(|| false);
    let mut building = use_signal(|| false);

    let repo_id_for_status = repo_id.clone();
    let mut embedding_status = use_resource(move || {
        let rid = repo_id_for_status.clone();
        async move { fetch_embedding_status(rid).await.ok() }
    });

    let has_embeddings = {
        let status = embedding_status.read();
        match &*status {
            Some(Some(resp)) => resp
                .data
                .as_ref()
                .map(|d| d.status == "completed")
                .unwrap_or(false),
            _ => false,
        }
    };

    let embedding_status_text = {
        let status = embedding_status.read();
        match &*status {
            Some(Some(resp)) => match &resp.data {
                Some(d) => match d.status.as_str() {
                    "completed" => format!(
                        "Embeddings ready: {}/{} chunks",
                        d.embedded_chunks, d.total_chunks
                    ),
                    "running" => format!(
                        "Building embeddings: {}/{}...",
                        d.embedded_chunks, d.total_chunks
                    ),
                    "failed" => format!(
                        "Embedding build failed: {}",
                        d.error_message.as_deref().unwrap_or("unknown error")
                    ),
                    s => format!("Status: {s}"),
                },
                None => "No embeddings built yet".to_string(),
            },
            Some(None) => "Failed to check embedding status".to_string(),
            None => "Checking embedding status...".to_string(),
        }
    };

    let repo_id_for_build = repo_id.clone();
    let on_build = move |_| {
        let rid = repo_id_for_build.clone();
        building.set(true);
        spawn(async move {
            let _ = trigger_embedding_build(rid).await;
            building.set(false);
            embedding_status.restart();
        });
    };

    let repo_id_for_send = repo_id.clone();
    let mut do_send = move || {
        let text = input_text.read().trim().to_string();
        if text.is_empty() || *loading.read() {
            return;
        }

        let rid = repo_id_for_send.clone();
        let user_msg = text.clone();

        // Add user message to UI
        messages.write().push(UiChatMessage {
            role: "user".to_string(),
            content: user_msg.clone(),
            sources: Vec::new(),
        });
        input_text.set(String::new());
        loading.set(true);

        spawn(async move {
            // Build history from existing messages
            let history: Vec<ChatHistoryMessage> = messages
                .read()
                .iter()
                .filter(|m| m.role == "user" || m.role == "assistant")
                .rev()
                .skip(1) // skip the message we just added
                .take(10) // limit history
                .collect::<Vec<_>>()
                .into_iter()
                .rev()
                .map(|m| ChatHistoryMessage {
                    role: m.role.clone(),
                    content: m.content.clone(),
                })
                .collect();

            match send_chat_message(rid, user_msg, history).await {
                Ok(resp) => {
                    messages.write().push(UiChatMessage {
                        role: "assistant".to_string(),
                        content: resp.data.message,
                        sources: resp.data.sources,
                    });
                }
                Err(e) => {
                    messages.write().push(UiChatMessage {
                        role: "assistant".to_string(),
                        content: format!("Error: {e}"),
                        sources: Vec::new(),
                    });
                }
            }
            loading.set(false);
        });
    };

    let mut do_send_click = do_send.clone();

    rsx! {
        PageHeader { title: "AI Chat" }

        // Embedding status banner
        div { class: "chat-embedding-banner",
            span { "{embedding_status_text}" }
            button {
                class: "btn btn-sm",
                disabled: *building.read(),
                onclick: on_build,
                if *building.read() { "Building..." } else { "Build Embeddings" }
            }
        }

        div { class: "chat-container",
            // Message list
            div { class: "chat-messages",
                if messages.read().is_empty() && !*loading.read() {
                    div { class: "chat-empty",
                        h3 { "Ask anything about your codebase" }
                        p { "Build embeddings first, then ask questions about functions, architecture, patterns, and more." }
                    }
                }
                for (i, msg) in messages.read().iter().enumerate() {
                    {
                        let class = if msg.role == "user" {
                            "chat-message chat-message-user"
                        } else {
                            "chat-message chat-message-assistant"
                        };
                        let content = msg.content.clone();
                        let sources = msg.sources.clone();
                        rsx! {
                            div { class: class, key: "{i}",
                                div { class: "chat-message-role",
                                    if msg.role == "user" { "You" } else { "Assistant" }
                                }
                                div { class: "chat-message-content", "{content}" }
                                if !sources.is_empty() {
                                    div { class: "chat-sources",
                                        span { class: "chat-sources-label", "Sources:" }
                                        for src in sources {
                                            div { class: "chat-source-card",
                                                div { class: "chat-source-header",
                                                    span { class: "chat-source-name",
                                                        "{src.qualified_name}"
                                                    }
                                                    span { class: "chat-source-location",
                                                        "{src.file_path}:{src.start_line}-{src.end_line}"
                                                    }
                                                }
                                                pre { class: "chat-source-snippet",
                                                    code { "{src.snippet}" }
                                                }
                                            }
                                        }
                                    }
                                }
                            }
                        }
                    }
                }
                if *loading.read() {
                    div { class: "chat-message chat-message-assistant",
                        div { class: "chat-message-role", "Assistant" }
                        div { class: "chat-message-content chat-typing", "Thinking..." }
                    }
                }
            }

            // Input area
            div { class: "chat-input-area",
                textarea {
                    class: "chat-input",
                    placeholder: "Ask about your codebase...",
                    value: "{input_text}",
                    disabled: !has_embeddings,
                    oninput: move |e| input_text.set(e.value()),
                    onkeydown: move |e: Event<KeyboardData>| {
                        if e.key() == Key::Enter && !e.modifiers().shift() {
                            e.prevent_default();
                            do_send();
                        }
                    },
                }
                button {
                    class: "btn chat-send-btn",
                    disabled: *loading.read() || !has_embeddings,
                    onclick: move |_| do_send_click(),
                    "Send"
                }
            }
        }
    }
}