"""
app.py — Streamlit entry point.

Responsibilities:
  - Constants and pipeline initialisation
  - Page header and file upload
  - Demo / Analyze buttons and analysis loop (via state.py)
  - Per-file view: sidebar + tabs (via ui.py and utils.py)
"""

import os
import random
import tempfile
from pathlib import Path

import streamlit as st
import torch
import pandas as pd
from sonogram import Sonogram

import sonogram_utility as su
import utils
import ui
from state import (
    init_session_state,
    get_display_name, apply_speaker_renames_to_df, convert_df,
    updateMultiSelect, store_speaker_clips, register_file, analyze,
    load_demo_single, load_demo_single_sample, load_demo_multi, run_analysis_loop,
    build_all_csv_zip,
    build_xml_download,
    build_all_xml_zip,
    rename_sample_demo_speakers,
)

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------

SUPPORTED_FILE_TYPES = (".wav", ".mp3", ".mp4", ".txt", ".rttm", ".csv")
ENABLE_DENOISE = False
EARLY_CLEANUP  = True
GAIN_WINDOW    = 4
MINIMUM_GAIN   = -45
MAXIMUM_GAIN   = -5
ATTEN_LIM_DB   = 3

PLOTLY_CONFIG = {"displayModeBar": True, "modeBarButtonsToRemove": []}

PARQUET_DATASET_DIR = Path("parquet_dataset")
PARQUET_DATASET_DIR.mkdir(parents=True, exist_ok=True)

DEMO_PATH        = "audioSamples/sample.rttm"
DEMO_SAMPLE_PATH = "audioSamples/sample_short.rttm"
MULTI_DEMO_PATHS = [
    f"audioSamples/class{i:02d}.rttm" for i in range(1, 11)
]

# ---------------------------------------------------------------------------
# Pipeline initialisation  (once per server process)
# ---------------------------------------------------------------------------

torch.classes.__path__ = [os.path.join(torch.__path__[0], torch.classes.__file__)]

if "pipeline" not in st.session_state or st.session_state.pipeline is None:
    pipeline = Sonogram()
    st.session_state.pipeline = pipeline
pipeline = st.session_state.pipeline
# ---------------------------------------------------------------------------
# Session state + global styles
# ---------------------------------------------------------------------------

init_session_state()

# Uploader key is rotated on reset to force the widget to clear
if "uploader_key" not in st.session_state:
    st.session_state.uploader_key = 0

st.markdown(
    "<style>"
    "details > summary { font-size: 1rem; font-weight: 500; }"
    ".stTabs [data-baseweb='tab'] { font-size: 1rem; }"
    ".stFileUploader label { font-size: 1rem; }"
    "[data-baseweb='tag'] span { cursor: default; }"
    "[data-baseweb='menu'] li { white-space: nowrap; overflow: hidden; text-overflow: ellipsis; max-width: 260px; }"
    "[data-baseweb='menu'] li:hover { overflow: visible; white-space: normal; z-index: 9999; }"
    "</style>",
    unsafe_allow_html=True,
)

# ---------------------------------------------------------------------------
# Page header
# ---------------------------------------------------------------------------

st.title("Instructor Support Tool")
if not (pipeline.isGPU or pipeline.isTPU):
    st.warning('''
        TOOL CURRENTLY USING CPU, ANALYSIS EXTREMELY SLOW  
        After performing one action, a RUNNING... icon will appear at the top right corner.  
        Please wait for this icon to disappear before performing another action.
    ''')

DEMO_VIDEO_URL = "https://www.youtube.com/watch?v=n7Z_BzU7yOY"

with st.expander("TL;DR"):
    st.write("Please watch this video for a walkthrough of this app:")
    st.video(DEMO_VIDEO_URL)
    st.write(
        'If you would like to see a sample result or multiple sample results generated from '
        'real classroom audio, select "Single File Demo (Sample)", "Single File Demo (Full)", '
        'or "Multiple Files Demo" on the left sidebar.'
    )
    st.markdown(
        "<p style='margin-bottom:4px;'>Keep in mind that this is a very early draft of the tool. "
        "Please be patient with any bugs/errors, and email Thien Duong at "
        "<a href='mailto:tduong@ualr.edu'>tduong@ualr.edu</a> if you need help using the tool!</p>"
        "<p style='margin-top:8px; margin-bottom:4px;'>Would you like additional data, charts, or features? "
        "<a href='https://forms.gle/A32CdfGYSZoMPyyX9'>Tell us about our project!</a></p>"
        "<p style='margin-bottom:4px;'>If you would like to learn more or work with us, contact Dr. Mark Baillie at "
        "<a href='mailto:mtbaillie@ualr.edu'>mtbaillie@ualr.edu</a></p>",
        unsafe_allow_html=True,
    )

with st.expander("Instructions and additional details"):
    st.write("Thank you for viewing our experimental app! The overall presentations and features are expected to be improved over time.")
    st.write("To use this app:\n1. Upload an audio file for live analysis. Alternatively, upload an already generated [rttm file](https://stackoverflow.com/questions/30975084/rttm-file-format)")
    st.write("2. Press Analyze All. No data is saved on our side.")
    st.write("3. Use the sidebar to select your file. Multiple files are supported for more comprehensive analysis.")
    st.write("4. Use the tabs to view different visualizations. Each can be downloaded.")
    st.write("4a. Graphs are built with [plotly](https://plotly.com/). Double-click to reset. [More examples](https://plotly.com/python/basic-charts/).")

with st.expander("Preliminary FAQ"):
    st.write("**1. I tried analyzing a file, but the page refreshed and nothing happened! Why?**")
    st.write("You may need to select a file using the sidebar on the left.")
    st.write("**2. I don't see a sidebar! Where is it?**")
    st.write("Press the '>' in the upper left to expand the sidebar.")
    st.write("**3. I still don't have a file to select in the dropdown! Why?**")
    st.write("Your file may be too large. We currently support approximately 1.5 hours of audio.")
    st.write("**4. I want to view my previously analyzed data. How?**")
    st.write("Download a CSV copy from the Download tab and re-upload it later.")
    st.write("**5. The app is extremely slow. What is wrong?**")
    st.write("We are securing funding for permanent GPU access. Until then, CPU analysis may take a very long time.")

# ---------------------------------------------------------------------------
# File upload
# ---------------------------------------------------------------------------

uploaded_file_paths = st.file_uploader(
    "Upload an audio of classroom activity to analyze",
    accept_multiple_files=True,
    key=f"uploader_{st.session_state.uploader_key}",
)

# Make temp directory
temp_dir = tempfile.mkdtemp()

# If files have been uploaded
if uploaded_file_paths:
    for uploaded_file in uploaded_file_paths:
        # Check for supported file types via extension
        if not uploaded_file.name.lower().endswith(SUPPORTED_FILE_TYPES):
            st.error(f"File must be of type: {SUPPORTED_FILE_TYPES}")
            continue
        # Save file locally to server
        fname = uploaded_file.name
        path  = os.path.join(temp_dir, fname)
        with open(path, "wb") as f:
            f.write(uploaded_file.getvalue())
        # Add file to list of known files
        if fname not in st.session_state.file_names:
            register_file(fname)
        st.session_state.file_paths[fname] = path
    st.session_state.valid_files = list(st.session_state.file_names)

file_names      = st.session_state.file_names
file_paths_dict = st.session_state.file_paths

# ---------------------------------------------------------------------------
# Sidebar: demo buttons
# ---------------------------------------------------------------------------

isDemo = False

if st.sidebar.button("Single File Demo (Sample)"):
    load_demo_single_sample(DEMO_SAMPLE_PATH)
    isDemo = True
    st.session_state["_sample_demo_loaded"] = True

# Only show between Sample and Full buttons once Sample has been loaded
_sample_demo_fname = DEMO_SAMPLE_PATH.split("/")[-1]
if (st.session_state.get("_sample_demo_loaded", False)
        and _sample_demo_fname in st.session_state.file_names):
    if st.sidebar.button("Rename Sample Speakers"):
        rename_sample_demo_speakers(_sample_demo_fname)

if st.sidebar.button("Single File Demo (Full)"):
    load_demo_single(DEMO_PATH)
    isDemo = True

if st.sidebar.button("Multiple Files Demo"):
    load_demo_multi(MULTI_DEMO_PATHS)
    isDemo = True

st.sidebar.caption(
    "Single File Demo (Sample) analyzes a 10-minute sample, while "
    "Single File Demo (Full) analyzes the entire file. "
    "Multiple Files Demo loads ten recording samples."
)
st.sidebar.caption(
    """After Single File Demo (Sample) is loaded, the Rename Sample Speakers button will appear.
    Click on Rename Sample Speakers to try out the rename feature.
""")

# ---------------------------------------------------------------------------
# Analyze All / Reset buttons
# ---------------------------------------------------------------------------

if len(file_names) == 0:
    st.text("Upload file(s) to enable analysis")
else:
    col_analyze, col_spacer, col_reset = st.columns([3, 5, 2])
    with col_analyze:
        if st.button("Analyze All New Audio", key="button_all"):
            st.session_state.analyzeAllToggle = True
    with col_reset:
        if st.button("🗑️ Reset App", key="button_reset", type="secondary", use_container_width=True):
            next_key = st.session_state.uploader_key + 1
            for key in list(st.session_state.keys()):
                del st.session_state[key]
            st.session_state.uploader_key = next_key
            st.rerun()

# ---------------------------------------------------------------------------
# Analysis loop
# ---------------------------------------------------------------------------

if st.session_state.analyzeAllToggle:
    run_analysis_loop(
        file_names, file_paths_dict
    )

# ---------------------------------------------------------------------------
# File selector
# ---------------------------------------------------------------------------

currFile = st.sidebar.selectbox(
    "Current File", file_names, on_change=updateMultiSelect, key="select_currFile"
)
if isDemo:
    currFile = file_names[0]

st.sidebar.divider()

if currFile is None:
    st.write("Select a file to view from the sidebar")

# ---------------------------------------------------------------------------
# Per-file analysis view
# ---------------------------------------------------------------------------

try:
    if currFile is None:
        raise ValueError("No file selected")

    st.session_state.resetResult = False
    currPlainName = currFile.rsplit(".", 1)[0]

    if not (
        currFile in st.session_state.results
        and currFile in st.session_state.summaries
        and len(st.session_state.results[currFile]) > 0
        and st.session_state.summaries[currFile].get("speakers_dataFrame") is not None
    ):
        raise ValueError("File not yet analyzed")

    st.header(f"Analysis of file {currFile}")
    TAB_NAMES = ["Population", "Speakers & Roles", "Pie Chart", "Sunburst",
                 "Treemap", "Time Spoken", "Timeline", "Download"]
    populationTab, renameTab, pie2, sunburst1, treemap1, bar1, timeline, dataTab = st.tabs(TAB_NAMES)

    currAnnotation, currTotalTime = st.session_state.results[currFile]
    speakerNames       = currAnnotation.labels()
    speakers_dataFrame = st.session_state.summaries[currFile]["speakers_dataFrame"]
    currDF, _          = su.annotationToSimpleDataFrame(currAnnotation)
    unusedSpeakers     = st.session_state.unusedSpeakers[currFile]
    # categorySelect is global tokens; extract raw speaker IDs for currFile
    _prefix = currFile + ": "
    categorySelections = [
        [t[len(_prefix):] for t in tokens if t.startswith(_prefix)]
        for tokens in st.session_state.categorySelect
    ]

    _saved_renames       = st.session_state.speakerRenames.get(currFile, {})
    raw_to_display       = {sp: _saved_renames.get(sp, sp) for sp in speakerNames}
    all_speakers_display = [raw_to_display[sp] for sp in speakerNames]

    catTypeColors = utils.colorsCSS(3)
    # Assign each speaker a fixed color by their position in speakerNames so
    # the same speaker always gets the same color across all charts.
    # Map by display name so renames are also covered.
    speaker_color_map = {
        get_display_name(sp, currFile): utils._SPEAKER_PALETTE[i % len(utils._SPEAKER_PALETTE)]
        for i, sp in enumerate(speakerNames)
    }
    speakerColors = list(speaker_color_map.values())
    catColors     = utils.colorsCSS(len(st.session_state.categories))

    # Rebuild live df4
    # Guard: trim categorySelections to match categories length in case they
    # are momentarily out of sync during rapid role assignments.
    nameList    = st.session_state.categories
    _selections = categorySelections[:len(nameList)]
    while len(_selections) < len(nameList):
        _selections.append([])
    valueList   = [su.sumTimes(currAnnotation.subset(s)) for s in _selections]
    categorySelections = _selections
    extraNames  = list(unusedSpeakers)
    extraValues = [su.sumTimes(currAnnotation.subset([sp])) for sp in unusedSpeakers]
    st.session_state.summaries[currFile]["df4"] = pd.DataFrame(
        {"names": nameList + extraNames, "values": valueList + extraValues}
    )

    # Build all_speaker_tokens for rename sidebar
    # raw token    : "file: SPEAKER_00"  (stored in data model)
    # display token: "file: John"        (shown in dropdowns)
    all_speaker_tokens = [
        f"{fn}: {sp}"
        for fn in st.session_state.file_names
        if fn in st.session_state.results and len(st.session_state.results[fn]) == 2
        for sp in st.session_state.results[fn][0].labels()
    ]
    # Map display label -> raw token so ui.py can translate back after selection.
    # If two speakers in the same file share a display name, append the raw ID
    # to disambiguate both entries.
    token_display_map = {}
    for fn in st.session_state.file_names:
        if not (fn in st.session_state.results and len(st.session_state.results[fn]) == 2):
            continue
        # Display label puts speaker name first so it's visible when truncated.
        # Raw token keeps "fn: sp" format — all parsing logic depends on this.
        sp_labels = {sp: f"{get_display_name(sp, fn)}: {fn}"
                     for sp in st.session_state.results[fn][0].labels()}
        label_counts = {}
        for disp in sp_labels.values():
            label_counts[disp] = label_counts.get(disp, 0) + 1
        for sp, disp in sp_labels.items():
            raw = f"{fn}: {sp}"
            if label_counts[disp] > 1:
                token_display_map[f"{get_display_name(sp, fn)} ({sp}): {fn}"] = raw
            else:
                token_display_map[disp] = raw
    display_speaker_tokens = list(token_display_map.keys())

    # -----------------------------------------------------------------------
    # Sidebar
    # -----------------------------------------------------------------------

    ui.render_role_sidebar(display_speaker_tokens, token_display_map)
    ui.render_rename_sidebar(currFile, speakerNames, display_speaker_tokens, token_display_map)

    # -----------------------------------------------------------------------
    # Tab: Data
    # -----------------------------------------------------------------------

    # -----------------------------------------------------------------------
    # Tab: Population
    # -----------------------------------------------------------------------

    with populationTab:
        st.markdown("Enter the number of students present in each recording.")
        pop_header_cols = st.columns([4, 2, 1, 1])
        pop_header_cols[0].markdown(
            "<p style='margin:0; font-weight:600;'>File Name</p>",
            unsafe_allow_html=True,
        )
        pop_header_cols[1].markdown(
            "<p style='margin:0; font-weight:600;'>Number of Students</p>",
            unsafe_allow_html=True,
        )
        st.markdown("<hr style='margin-top:2px; margin-bottom:4px;'>",
                    unsafe_allow_html=True)
        for fn in st.session_state.file_names:
            if not (fn in st.session_state.results
                    and len(st.session_state.results[fn]) == 2):
                continue
            plain    = fn.rsplit(".", 1)[0]
            edit_key = f"pop_editing_{fn}"
            input_key = f"pop_input_{fn}"
            current_pop = st.session_state.studentPopulations.get(fn)

            row_cols = st.columns([4, 2, 1, 1])
            row_cols[0].write(plain)

            if st.session_state.get(edit_key, False):
                # Render input and buttons inside the same row columns so
                # nothing spills below the row
                row_cols[1].number_input(
                    "Students", min_value=0, step=1,
                    value=int(current_pop) if current_pop is not None else 0,
                    key=input_key, label_visibility="collapsed",
                )
                if row_cols[2].button("✓", key=f"pop_confirm_{fn}"):
                    st.session_state.studentPopulations[fn] = int(
                        st.session_state.get(input_key, current_pop or 0)
                    )
                    st.session_state[edit_key] = False
                    st.rerun()
                if row_cols[3].button("✕", key=f"pop_cancel_{fn}"):
                    st.session_state[edit_key] = False
                    st.rerun()
            else:
                row_cols[1].write(str(current_pop) if current_pop is not None else "—")
                if row_cols[2].button("✎", key=f"pop_edit_btn_{fn}",
                                      help="Edit student count"):
                    st.session_state[edit_key] = True
                    st.rerun()

    # -----------------------------------------------------------------------
    # Tab: Download
    # -----------------------------------------------------------------------

    with dataTab:
        # Build raw-speaker -> role lookup against the RAW currDF (before renames
        # are applied) so the map keys match the original SPEAKER_## labels.
        raw_to_role = {
            token.split(": ", 1)[1]: st.session_state.categories[i]
            for i, tokens in enumerate(st.session_state.categorySelect)
            for token in tokens
            if token.startswith(f"{currFile}: ")
        }
        displayDF = currDF.copy()
        displayDF["Role"] = displayDF["Resource"].map(raw_to_role).fillna("")
        displayDF = apply_speaker_renames_to_df(displayDF, currFile, column="Resource")
        displayDF = displayDF.drop(columns=["Task"], errors="ignore")
        displayDF = displayDF.rename(columns={"Resource": "Speaker"})
        if "Start" in displayDF.columns:
            displayDF = displayDF.sort_values("Start").reset_index(drop=True)

        xml_bytes = build_xml_download(currFile)
        st.download_button(
            f"Download {currPlainName}.xml", xml_bytes,
            f"sonogram-analysis-{currPlainName}.xml", "application/xml",
            key="download-xml", on_click="ignore",
        )

        analyzed_count = sum(
            1 for r in st.session_state.results.values() if len(r) == 2
        )
        zip_bytes = build_all_xml_zip() if analyzed_count > 1 else b""
        st.download_button(
            "Download all analyzed datas.xml in .zip archive", zip_bytes,
            "sonogram-analysis-all.zip", "application/zip",
            key="download-all-zip", on_click="ignore",
            disabled=(analyzed_count <= 1),
        )

    # -----------------------------------------------------------------------
    # Tab: Rename Speaker
    # -----------------------------------------------------------------------

    with renameTab:
        ui.render_speaker_samples_tab(speakerNames, raw_to_display, currFile)

    # -----------------------------------------------------------------------
    # Charts
    # -----------------------------------------------------------------------

    df4 = st.session_state.summaries[currFile]["df4"].copy()
    df5 = st.session_state.summaries[currFile]["df5"].copy()
    df2 = st.session_state.summaries[currFile]["df2"].copy()

    ui.render_chart(
        utils.build_fig_pie2(df4, speakerNames, speaker_color_map, catColors, get_display_name, currFile),
        pie2,
        "ascn_pie2.pdf", "ascn_pie2.svg",
        f"sonogram-speaker-percent-{currPlainName}.pdf",
        f"sonogram-speaker-percent-{currPlainName}.svg",
        "download-pdf2", "download-svg2", PLOTLY_CONFIG,
    )
    ui.render_chart(
        utils.build_fig_sunburst(df5, catTypeColors, speaker_color_map, get_display_name, currFile),
        sunburst1,
        "ascn_sunburst.pdf", "ascn_sunburst.svg",
        f"sonogram-speaker-categories-{currPlainName}.pdf",
        f"sonogram-speaker-categories-{currPlainName}.svg",
        "download-pdf3", "download-svg3", PLOTLY_CONFIG,
    )
    with sunburst1:
        sv_pct = df5.loc[df5["ids"] == "OV", "percentiles"].values
        mv_pct = df5.loc[df5["ids"] == "MV", "percentiles"].values
        sv_nonzero = len(sv_pct) > 0 and sv_pct[0] > 0
        mv_nonzero = len(mv_pct) > 0 and mv_pct[0] > 0
        # Only render individual charts if BOTH categories are non-zero.
        # If either is 0%, the individual chart would be identical to the
        # combination sunburst, so skip both to avoid redundancy.
        if sv_nonzero and mv_nonzero:
            fig_sun_single = utils.build_fig_sunburst_single(
                df5, speaker_color_map, get_display_name, currFile)
            if fig_sun_single is not None:
                st.plotly_chart(fig_sun_single, use_container_width=True,
                                config=PLOTLY_CONFIG)
            fig_sun_multi = utils.build_fig_sunburst_multi(
                df5, speaker_color_map, get_display_name, currFile)
            if fig_sun_multi is not None:
                st.plotly_chart(fig_sun_multi, use_container_width=True,
                                config=PLOTLY_CONFIG)
    ui.render_chart(
        utils.build_fig_treemap(df5, catTypeColors, speaker_color_map, get_display_name, currFile),
        treemap1,
        "ascn_treemap.pdf", "ascn_treemap.svg",
        f"sonogram-treemap-{currPlainName}.pdf",
        f"sonogram-treemap-{currPlainName}.svg",
        "download-pdf4", "download-svg4", PLOTLY_CONFIG,
    )
    ui.render_chart(
        utils.build_fig_timeline(speakers_dataFrame, currTotalTime, speaker_color_map, get_display_name, currFile,
                                 mv_intervals=st.session_state.summaries[currFile].get("mv_intervals", [])),
        timeline,
        "ascn_timeline.pdf", "ascn_timeline.svg",
        f"sonogram-timeline-{currPlainName}.pdf",
        f"sonogram-timeline-{currPlainName}.svg",
        "download-pdf5", "download-svg5", PLOTLY_CONFIG,
    )
    ui.render_chart(
        utils.build_fig_bar(df2, speakerNames, catColors, speaker_color_map, get_display_name, currFile,
                            mv_per_speaker=st.session_state.summaries[currFile].get("mv_per_speaker", {})),
        bar1,
        "ascn_bar.pdf", "ascn_bar.svg",
        f"sonogram-speaker-time-{currPlainName}.pdf",
        f"sonogram-speaker-time-{currPlainName}.svg",
        "download-pdf6", "download-svg6", PLOTLY_CONFIG,
    )

except ValueError:
    pass

# ---------------------------------------------------------------------------
# Multi-file summary + footer
# ---------------------------------------------------------------------------

ui.render_multifile_summary(PLOTLY_CONFIG)