"""
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 tempfile
from pathlib import Path

import streamlit as st
import torch
import pandas as pd
from pyannote.audio import Pipeline
from pyannote.audio.pipelines import SpeakerDiarization
from pyannote.audio.models.segmentation import PyanNet
from pyannote.audio import Inference
from pyannote.pipeline.parameter import ParamDict

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,
)

# ---------------------------------------------------------------------------
# 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        = "sample.rttm"
DEMO_SAMPLE_PATH = "sample_short.rttm"
MULTI_DEMO_PATHS = [
    "audioSamples/media-afc-cal-afc1986022_sr01a05.rttm",
    "audioSamples/media-afc-cal-afc1986022_sr34a01.rttm",
    "audioSamples/media-afc-cal-afc1986022_sr14b02.rttm",
    "audioSamples/media-afc-cal-afc1986022_sr52a02.rttm",
    "audioSamples/media-afc-cal-afc1986022_sr14b01.rttm",
]

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

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

isGPU  = torch.cuda.is_available()
device = torch.device("cuda" if isGPU else "cpu")
print(f"Using {device}")

if ENABLE_DENOISE:
    from df import init_df
    dfModel, dfState, _ = init_df(model_base_dir="DeepFilterNet3")
    dfModel.to(device)
else:
    dfModel = dfState = None

# Load baseline model
pipeline = Pipeline.from_pretrained("pyannote/speaker-diarization-3.1")

# Generate replacement segmentation model
newSpecs = pipeline._segmentation.model.specifications
segModel = PyanNet.from_pretrained('20251208_Sonogram_Segmentation.ckpt')
segModel.specifications = newSpecs
segmentation_duration = segModel.specifications.duration

# Update baseline with training
pipeline._segmentation = Inference(
    segModel,
    duration=segmentation_duration,
    step=pipeline.segmentation_step * segmentation_duration,
    skip_aggregation=True,
    batch_size=1,
)

pipeline.to(device)

# ---------------------------------------------------------------------------
# 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 isGPU:
    st.warning("TOOL CURRENTLY USING CPU, ANALYSIS EXTREMELY SLOW")

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 Connor Young at "
    "<a href='mailto:czyoung@ualr.edu'>czyoung@ualr.edu</a> if you need help using the tool!</p>"
    "<p style='margin-top:16px; margin-bottom:0px;'>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>"
    "<hr style='margin-top:16px; margin-bottom:16px;'>",
    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("(Potentially) 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}",
)

temp_dir = tempfile.mkdtemp()

if uploaded_file_paths:
    for uploaded_file in uploaded_file_paths:
        if not uploaded_file.name.lower().endswith(SUPPORTED_FILE_TYPES):
            st.error(f"File must be of type: {SUPPORTED_FILE_TYPES}")
            continue
        fname = uploaded_file.name
        path  = os.path.join(temp_dir, fname)
        with open(path, "wb") as f:
            f.write(uploaded_file.getvalue())
        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

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. "
    "The latter can take a while to finish."
)

# ---------------------------------------------------------------------------
# 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, pipeline,
        ENABLE_DENOISE, EARLY_CLEANUP,
        GAIN_WINDOW, MINIMUM_GAIN, MAXIMUM_GAIN,
        dfModel, dfState, ATTEN_LIM_DB,
    )

# ---------------------------------------------------------------------------
# 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
    ):
        raise ValueError("File not yet analyzed")

    st.header(f"Analysis of file {currFile}")
    TAB_NAMES = ["Speakers & Roles", "Pie Chart", "Sunburst",
                 "Treemap", "Time Spoken", "Timeline", "Download"]
    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 = su.colorsCSS(3)
    allColors     = su.colorsCSS(len(speakerNames) + len(st.session_state.categories))
    speakerColors = allColors[:len(speakerNames)]
    catColors     = allColors[len(speakerNames):]

    # 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
        sp_labels = {sp: f"{fn}: {get_display_name(sp, 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"{fn}: {get_display_name(sp, fn)} ({sp})"] = 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
    # -----------------------------------------------------------------------

    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)
        csv = convert_df(displayDF)
        st.download_button(
            f"Download {currPlainName}.csv", csv,
            f"sonogram-analysis-{currPlainName}.csv", "text/csv",
            key="download-csv", on_click="ignore",
        )

        analyzed_count = sum(
            1 for r in st.session_state.results.values() if len(r) == 2
        )
        zip_bytes = build_all_csv_zip() if analyzed_count > 1 else b""
        st.download_button(
            "Download all analyzed data in .csv (.zip)", 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, speakerColors, 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, speakerColors, 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,
    )
    ui.render_chart(
        utils.build_fig_treemap(df5, catTypeColors, speakerColors, 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, speakerColors, get_display_name, currFile),
        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, speakerColors, get_display_name, currFile),
        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)