ui.py

"""
ui.py — Reusable Streamlit UI rendering helpers.

Keeps all layout/widget code that isn't tied to a single page out of app.py,
while keeping utils.py pure-Python and state.py focused on callbacks.
"""

import os

import streamlit as st
import plotly.express as px

import utils
from state import (
    updateCategoryOptions, removeCategory, addCategory,
    applyGlobalRenames, addGlobalRename, removeGlobalRename,
    _global_rename_key, randomize_speaker_clip, apply_inline_rename,
)


# ---------------------------------------------------------------------------
# Chart rendering
# ---------------------------------------------------------------------------

def render_chart(fig, tab, pdf_path, svg_path, pdf_name, svg_name, pdf_key, svg_key,
                 plotly_config=None):
    """Render a Plotly figure inside a tab with PDF/SVG download buttons."""
    cfg = plotly_config or {"displayModeBar": True, "modeBarButtonsToRemove": []}
    with tab:
        st.plotly_chart(fig, use_container_width=True, config=cfg)
        col_l, col_r = st.columns(2)
        try:
            fig.write_image(pdf_path)
            fig.write_image(svg_path)
        except Exception:
            pass
        with col_l:
            if os.path.exists(pdf_path):
                with open(pdf_path, "rb") as f:
                    st.download_button("Save As PDF", f, pdf_name, "application/pdf",
                                       key=pdf_key, on_click="ignore")
        with col_r:
            if os.path.exists(svg_path):
                with open(svg_path, "rb") as f:
                    st.download_button("Save As SVG", f, svg_name, "image/svg+xml",
                                       key=svg_key, on_click="ignore")


# ---------------------------------------------------------------------------
# Sidebar — categories section
# ---------------------------------------------------------------------------

def render_categories_sidebar(currFile, categorySelections, all_speakers_display,
                               raw_to_display):
    """Render category multiselects, remove buttons, and Add category input."""
    for i, category in enumerate(st.session_state.categories):
        ms_key       = f"multiselect_{category}"
        speakerSet   = categorySelections[i]
        default_disp = [raw_to_display.get(sp, sp) for sp in speakerSet]
        if ms_key not in st.session_state:
            st.session_state[ms_key] = default_disp
        st.sidebar.multiselect(
            category, all_speakers_display,
            key=ms_key, on_change=updateCategoryOptions, args=(currFile,),
        )
        st.sidebar.button(
            f"Remove {category}", key=f"remove_{category}",
            on_click=removeCategory, args=(i,),
        )
    st.sidebar.text_input("Add Category", key="categoryInput", on_change=addCategory)


# ---------------------------------------------------------------------------
# Sidebar — rename speakers section
# ---------------------------------------------------------------------------

def render_rename_sidebar(currFile, speakerNames, all_speaker_tokens):
    """Render the full Rename Speakers sidebar section."""
    st.sidebar.divider()
    st.sidebar.subheader("Rename Speakers")
    st.sidebar.markdown(
        "<p style='font-size:0.85rem; color:gray; margin-bottom:2px;'>"
        "Assign a name and select which speaker labels (across all files) it applies to. "
        "Changes apply to all matched speakers instantly.</p>",
        unsafe_allow_html=True,
    )

    st.sidebar.divider()

    def _on_grename_change(idx):
        st.session_state.globalRenames[idx]["speakers"] = list(
            st.session_state[_global_rename_key(idx)]
        )
        applyGlobalRenames()

    for idx, entry in enumerate(st.session_state.globalRenames):
        grkey = _global_rename_key(idx)
        if grkey not in st.session_state:
            st.session_state[grkey] = list(entry["speakers"])
        st.sidebar.markdown(f"**{entry['name']}**")
        st.sidebar.multiselect(
            f"Speakers for {entry['name']}", options=all_speaker_tokens,
            key=grkey, on_change=_on_grename_change, args=(idx,),
            label_visibility="collapsed",
        )
        st.sidebar.button(
            f"Remove '{entry['name']}'", key=f"remove_grename_{idx}",
            on_click=removeGlobalRename, args=(idx,),
        )

    st.sidebar.text_input(
        "Add rename", placeholder="e.g. John",
        key="globalRenameInput", on_change=addGlobalRename,
    )


# ---------------------------------------------------------------------------
# Data tab — custom scrollable table
# ---------------------------------------------------------------------------

def render_data_table(tableDF, speakerNames, raw_to_display, currFile):
    """Render the custom column-based scrollable table (Speaker, Start, Finish)."""
    other_cols = [c for c in tableDF.columns if c != "Speaker"]
    col_widths = [2] + [1] * len(other_cols)

    # Header
    header_cols = st.columns(col_widths)
    header_cols[0].markdown("**Speaker**")
    for i, col_name in enumerate(other_cols):
        header_cols[i + 1].markdown(f"**{col_name}**")

    st.markdown("<hr style='margin-top:2px; margin-bottom:4px;'>", unsafe_allow_html=True)

    with st.container(height=480, border=False):
        for _, row in tableDF.iterrows():
            row_cols   = st.columns(col_widths)
            display_sp = row["Speaker"]
            row_cols[0].write(display_sp)
            for i, col_name in enumerate(other_cols):
                row_cols[i + 1].write(row[col_name])


# ---------------------------------------------------------------------------
# Rename Speaker tab — speaker / audio sample table
# ---------------------------------------------------------------------------

def render_speaker_samples_tab(speakerNames, raw_to_display, currFile):
    """Render a table: Speaker (with inline ✎ rename + history dropdown) | Audio Sample | ↺ button."""
    file_samples = st.session_state.speakerClips.get(currFile, {})
    has_waveform = currFile in st.session_state.speakerWaveforms
    has_samples  = bool(file_samples)

    if "inline_rename_active" not in st.session_state:
        st.session_state.inline_rename_active = {}
    if "inline_rename_history" not in st.session_state:
        st.session_state.inline_rename_history = []

    # Header
    header_cols = st.columns([3, 3, 1])
    header_cols[0].markdown("**Speaker**")
    header_cols[1].markdown("**Audio Sample**")
    header_cols[2].markdown("**&nbsp;**", unsafe_allow_html=True)
    st.markdown("<hr style='margin-top:2px; margin-bottom:4px;'>", unsafe_allow_html=True)

    if not has_samples:
        st.info("Audio samples are only available for files analyzed from audio (not RTTM/CSV/TXT).")

    for sp in speakerNames:
        display_name = raw_to_display.get(sp, sp)
        edit_key     = (currFile, sp)
        is_editing   = st.session_state.inline_rename_active.get(edit_key, False)

        row_cols = st.columns([3, 3, 1])

        # --- Speaker cell ---
        if is_editing:
            select_key  = f"inline_rename_select_{currFile}_{sp}"
            input_key   = f"inline_rename_input_{currFile}_{sp}"
            confirm_key = f"inline_rename_confirm_{currFile}_{sp}"
            cancel_key  = f"inline_rename_cancel_{currFile}_{sp}"

            history     = st.session_state.inline_rename_history
            current_val = display_name if display_name != sp else ""

            with row_cols[0]:
                if history:
                    # Selectbox with index=None shows placeholder and lets user
                    # type to filter — single combined widget, no separate input.
                    options     = history
                    default_idx = history.index(current_val) if current_val in history else None
                    chosen = st.selectbox(
                        "Rename",
                        options=options,
                        index=default_idx,
                        key=select_key,
                        label_visibility="collapsed",
                        placeholder="Type or select a name…",
                    )
                    new_name = chosen or ""
                else:
                    # No history yet — plain text input
                    new_name = st.text_input(
                        "Rename",
                        value=current_val,
                        key=input_key,
                        label_visibility="collapsed",
                        placeholder=f"Rename {sp}…",
                    )

                btn_col1, btn_col2 = st.columns(2)
                if btn_col1.button("✓", key=confirm_key, help="Confirm rename"):
                    confirmed = new_name.strip()
                    if confirmed:
                        apply_inline_rename(currFile, sp, confirmed)
                        if confirmed not in st.session_state.inline_rename_history:
                            st.session_state.inline_rename_history.append(confirmed)
                    st.session_state.inline_rename_active[edit_key] = False
                    st.rerun()
                if btn_col2.button("✕", key=cancel_key, help="Cancel"):
                    st.session_state.inline_rename_active[edit_key] = False
                    st.rerun()
        else:
            with row_cols[0]:
                name_col, pencil_col = st.columns([4, 1])
                name_col.write(display_name)
                if pencil_col.button(
                    "✎", key=f"inline_rename_edit_{currFile}_{sp}",
                    help=f"Rename {sp}",
                ):
                    st.session_state.inline_rename_active[edit_key] = True
                    st.rerun()

        # --- Audio sample cell ---
        if sp in file_samples:
            row_cols[1].audio(file_samples[sp], format="audio/wav")
            sp_segs = st.session_state.speakerSegments.get(currFile, {}).get(sp, [])
            if has_waveform and sp_segs:
                if row_cols[2].button(
                    "↺", key=f"sample_randomize_{currFile}_{sp}",
                    help="Try a different audio sample for this speaker",
                ):
                    randomize_speaker_clip(currFile, sp)
                    st.rerun()
        else:
            row_cols[1].write("—")


# ---------------------------------------------------------------------------
# Multi-file summary expander
# ---------------------------------------------------------------------------

def render_multifile_summary(plotly_config=None):
    """Render the Multi-file Summary Data expander if enough files are analyzed."""
    cfg = plotly_config or {"displayModeBar": True, "modeBarButtonsToRemove": []}
    if not st.session_state.results:
        return

    with st.expander("Multi-file Summary Data"):
        st.header("Multi-file Summary Data")
        with st.spinner("Processing summary results..."):
            validNames = [
                fn for fn in st.session_state.file_names
                if fn in st.session_state.results
                and len(st.session_state.results[fn]) == 2
            ]
            if len(validNames) <= 1:
                return

            df6, allCategories = utils.build_multifile_category_df(
                validNames, st.session_state.results, st.session_state.summaries,
                st.session_state.categories, st.session_state.categorySelect,
            )
            st.plotly_chart(
                px.bar(df6, x="files", y=allCategories,
                       title="Time Spoken by Each Speaker in Each File"),
                use_container_width=True, config=cfg,
            )

            df7, _ = utils.build_multifile_voice_df(validNames, st.session_state.summaries)
            for sort_cols, ascending, title in [
                (["One Voice",  "Multi Voice"], True,  "Cross-file Voice Categories sorted for One Voice"),
                (["Multi Voice","One Voice"],   True,  "Cross-file Voice Categories sorted for Multi Voice"),
                (["No Voice",   "Multi Voice"], False, "Cross-file Voice Categories sorted for Any Voice"),
            ]:
                st.plotly_chart(
                    px.bar(df7.sort_values(by=sort_cols, ascending=ascending),
                           x="files", y=["One Voice", "Multi Voice", "No Voice"],
                           title=title),
                    use_container_width=True, config=cfg,
                )