app.py

from __future__ import annotations

import logging
import os
import time
from typing import Any

import pandas as pd
import streamlit as st

from charts import (
    build_class_small_multiples,
    build_latest_comparison_chart,
    build_trend_chart,
)
from config import AppConfig, load_config
from hf_client import HubClient, RepoRevision
from parsing import extract_performance_metrics_table
from transforms import (
    add_latest_flags,
    build_latest_snapshot_table,
    filter_history,
    normalize_metrics_dataframe,
    table_to_long_dataframe,
)

try:
    from streamlit_autorefresh import st_autorefresh
except Exception:  # pragma: no cover
    def st_autorefresh(*args: Any, **kwargs: Any) -> None:
        return None


CACHE_TTL_SECONDS = int(os.getenv("EVAL_CACHE_TTL_SECONDS", "600"))
LOG_LEVEL = os.getenv("EVAL_LOG_LEVEL", "INFO").upper()
logging.basicConfig(level=getattr(logging, LOG_LEVEL, logging.INFO))
logger = logging.getLogger("segmentation_eval_dashboard")


st.set_page_config(
    page_title="Segmentation Evaluation Dashboard",
    layout="wide",
)


@st.cache_data(ttl=CACHE_TTL_SECONDS, show_spinner="Loading metrics from Hugging Face Hub...")
def collect_dashboard_data(config: AppConfig) -> tuple[pd.DataFrame, pd.DataFrame, list[str]]:
    client = HubClient(config)
    repos = client.discover_dataset_repos()
    logger.info("Discovered %d candidate repos", len(repos))
    start_time = time.monotonic()

    history_frames: list[pd.DataFrame] = []
    status_rows: list[dict[str, Any]] = []

    for repo_id in repos:
        elapsed = time.monotonic() - start_time
        if elapsed >= config.max_load_seconds:
            status_rows.append(
                {
                    "repo_id": "__load_guard__",
                    "status": "partial",
                    "revisions_checked": 0,
                    "revisions_parsed": 0,
                    "points": 0,
                    "latest_timestamp": pd.NaT,
                    "message": f"Stopped early after {config.max_load_seconds}s to keep app responsive",
                }
            )
            logger.warning("Load guard triggered after %.1fs", elapsed)
            break

        revisions_checked = 0
        revisions_parsed = 0
        points = 0
        message = ""
        commit_error_message = ""

        try:
            commit_limit = config.max_commits_per_repo if config.include_history else 1
            revisions = client.list_dataset_commits(repo_id=repo_id, max_commits=max(1, commit_limit))
        except Exception as exc:  # pragma: no cover
            revisions = [RepoRevision(repo_id=repo_id, revision="main", committed_at=None)]
            commit_error_message = f"Commit listing failed ({type(exc).__name__}); fallback to main"
            logger.warning("Repo %s: %s", repo_id, commit_error_message)

        repo_frames: list[pd.DataFrame] = []
        seen_revisions: set[str] = set()
        for revision_info in revisions:
            revision = revision_info.revision
            if revision in seen_revisions:
                continue

            revisions_checked += 1
            seen_revisions.add(revision)

            readme = client.fetch_dataset_readme(repo_id=repo_id, revision=revision)
            if not readme:
                continue

            metrics_table = extract_performance_metrics_table(
                markdown=readme,
                heading=config.metrics_section_heading,
            )
            if metrics_table is None or metrics_table.empty:
                continue

            normalized = normalize_metrics_dataframe(metrics_table)
            if normalized.empty:
                continue

            long_df = table_to_long_dataframe(
                dataframe=normalized,
                repo_id=repo_id,
                revision=revision,
                committed_at=revision_info.committed_at,
                metrics=config.metrics,
            )
            if long_df.empty:
                continue

            revisions_parsed += 1
            points += len(long_df)
            repo_frames.append(long_df)

        if repo_frames:
            repo_history = pd.concat(repo_frames, ignore_index=True)
            history_frames.append(repo_history)
            latest_timestamp = repo_history["timestamp"].max()
            status = "ok"
            message = f"Parsed {revisions_parsed}/{revisions_checked} revisions"
            if commit_error_message:
                message = f"{message}. {commit_error_message}"
        else:
            latest_timestamp = pd.NaT
            status = "error" if commit_error_message else "no_metrics"
            message = commit_error_message or "No valid Performance Metrics table found"

        status_rows.append(
            {
                "repo_id": repo_id,
                "status": status,
                "revisions_checked": revisions_checked,
                "revisions_parsed": revisions_parsed,
                "points": points,
                "latest_timestamp": latest_timestamp,
                "message": message,
            }
        )
        logger.info(
            "Repo %s: status=%s revisions_checked=%d revisions_parsed=%d points=%d",
            repo_id,
            status,
            revisions_checked,
            revisions_parsed,
            points,
        )

    if history_frames:
        history_df = pd.concat(history_frames, ignore_index=True)
        history_df = history_df.sort_values(["timestamp", "repo_id", "revision", "class", "metric"])  # type: ignore[arg-type]
        history_df = add_latest_flags(history_df)
    else:
        history_df = pd.DataFrame(
            columns=["repo_id", "revision", "timestamp", "class", "metric", "value", "std", "is_latest"]
        )

    status_df = pd.DataFrame(status_rows)
    if not status_df.empty and "latest_timestamp" in status_df.columns:
        status_df["latest_timestamp"] = pd.to_datetime(status_df["latest_timestamp"], utc=True, errors="coerce")

    return history_df, status_df, repos


def main() -> None:
    config = load_config()

    if config.auto_refresh_seconds > 0:
        st_autorefresh(interval=config.auto_refresh_seconds * 1000, key="dashboard-refresh")

    st.title("Segmentation Evaluation Metrics Dashboard")
    st.caption("Source: Hugging Face dataset cards (`### Performance Metrics`).")

    with st.sidebar:
        st.header("Controls")
        if st.button("Refresh now", use_container_width=True):
            st.cache_data.clear()
            st.rerun()

        st.caption(
            "Filtering/discovery is configured by env vars: "
            "`HF_OWNER`, `EVAL_REPO_ALLOWLIST`, `EVAL_REPO_PREFIXES`, `EVAL_REPO_ID_CONTAINS`."
        )
        show_error_bands = st.checkbox("Show std error bands", value=True)
        show_small_multiples = st.checkbox("Show class small multiples", value=False)

    history_df, status_df, discovered_repos = collect_dashboard_data(config)
    logger.info(
        "UI load complete: discovered=%d parsed_rows=%d status_rows=%d",
        len(discovered_repos),
        len(history_df),
        len(status_df),
    )

    if not discovered_repos:
        st.warning("No dataset repositories were discovered. Check owner/prefix/allowlist environment variables.")
        st.stop()

    available_repos = sorted(history_df["repo_id"].dropna().unique().tolist()) if not history_df.empty else discovered_repos
    available_classes = sorted(history_df["class"].dropna().unique().tolist()) if not history_df.empty else []
    available_metrics = sorted(history_df["metric"].dropna().unique().tolist()) if not history_df.empty else list(config.metrics)

    with st.sidebar:
        selected_repos = st.multiselect("Repos", options=available_repos, default=available_repos)
        selected_classes = st.multiselect("Classes", options=available_classes, default=available_classes)
        default_metrics = [metric for metric in config.metrics if metric in available_metrics] or available_metrics
        selected_metrics = st.multiselect("Metrics", options=available_metrics, default=default_metrics)

    filtered_history = filter_history(
        history_df=history_df,
        repos=selected_repos,
        classes=selected_classes,
        metrics=selected_metrics,
    )

    latest_filtered = filtered_history[filtered_history["is_latest"]].copy()

    success_repos = int((status_df["status"] == "ok").sum()) if not status_df.empty else 0
    latest_runs = latest_filtered[["repo_id", "revision"]].drop_duplicates().shape[0] if not latest_filtered.empty else 0
    latest_f1 = latest_filtered[latest_filtered["metric"] == "F1"]["value"].mean() if not latest_filtered.empty else float("nan")
    latest_timestamp = filtered_history["timestamp"].max() if not filtered_history.empty else pd.NaT

    col1, col2, col3, col4 = st.columns(4)
    col1.metric("Discovered repos", len(discovered_repos))
    col2.metric("Repos parsed", success_repos)
    col3.metric("Latest runs in view", latest_runs)
    col4.metric("Latest mean F1", f"{latest_f1:.3f}" if pd.notna(latest_f1) else "n/a")

    if pd.notna(latest_timestamp):
        st.caption(f"Most recent datapoint: {pd.to_datetime(latest_timestamp).strftime('%Y-%m-%d %H:%M:%S %Z')}")

    st.subheader("Trend Over Time")
    trend_chart = build_trend_chart(filtered_history, show_error_bands=show_error_bands)
    if trend_chart is None:
        st.info("No trend data available for the active filters.")
    else:
        st.altair_chart(trend_chart, use_container_width=True)

    st.subheader("Latest Snapshot")
    snapshot_table = build_latest_snapshot_table(filtered_history)
    if snapshot_table.empty:
        st.info("No latest snapshot rows for current filters.")
    else:
        st.dataframe(snapshot_table, use_container_width=True, hide_index=True)

    st.subheader("Repo Comparison (Latest)")
    comparison_metric_options = selected_metrics or available_metrics
    comparison_metric = None
    if comparison_metric_options:
        comparison_metric = st.selectbox(
            "Metric for repo comparison",
            options=comparison_metric_options,
            index=0,
        )

    if comparison_metric is not None:
        comparison_chart = build_latest_comparison_chart(filtered_history, metric=comparison_metric)
        if comparison_chart is None:
            st.info("No latest comparison data available for this metric.")
        else:
            st.altair_chart(comparison_chart, use_container_width=True)

    if show_small_multiples and comparison_metric:
        st.subheader("Per-Class Small Multiples")
        small_multiples_chart = build_class_small_multiples(filtered_history, metric=comparison_metric)
        if small_multiples_chart is None:
            st.info("No small-multiple data available for this metric.")
        else:
            st.altair_chart(small_multiples_chart, use_container_width=True)

    st.subheader("Export")
    csv_bytes = filtered_history.to_csv(index=False).encode("utf-8")
    st.download_button(
        "Export filtered history as CSV",
        data=csv_bytes,
        file_name="segmentation_metrics_history.csv",
        mime="text/csv",
        use_container_width=False,
    )

    with st.expander("Health / Status", expanded=False):
        if status_df.empty:
            st.info("No status rows to display.")
        else:
            st.dataframe(status_df.sort_values(["status", "repo_id"]), use_container_width=True, hide_index=True)

    with st.expander("Raw Parsed Data", expanded=False):
        if filtered_history.empty:
            st.info("No parsed rows for current filters.")
        else:
            st.dataframe(
                filtered_history.sort_values(["timestamp", "repo_id", "class", "metric"], ascending=[False, True, True, True]),
                use_container_width=True,
                hide_index=True,
            )


if __name__ == "__main__":
    main()