ui/components.py

"""
DocMind — Reusable Streamlit UI Components

Provides render functions for chat messages, source cards,
grounding bars, confidence badges, and pipeline progress.
"""

import streamlit as st
from typing import Dict, List, Optional

from pipeline.chunker import ChunkMetadata
from pipeline.grounding import ConfidenceLevel, GroundingResult, SentenceScore
from pipeline.retriever import RetrievalStats


# ── Chat Messages ───────────────────────────────────────────────────

def render_chat_message(
    role: str,
    content: str,
    grounding_result: Optional[GroundingResult] = None,
    sources: Optional[List[ChunkMetadata]] = None,
    doc_index_map: Optional[Dict[str, int]] = None,
) -> None:
    """
    Render a single chat message with optional grounding info and sources.

    Args:
        role: "user" or "bot"
        content: The message text
        grounding_result: Optional grounding gate result for bot messages
        sources: Optional list of source chunks for bot messages
        doc_index_map: Optional mapping of doc_id → color index
    """
    if role == "user":
        avatar = "👤"
        bubble_class = "user"
    else:
        avatar = "🧠"
        bubble_class = "bot"

    st.markdown(f"""
    <div class="chat-msg {role}">
        <div class="chat-avatar {role}">{avatar}</div>
        <div class="chat-bubble {bubble_class}">{content}</div>
    </div>
    """, unsafe_allow_html=True)

    # For bot messages, show grounding info and sources
    if role == "bot" and grounding_result and not grounding_result.is_refused:
        render_grounding_bar(grounding_result)
        render_confidence_badge(grounding_result.confidence)

        # Show per-sentence scores in an expander
        if grounding_result.sentence_scores:
            with st.expander("📊 Per-sentence grounding scores"):
                for ss in grounding_result.sentence_scores:
                    _render_sentence_score(ss)

    if role == "bot" and sources:
        with st.expander(f"📎 Source chunks ({len(sources)})"):
            for chunk in sources:
                doc_idx = 0
                if doc_index_map and chunk.doc_id in doc_index_map:
                    doc_idx = doc_index_map[chunk.doc_id]
                render_source_card(chunk, doc_idx)


# ── Source Cards ────────────────────────────────────────────────────

def render_source_card(chunk: ChunkMetadata, doc_color_index: int = 0) -> None:
    """Render a collapsible source chunk preview with document color tag."""
    tag_class = f"doc-tag-{doc_color_index % 3}"
    preview_text = chunk.text[:300] + ("..." if len(chunk.text) > 300 else "")

    st.markdown(f"""
    <div class="source-card">
        <div class="source-card-header">
            <span class="source-tag {tag_class}">{chunk.doc_name}</span>
            <span>📄 Page {chunk.page_num}</span>
            <span style="color: #64748B;">|</span>
            <span style="color: #64748B;">{chunk.chunk_id}</span>
        </div>
        <div>{preview_text}</div>
    </div>
    """, unsafe_allow_html=True)


# ── Grounding Bar ──────────────────────────────────────────────────

def render_grounding_bar(grounding_result: GroundingResult) -> None:
    """Render an animated grounding score progress bar."""
    score = grounding_result.overall_score
    pct = max(0, min(100, int(score * 100)))

    if grounding_result.confidence == ConfidenceLevel.HIGH:
        fill_class = "high"
    elif grounding_result.confidence == ConfidenceLevel.MODERATE:
        fill_class = "moderate"
    else:
        fill_class = "low"

    st.markdown(f"""
    <div class="grounding-bar-container">
        <div class="grounding-bar-label">
            <span>Grounding Score</span>
            <span>{score:.1%}</span>
        </div>
        <div class="grounding-bar-track">
            <div class="grounding-bar-fill {fill_class}" style="width: {pct}%;"></div>
        </div>
    </div>
    """, unsafe_allow_html=True)


# ── Confidence Badge ───────────────────────────────────────────────

def render_confidence_badge(level: ConfidenceLevel) -> None:
    """Render a confidence level badge."""
    badges = {
        ConfidenceLevel.HIGH: ("✅ High Confidence", "badge-high"),
        ConfidenceLevel.MODERATE: ("⚠️ Moderate Confidence", "badge-moderate"),
        ConfidenceLevel.LOW: ("❌ Insufficient Grounding", "badge-low"),
    }
    text, css_class = badges.get(level, ("❓ Unknown", "badge-low"))

    st.markdown(f"""
    <div class="badge {css_class}">{text}</div>
    """, unsafe_allow_html=True)


def _render_sentence_score(ss: SentenceScore) -> None:
    """Render a single sentence's grounding score."""
    if ss.confidence == ConfidenceLevel.HIGH:
        color = "#34D399"
        icon = "✅"
    elif ss.confidence == ConfidenceLevel.MODERATE:
        color = "#FBBF24"
        icon = "⚠️"
    else:
        color = "#F87171"
        icon = "❌"

    st.markdown(f"""
    <div style="padding: 0.3rem 0; border-bottom: 1px solid rgba(148,163,184,0.08);">
        <span>{icon}</span>
        <span style="color: {color}; font-weight: 600; font-size: 0.8rem;">
            {ss.entailment_score:.1%}
        </span>
        <span style="color: #94A3B8; font-size: 0.78rem; margin-left: 0.3rem;">
            [{ss.chunk_id}]
        </span>
        <br>
        <span style="color: #CBD5E1; font-size: 0.82rem;">{ss.sentence}</span>
    </div>
    """, unsafe_allow_html=True)


# ── Retrieval Stats ────────────────────────────────────────────────

def render_retrieval_stats(stats: RetrievalStats) -> None:
    """Render retrieval debug information."""
    st.markdown(f"""
    <div class="debug-panel">
        <strong>Retrieval Stats</strong><br>
        BM25: {stats.bm25_hits} hits &nbsp;|&nbsp;
        Dense: {stats.dense_hits} hits &nbsp;|&nbsp;
        After RRF: {stats.rrf_results} &nbsp;|&nbsp;
        Latency: {stats.latency_ms:.0f}ms
    </div>
    """, unsafe_allow_html=True)


# ── Document Status ────────────────────────────────────────────────

def render_document_status(
    doc_name: str,
    chunk_count: int,
    page_count: int,
    doc_color_index: int = 0,
) -> None:
    """Render a sidebar document status card."""
    tag_class = f"doc-tag-{doc_color_index % 3}"

    st.markdown(f"""
    <div class="doc-status">
        <div class="doc-status-name">
            <span class="source-tag {tag_class}">&nbsp;</span>
            📄 {doc_name}
        </div>
        <div class="doc-status-meta">
            {chunk_count} chunks • {page_count} pages
        </div>
    </div>
    """, unsafe_allow_html=True)


# ── Pipeline Progress ──────────────────────────────────────────────

def render_pipeline_progress(stages: List[dict]) -> None:
    """
    Render a multi-stage pipeline progress indicator.

    Each stage is a dict: {"name": str, "status": "done"|"active"|"pending"}
    """
    icons = {"done": "✅", "active": "⏳", "pending": "⏸️"}
    html_parts = []
    for stage in stages:
        icon = icons.get(stage["status"], "⏸️")
        css_class = stage["status"]
        html_parts.append(
            f'<div class="pipeline-step {css_class}">{icon} {stage["name"]}</div>'
        )

    st.markdown("\n".join(html_parts), unsafe_allow_html=True)


# ── Empty State / Hero ─────────────────────────────────────────────

def render_empty_state() -> None:
    """Render the advanced empty state (Hero Section) when no documents are uploaded."""
    st.html("""
    <div class="hero-container">
        <h2 style="font-size: 2.2rem; font-weight: 700; color: #E2E8F0; margin-bottom: 0.5rem;">Welcome to DocMind</h2>
        <p style="color: #94A3B8; font-size: 1.1rem; max-width: 600px; margin: 0 auto;">
            Upload your documents in the sidebar to securely chat, summarize, and compare content using Grounded RAG.
        </p>
        
        <div class="feature-grid">
            <div class="feature-card">
                <div class="feature-icon">⚡</div>
                <div class="feature-title">Hybrid Retrieval</div>
                <div class="feature-desc">
                    Combines dense vector search (BGE-M3) with sparse keyword search (BM25) to ensure maximum context recall across your documents.
                </div>
            </div>
            
            <div class="feature-card">
                <div class="feature-icon">🛡️</div>
                <div class="feature-title">NLI Grounding</div>
                <div class="feature-desc">
                    Every generated sentence is rigorously verified against the source context using Natural Language Inference models to eliminate hallucinations.
                </div>
            </div>
            
            <div class="feature-card">
                <div class="feature-icon">🧠</div>
                <div class="feature-title">Llama 3 Powered</div>
                <div class="feature-desc">
                    Utilizes Groq's lightning-fast inference for deep reasoning, automated summarization, and side-by-side document comparison.
                </div>
            </div>
        </div>
    </div>
    """)


# ── Comparison Table ───────────────────────────────────────────────

def render_comparison_table(comparison_text: str) -> None:
    """Render a document comparison result."""
    st.markdown(comparison_text, unsafe_allow_html=True)


# ── Dashboard Metrics ──────────────────────────────────────────────

def render_dashboard_metrics(doc_count: int, chunk_count: int, memory_mb: float = 0.0) -> None:
    """Render a row of top-level metric cards for the dashboard."""
    st.markdown(f"""
    <div class="metric-row">
        <div class="metric-card">
            <div class="metric-title">Documents</div>
            <div class="metric-value">{doc_count}</div>
            <div class="metric-subtitle">Currently Indexed</div>
        </div>
        <div class="metric-card">
            <div class="metric-title">Knowledge Chunks</div>
            <div class="metric-value">{chunk_count}</div>
            <div class="metric-subtitle">Vector DB + BM25</div>
        </div>
        <div class="metric-card">
            <div class="metric-title">System Status</div>
            <div class="metric-value">Ready</div>
            <div class="metric-subtitle" style="color: #34D399;">Hybrid Retrieval Online</div>
        </div>
    </div>
    """, unsafe_allow_html=True)


# ── Key Point Card ─────────────────────────────────────────────────

def render_keypoint_card(text: str, page_ref: str = "") -> None:
    """Render a styled HTML card for a key point."""
    page_html = f" <span style='color:#64748B; font-size:0.8rem;'>{page_ref}</span>" if page_ref else ""
    st.markdown(f"""
    <div class="keypoint-card">
        <div class="keypoint-icon">🎯</div>
        <div>
            {text}{page_html}
        </div>
    </div>
    """, unsafe_allow_html=True)