#!/usr/bin/env python3
"""
UHT Discovery - NiceGUI web interface.
"""

import os
import shutil
import inspect
import tempfile
import threading
import zipfile
from contextlib import contextmanager
from pathlib import Path
from typing import Callable, List, Optional

# NiceGUI runs workflow code in worker threads; force a headless Matplotlib backend
# to avoid macOS AppKit crashes ("NSWindow should only be instantiated on main thread").
os.environ.setdefault("MPLBACKEND", "Agg")

import pandas as pd
import plotly.express as px
import plotly.graph_objects as go
import numpy as np
from fastapi import FastAPI
from fastapi.staticfiles import StaticFiles
from nicegui import app, run, ui

from uht_discovery.run_isolation import RunIsolationContext, generate_run_id, get_workdir_base

APPLE_CSS = """
:root {
    --font-sans: -apple-system, BlinkMacSystemFont, "SF Pro Display", "Helvetica Neue", Helvetica, Arial, sans-serif;
    --bg-primary: #ffffff;
    --bg-secondary: #f5f5f7;
    --bg-glass: rgba(255, 255, 255, 0.72);
    --bg-glass-heavy: rgba(255, 255, 255, 0.85);
    --text-primary: #1d1d1f;
    --text-secondary: #6e6e73;
    --accent: #0071e3;
    --accent-hover: #0077ed;
    --border: rgba(0, 0, 0, 0.08);
    --border-input: rgba(0, 0, 0, 0.12);
    --shadow-md: 0 4px 14px rgba(0, 0, 0, 0.06);
    --radius-input: 8px;
    --radius-card: 12px;
    --sidebar-width: 260px;
    --content-max-width: 1100px;
}
[data-theme="dark"] {
    --bg-primary: #1d1d1f;
    --bg-secondary: #2d2d2f;
    --bg-glass: rgba(29, 29, 31, 0.72);
    --bg-glass-heavy: rgba(29, 29, 31, 0.85);
    --text-primary: #f5f5f7;
    --text-secondary: #a1a1a6;
    --accent: #2997ff;
    --accent-hover: #40a9ff;
    --border: rgba(255, 255, 255, 0.08);
    --border-input: rgba(255, 255, 255, 0.12);
}
body, .nicegui-content {
    font-family: var(--font-sans) !important;
    background: var(--bg-secondary) !important;
    color: var(--text-primary) !important;
    margin: 0;
}
/* --- Sidebar --- */
.apple-sidebar {
    position: fixed;
    top: 0;
    left: 0;
    width: var(--sidebar-width);
    height: 100vh;
    background: linear-gradient(180deg, var(--bg-glass-heavy) 0%, var(--bg-secondary) 100%);
    border-right: 1px solid var(--border);
    z-index: 100;
    display: flex;
    flex-direction: column;
    overflow-y: auto;
    padding: 24px 16px 16px;
    box-sizing: border-box;
}
.apple-sidebar .sidebar-brand {
    font-size: 20px;
    font-weight: 700;
    color: var(--text-primary);
    margin-bottom: 4px;
}
.apple-sidebar .sidebar-subtitle {
    font-size: 12px;
    color: var(--text-secondary);
    margin-bottom: 22px;
}
.apple-sidebar .sidebar-section {
    font-size: 11px;
    font-weight: 600;
    text-transform: uppercase;
    color: var(--text-secondary);
    margin: 16px 0 8px 8px;
}
.apple-sidebar .sidebar-item {
    display: flex;
    align-items: center;
    padding: 8px 12px;
    border-radius: 8px;
    font-size: 14px;
    font-weight: 500;
    color: var(--text-primary);
    text-decoration: none;
    margin-bottom: 2px;
    transition: background 0.15s ease;
}
.apple-sidebar .sidebar-item:hover {
    background: var(--border);
}
.apple-sidebar .sidebar-item.active {
    background: var(--accent);
    color: #ffffff;
}
.sidebar-spacer {
    flex: 1;
}
.theme-toggle {
    display: flex;
    align-items: center;
    gap: 8px;
    padding: 8px 12px;
    border-radius: 8px;
    color: var(--text-secondary);
    cursor: pointer;
    transition: background 0.15s ease;
}
.theme-toggle:hover {
    background: var(--border);
}
/* --- Content --- */
.apple-content {
    margin-left: var(--sidebar-width);
    padding: 32px 36px;
    min-height: 100vh;
    box-sizing: border-box;
}
.apple-content-inner {
    max-width: var(--content-max-width);
    margin: 0 auto;
}
/* --- Page header --- */
.apple-page-header {
    margin-bottom: 24px;
}
.apple-page-header h1 {
    font-size: 28px;
    font-weight: 700;
    color: var(--text-primary);
    margin: 0 0 4px 0;
    letter-spacing: -0.3px;
}
.apple-page-header p {
    font-size: 14px;
    color: var(--text-secondary);
    margin: 0;
    line-height: 1.5;
}
/* --- Cards --- */
.apple-card {
    background: var(--bg-glass);
    border: 1px solid var(--border);
    border-radius: var(--radius-card);
    box-shadow: var(--shadow-md);
    padding: 24px 26px;
    margin-bottom: 18px;
    animation: fadeSlideUp 0.3s ease-out both;
}
.apple-card-title {
    font-size: 17px;
    font-weight: 600;
    color: var(--text-primary);
    margin: 0 0 4px 0;
}
.apple-card-subtitle {
    font-size: 13px;
    color: var(--text-secondary);
    margin: 0 0 14px 0;
}
/* --- Form grid (2-column) --- */
.apple-form-grid {
    display: grid;
    grid-template-columns: 1fr 1fr;
    gap: 12px 20px;
    align-items: start;
}
.apple-form-grid .full-width {
    grid-column: 1 / -1;
}
/* --- Section divider inside cards --- */
.apple-section-divider {
    border: none;
    border-top: 1px solid var(--border);
    margin: 18px 0 16px 0;
}
/* --- Section label --- */
.apple-section-label {
    font-size: 11px;
    font-weight: 600;
    text-transform: uppercase;
    letter-spacing: 0.5px;
    color: var(--text-secondary);
    margin: 0 0 10px 0;
}
/* --- Help text --- */
.apple-help-text {
    font-size: 12px;
    color: var(--text-secondary);
    margin: -4px 0 8px 0;
    line-height: 1.4;
}
/* --- Fields --- */
.apple-field .q-field__control {
    border: 1px solid var(--border-input) !important;
    border-radius: var(--radius-input) !important;
    background: var(--bg-primary) !important;
}
/* --- Buttons --- */
.apple-btn {
    background: var(--accent) !important;
    color: #fff !important;
    border-radius: 999px !important;
    padding: 9px 18px !important;
    font-weight: 600 !important;
    transition: transform 0.15s ease, box-shadow 0.15s ease, background 0.15s ease !important;
}
.apple-btn:hover {
    background: var(--accent-hover) !important;
    transform: translateY(-1px);
    box-shadow: 0 4px 12px rgba(0, 113, 227, 0.25);
}
.apple-btn:active {
    transform: translateY(0);
}
/* --- Upload zone --- */
.apple-upload .q-uploader {
    border: 2px dashed var(--border-input);
    border-radius: 10px;
    background: var(--bg-primary);
    transition: border-color 0.2s ease, background 0.2s ease;
}
.apple-upload .q-uploader:hover {
    border-color: var(--accent);
    background: rgba(0, 113, 227, 0.03);
}
.apple-status {
    color: var(--text-secondary);
    font-size: 13px;
}
.apple-results a {
    color: var(--accent);
    text-decoration: none;
}
.apple-loading {
    display: flex;
    align-items: center;
    gap: 10px;
    color: var(--text-secondary);
    font-size: 13px;
    margin-top: 8px;
}
.apple-plot-wrap {
    margin-top: 16px;
}
/* --- Progress stepper --- */
.progress-stepper {
    display: flex;
    align-items: flex-start;
    gap: 0;
    margin: 16px 0;
    overflow-x: auto;
}
.progress-stepper .step {
    display: flex;
    flex-direction: column;
    align-items: center;
    flex: 1;
    min-width: 100px;
    position: relative;
}
.progress-stepper .step-dot {
    width: 14px;
    height: 14px;
    border-radius: 50%;
    background: var(--border-input);
    border: 2px solid var(--border-input);
    z-index: 1;
    transition: all 0.3s ease;
}
.progress-stepper .step.completed .step-dot {
    background: var(--accent);
    border-color: var(--accent);
}
.progress-stepper .step.active .step-dot {
    background: var(--accent);
    border-color: var(--accent);
    animation: stepPulse 1.5s ease-in-out infinite;
}
.progress-stepper .step-label {
    font-size: 11px;
    font-weight: 500;
    color: var(--text-secondary);
    text-align: center;
    margin-top: 6px;
    line-height: 1.3;
}
.progress-stepper .step.active .step-label {
    color: var(--accent);
    font-weight: 600;
}
.progress-stepper .step.completed .step-label {
    color: var(--text-primary);
}
.progress-stepper .step-time {
    font-size: 10px;
    color: var(--text-secondary);
    margin-top: 2px;
}
.progress-stepper .step-line {
    position: absolute;
    top: 7px;
    left: 50%;
    right: -50%;
    height: 2px;
    background: var(--border-input);
    z-index: 0;
}
.progress-stepper .step.completed .step-line {
    background: var(--accent);
}
.progress-stepper .step:last-child .step-line {
    display: none;
}
@keyframes stepPulse {
    0%, 100% { box-shadow: 0 0 0 0 rgba(0, 113, 227, 0.4); }
    50% { box-shadow: 0 0 0 6px rgba(0, 113, 227, 0); }
}
/* --- Helix spinner --- */
.helix-spinner {
    display: inline-block;
    width: 28px;
    height: 28px;
    position: relative;
}
.helix-spinner::before,
.helix-spinner::after {
    content: '';
    position: absolute;
    width: 100%;
    height: 100%;
    border-radius: 50%;
    border: 3px solid transparent;
    animation: helixSpin 1.2s linear infinite;
}
.helix-spinner::before {
    border-top-color: var(--accent);
    border-bottom-color: var(--accent);
}
.helix-spinner::after {
    border-left-color: var(--accent);
    border-right-color: var(--accent);
    animation-delay: -0.6s;
    opacity: 0.5;
}
@keyframes helixSpin {
    0% { transform: rotate(0deg); }
    100% { transform: rotate(360deg); }
}
/* --- Entrance animation --- */
@keyframes fadeSlideUp {
    from { opacity: 0; transform: translateY(12px); }
    to { opacity: 1; transform: translateY(0); }
}
/* --- Info panel (expansion inside card) --- */
.apple-info-panel .q-expansion-item__container {
    border: none;
    background: transparent;
}
.apple-info-panel .q-item {
    padding: 0;
    min-height: 28px;
}
.apple-info-panel .q-item__label {
    font-size: 12px;
    font-weight: 600;
    color: var(--text-secondary);
}
.apple-info-panel .q-expansion-item__content {
    font-size: 12px;
    color: var(--text-secondary);
    line-height: 1.5;
    padding: 4px 0 0 0;
}
/* --- Responsive --- */
@media (max-width: 980px) {
    .apple-sidebar {
        position: static;
        width: 100%;
        height: auto;
        border-right: none;
        border-bottom: 1px solid var(--border);
    }
    .apple-content {
        margin-left: 0;
        padding: 18px;
    }
}
@media (max-width: 700px) {
    .apple-form-grid {
        grid-template-columns: 1fr;
    }
    .apple-page-header h1 {
        font-size: 26px;
    }
}
"""

_NAV_GROUPS: list[tuple[str, list[tuple[str, str, str]]]] = [
    (
        "Workflows",
        [
            ("/", "BLASTer", "travel_explore"),
            ("/trim", "Trim", "content_cut"),
            ("/cluster", "pLM-clust v2", "hub"),
        ],
    )
]


class UploadBuffer:
    ALLOWED_SUFFIXES = {".fasta", ".fa", ".faa"}

    def __init__(self) -> None:
        self.files: List[Path] = []
        self.temp_dirs: List[Path] = []

    def clear(self) -> None:
        self.files.clear()
        for tmp in self.temp_dirs:
            try:
                resolved = tmp.resolve()
                tmp_root = Path(tempfile.gettempdir()).resolve()
                if tmp_root in resolved.parents and resolved.name.startswith("uht_gui_upload_"):
                    shutil.rmtree(resolved, ignore_errors=True)
            except Exception:
                pass
        self.temp_dirs.clear()

    def add_uploaded_file(self, name: str, content: bytes) -> Path:
        filename = Path(name).name
        suffix = Path(filename).suffix.lower()
        if suffix not in self.ALLOWED_SUFFIXES:
            raise ValueError(f"Unsupported file type: {suffix or 'unknown'}")
        tmp_dir = Path(tempfile.mkdtemp(prefix="uht_gui_upload_"))
        tmp_path = tmp_dir / filename
        with open(tmp_path, "wb") as f:
            f.write(content)
        self.temp_dirs.append(tmp_dir)
        self.files.append(tmp_path)
        return tmp_path

    def has_files(self) -> bool:
        return len(self.files) > 0


MAX_UPLOAD_FILE_SIZE = 20 * 1024 * 1024
MAX_UPLOAD_FILES = 20
_ACTIVE_RUN_LOCK = threading.Lock()


def _write_uploads(buffer: UploadBuffer, dest_dir: Path) -> List[Path]:
    dest_dir.mkdir(parents=True, exist_ok=True)
    written: List[Path] = []
    for source_path in buffer.files:
        dest_path = dest_dir / source_path.name
        shutil.copy2(source_path, dest_path)
        written.append(dest_path)
    return written

def _list_result_files(run_dir: Path, module: str) -> List[Path]:
    results_dir = run_dir / "results" / module
    if not results_dir.exists():
        return []
    return [p for p in results_dir.rglob("*") if p.is_file()]


def _relative_to_workdir(path: Path) -> str:
    base = get_workdir_base().resolve()
    try:
        rel = path.resolve().relative_to(base)
    except Exception:
        rel = path.name
    return f"/files/{rel.as_posix()}"


def _projection_files(results_root: Path, prefix: str) -> List[Path]:
    return sorted(results_root.rglob(f"{prefix}_coordinates_*.csv"))


def _projection_figure(df: pd.DataFrame, title: str) -> object:
    safe_df = df.copy()
    safe_df["cluster"] = safe_df["cluster"].astype(str)
    hover_cols = [col for col in ["header", "nll_score"] if col in safe_df.columns]
    fig = px.scatter(
        safe_df,
        x="Dim1",
        y="Dim2",
        color="cluster",
        hover_data=hover_cols,
        title=title,
        color_discrete_sequence=px.colors.qualitative.Bold,
    )
    fig.update_traces(marker={"size": 8, "opacity": 0.85})
    fig.update_layout(
        paper_bgcolor="rgba(0,0,0,0)",
        plot_bgcolor="rgba(0,0,0,0)",
        legend_title_text="Cluster",
        margin={"l": 16, "r": 16, "t": 50, "b": 16},
    )
    return fig


def _fasta_lengths(fasta_path: Path) -> List[int]:
    lengths: List[int] = []
    curr = 0
    with open(fasta_path, "r", encoding="utf-8", errors="ignore") as handle:
        for line in handle:
            line = line.strip()
            if not line:
                continue
            if line.startswith(">"):
                if curr > 0:
                    lengths.append(curr)
                curr = 0
            else:
                curr += len(line)
    if curr > 0:
        lengths.append(curr)
    return lengths


def _smart_bin_count(lengths: List[int]) -> int:
    if len(lengths) < 2:
        return 10
    arr = np.array(lengths, dtype=float)
    q75, q25 = np.percentile(arr, [75, 25])
    iqr = q75 - q25
    if iqr <= 0:
        return int(max(10, min(50, np.sqrt(len(arr)))))
    bin_width = 2 * iqr / np.cbrt(len(arr))
    if bin_width <= 0:
        return int(max(10, min(50, np.sqrt(len(arr)))))
    bins = int(np.ceil((arr.max() - arr.min()) / bin_width))
    return max(10, min(80, bins))


def _trim_histogram_figure(lengths: List[int], filename: str) -> object:
    arr = np.array(lengths, dtype=float)
    bins = _smart_bin_count(lengths)
    q05, q95 = np.percentile(arr, [5, 95])
    fig = go.Figure()
    fig.add_trace(
        go.Histogram(
            x=arr,
            nbinsx=bins,
            marker={"color": "#0071e3", "line": {"color": "rgba(0,0,0,0.15)", "width": 1}},
            opacity=0.9,
            name="Sequence length",
        )
    )
    fig.add_vline(x=float(q05), line={"color": "#f97316", "dash": "dash"}, annotation_text="Q5")
    fig.add_vline(x=float(q95), line={"color": "#16a34a", "dash": "dash"}, annotation_text="Q95")
    fig.update_layout(
        title=f"Length Distribution: {filename}",
        xaxis_title="Sequence length (aa)",
        yaxis_title="Count",
        paper_bgcolor="rgba(0,0,0,0)",
        plot_bgcolor="rgba(0,0,0,0)",
        bargap=0.03,
        margin={"l": 16, "r": 16, "t": 50, "b": 16},
    )
    return fig


def _render_trim_histogram(
    fasta_path: Path,
    container: ui.column,
    low_input: ui.number,
    high_input: ui.number,
    status: ui.label,
) -> None:
    lengths = _fasta_lengths(fasta_path)
    container.clear()
    if not lengths:
        with container:
            ui.label("No sequences found in uploaded FASTA.").classes("apple-status")
        status.set_text("Uploaded file contains no FASTA sequences.")
        return

    arr = np.array(lengths, dtype=float)
    q05, q95 = np.percentile(arr, [5, 95])
    low_input.value = int(round(q05))
    high_input.value = int(round(q95))

    with container:
        ui.label("Length Histogram (smart bins + suggested thresholds)").classes("apple-status")
        ui.plotly(_trim_histogram_figure(lengths, fasta_path.name)).classes("w-full apple-plot-wrap")

    status.set_text(
        f"Histogram ready for {fasta_path.name}. Suggested range: {int(round(q05))}-{int(round(q95))}."
    )


def _render_projection_section(results_root: Path) -> None:
    tsne_files = _projection_files(results_root, "tsne")
    umap_files = _projection_files(results_root, "umap")
    if not tsne_files and not umap_files:
        return

    ui.label("Interactive Projections").classes("apple-card-subtitle")

    with ui.tabs().classes("w-full") as plot_tabs:
        tsne_tab = ui.tab("t-SNE")
        umap_tab = ui.tab("UMAP")

    with ui.tab_panels(plot_tabs, value=tsne_tab).classes("w-full apple-plot-wrap"):
        with ui.tab_panel(tsne_tab):
            if not tsne_files:
                ui.label("No t-SNE coordinates available for this run.").classes("apple-status")
            else:
                for csv_path in tsne_files:
                    df = pd.read_csv(csv_path)
                    title = f"t-SNE: {csv_path.parent.name}"
                    ui.plotly(_projection_figure(df, title)).classes("w-full")

        with ui.tab_panel(umap_tab):
            if not umap_files:
                ui.label("No UMAP coordinates available for this run.").classes("apple-status")
            else:
                for csv_path in umap_files:
                    df = pd.read_csv(csv_path)
                    title = f"UMAP: {csv_path.parent.name}"
                    ui.plotly(_projection_figure(df, title)).classes("w-full")


def _render_results(container: ui.column, run_dir: Path, module: str) -> None:
    container.clear()
    results_root = run_dir / "results" / module
    files = _list_result_files(run_dir, module)
    with container:
        if not files:
            ui.label("No result files found.").classes("apple-status")
            return
        bundle_path = run_dir / f"{module}_results_bundle.zip"
        with zipfile.ZipFile(bundle_path, "w", zipfile.ZIP_DEFLATED) as zf:
            for file_path in files:
                zf.write(file_path, arcname=file_path.relative_to(run_dir))

        with ui.row().classes("items-center gap-3"):
            ui.label(f"Results ready ({len(files)} files bundled)").classes("apple-status")
            ui.button(
                "Download All Results (.zip)",
                on_click=lambda p=str(bundle_path): ui.download(p),
            ).classes("apple-btn").props("unelevated no-caps")
        if module == "plmclustv2":
            _render_projection_section(results_root)


def _ensure_workdir() -> Path:
    base = get_workdir_base()
    base.mkdir(parents=True, exist_ok=True)
    return base


@contextmanager
def apple_card(title: str, subtitle: str = ""):
    with ui.element("div").classes("apple-card"):
        if title:
            ui.label(title).classes("apple-card-title")
        if subtitle:
            ui.label(subtitle).classes("apple-card-subtitle")
        yield


def apple_input(label: str, value: str = "", placeholder: str = "") -> ui.input:
    el = ui.input(label=label, value=value, placeholder=placeholder)
    el.classes("apple-field w-full")
    el.props("outlined dense")
    return el


def apple_number(label: str, value: float, min_val: Optional[float] = None,
                 max_val: Optional[float] = None, step: float = 1.0) -> ui.number:
    kwargs = {"label": label, "value": value, "step": step}
    if min_val is not None:
        kwargs["min"] = min_val
    if max_val is not None:
        kwargs["max"] = max_val
    el = ui.number(**kwargs)
    el.classes("apple-field w-full")
    el.props("outlined dense")
    return el


def apple_select(label: str, options: List[str], value: str) -> ui.select:
    el = ui.select(options, value=value, label=label)
    el.classes("apple-field w-full")
    el.props("outlined dense")
    return el


def apple_upload(
    buffer: UploadBuffer,
    label: str,
    on_uploaded: Optional[Callable[[Path], None]] = None,
) -> ui.upload:
    async def _read_upload_bytes(source: object) -> Optional[bytes]:
        if source is None:
            return None
        if isinstance(source, (bytes, bytearray)):
            return bytes(source)
        read_fn = getattr(source, "read", None)
        if callable(read_fn):
            data = read_fn()
            if inspect.isawaitable(data):
                data = await data
            if isinstance(data, str):
                return data.encode("utf-8")
            if isinstance(data, (bytes, bytearray)):
                return bytes(data)
        return None

    def _extract_upload_parts(e: object) -> tuple[str, Optional[int], object]:
        file_obj = getattr(e, "file", None)
        if file_obj is not None:
            name = (
                getattr(file_obj, "name", None)
                or getattr(file_obj, "filename", None)
                or "upload.fasta"
            )
            size_attr = getattr(file_obj, "size", None)
            if callable(size_attr):
                size = int(size_attr())
            elif size_attr is not None:
                size = int(size_attr)
            else:
                size = None
            return str(name), size, file_obj

        name = getattr(e, "name", None) or getattr(e, "filename", None) or "upload.fasta"
        size_attr = getattr(e, "size", None)
        if callable(size_attr):
            size = int(size_attr())
        elif size_attr is not None:
            size = int(size_attr)
        else:
            size = None
        args_payload = getattr(e, "args", None)
        content = getattr(e, "content", None) or getattr(e, "data", None) or getattr(e, "bytes", None)
        if content is None and isinstance(args_payload, dict):
            if not getattr(e, "name", None) and not getattr(e, "filename", None):
                name = args_payload.get("name") or args_payload.get("filename") or name
            if size is None and args_payload.get("size") is not None:
                size = int(args_payload["size"])
            content = args_payload.get("content") or args_payload.get("data") or args_payload.get("bytes")
        return str(name), size, content

    async def handle_upload(e) -> None:
        try:
            file_name, size, source = _extract_upload_parts(e)
            content = await _read_upload_bytes(source)
            if content is None:
                ui.notify("Upload failed: unsupported upload event payload", color="negative")
                return
            if size is None:
                size = len(content)
            if size > MAX_UPLOAD_FILE_SIZE:
                size_mb = round(size / (1024 * 1024), 2)
                ui.notify(f"File too large: {size_mb} MB (max 20 MB)", color="negative")
                return
            uploaded_path = buffer.add_uploaded_file(file_name, content)
            if on_uploaded is not None:
                on_uploaded(uploaded_path)
            ui.notify(f"Uploaded: {Path(file_name).name}", color="positive")
        except ValueError as err:
            ui.notify(str(err), color="negative")
            ui.notify("Accepted extensions: .fasta, .fa, .faa", color="warning")

    upload = ui.upload(
        label=label,
        multiple=True,
        on_upload=handle_upload,
        auto_upload=True,
        max_file_size=MAX_UPLOAD_FILE_SIZE,
        max_files=MAX_UPLOAD_FILES,
        on_rejected=lambda: ui.notify(
            "Upload rejected. Use FASTA files only and keep each file under 20 MB.",
            color="negative",
        ),
    ).classes("apple-upload")
    upload.props("accept=.fasta,.fa,.faa")
    return upload


def apple_button(label: str, on_click):
    return ui.button(label, on_click=on_click).classes("apple-btn").props("unelevated no-caps")


def _run_blaster(buffer: UploadBuffer, num_hits: int, blast_db: str, evalue: float,
                 email: str, api_key: Optional[str], status: ui.label,
                 results_container: ui.column) -> None:
    from uht_discovery.core.BLASTer import run as blaster_run

    if not buffer.files:
        status.set_text("Please upload at least one FASTA file.")
        return
    run_id = generate_run_id()
    status.set_text(f"Run {run_id} started...")
    with RunIsolationContext(run_id) as run_ctx:
        _write_uploads(buffer, run_ctx.get_input_path("blaster"))
        cfg = {
            "blaster_project_directory": run_id,
            "blaster_num_hits": num_hits,
            "blaster_blast_db": blast_db,
            "blaster_evalue": evalue,
            "blaster_email": email,
            "blaster_api_key": api_key,
        }
        blaster_run(cfg)
        status.set_text(f"Run {run_id} completed.")
        _render_results(results_container, run_ctx.run_dir, "blaster")


def _run_trim(buffer: UploadBuffer, auto_mode: bool, low: Optional[int],
              high: Optional[int], status: ui.label,
              results_container: ui.column) -> None:
    from uht_discovery.core.trim import run as trim_run

    if not buffer.files:
        status.set_text("Please upload at least one FASTA file.")
        return
    run_id = generate_run_id()
    status.set_text(f"Run {run_id} started...")
    with RunIsolationContext(run_id) as run_ctx:
        _write_uploads(buffer, run_ctx.get_input_path("trim"))
        cfg = {
            "trim_project_directory": run_id,
            "auto_mode": auto_mode,
            "trim_low": low,
            "trim_high": high,
        }
        trim_run(cfg)
        status.set_text(f"Run {run_id} completed.")
        _render_results(results_container, run_ctx.run_dir, "trim")


def _run_plmclust(buffer: UploadBuffer, cluster_number: str, silhouette_min: int,
                  silhouette_max: int, keep_separate: bool,
                  status: ui.label, results_container: ui.column) -> None:
    from uht_discovery.core.plmclustv2 import run as plmclustv2_run

    if not buffer.files:
        status.set_text("Please upload at least one FASTA file.")
        return
    run_id = generate_run_id()
    status.set_text(f"Run {run_id} started...")
    with RunIsolationContext(run_id) as run_ctx:
        _write_uploads(buffer, run_ctx.get_input_path("plmclustv2"))
        cfg = {
            "plmclustv2_project_directory": run_id,
            "plmclustv2_cluster_number": cluster_number,
            "silhouette_range": [silhouette_min, silhouette_max],
            "plmclustv2_keepseparate": keep_separate,
        }
        plmclustv2_run(cfg)
        status.set_text(f"Run {run_id} completed.")
        _render_results(results_container, run_ctx.run_dir, "plmclustv2")


def _sync_data_theme() -> None:
    """Set HTML data-theme attribute to match current dark_mode storage value."""
    try:
        is_dark = app.storage.user.get("dark_mode", False)
    except RuntimeError:
        is_dark = False
    theme = "dark" if is_dark else "light"
    ui.run_javascript(f"document.documentElement.setAttribute('data-theme', '{theme}')")


def _build_sidebar(current_path: str) -> None:
    with ui.element("nav").classes("apple-sidebar"):
        ui.label("UHT Discovery").classes("sidebar-brand")
        ui.label("Protein workflow interface").classes("sidebar-subtitle")

        for group_label, items in _NAV_GROUPS:
            ui.label(group_label).classes("sidebar-section")
            for route, label, icon_name in items:
                active = " active" if current_path == route else ""
                with ui.link(target=route).classes(f"sidebar-item{active}").style("text-decoration: none;"):
                    ui.icon(icon_name).style("font-size: 18px;")
                    ui.label(label).style("margin-left: 10px;")

        ui.element("div").classes("sidebar-spacer")

        # Persist dark mode across page navigations via bind_value
        try:
            app.storage.user.setdefault("dark_mode", False)
        except RuntimeError:
            pass
        dark = ui.dark_mode().bind_value(app.storage.user, "dark_mode")

        def toggle_theme() -> None:
            try:
                app.storage.user["dark_mode"] = not app.storage.user.get("dark_mode", False)
            except RuntimeError:
                pass
            _sync_data_theme()

        # Apply data-theme on page load
        _sync_data_theme()

        with ui.element("div").classes("theme-toggle").on("click", toggle_theme):
            ui.icon("dark_mode").style("font-size: 18px;")
            ui.label("Toggle Dark Mode")


def _page_wrapper(current_path: str):
    ui.add_head_html(f"<style>{APPLE_CSS}</style>")
    _build_sidebar(current_path)
    return ui.element("main").classes("apple-content")


# ---------------------------------------------------------------------------
# Helper UI components
# ---------------------------------------------------------------------------

def _page_header(title: str, description: str) -> None:
    with ui.element("div").classes("apple-page-header"):
        ui.html(f"<h1>{title}</h1>")
        ui.html(f"<p>{description}</p>")


def _info_panel(title: str, lines: list[str]) -> None:
    with ui.element("div").classes("apple-info-panel"):
        with ui.expansion(title, icon="info").props("dense"):
            for line in lines:
                ui.html(f"<p style='margin:2px 0;'>{line}</p>")


def _section_divider() -> None:
    ui.html("<hr class='apple-section-divider'>")


def _help_text(text: str) -> None:
    ui.label(text).classes("apple-help-text")


def _progress_stepper(steps: list[tuple[str, str]], active_index: int) -> ui.element:
    container = ui.element("div").classes("progress-stepper")
    with container:
        for i, (label, time_est) in enumerate(steps):
            cls = "step"
            if i < active_index:
                cls += " completed"
            elif i == active_index:
                cls += " active"
            with ui.element("div").classes(cls):
                if i < len(steps) - 1:
                    ui.element("div").classes("step-line")
                ui.element("div").classes("step-dot")
                ui.label(label).classes("step-label")
                ui.label(time_est).classes("step-time")
    return container


# ---------------------------------------------------------------------------
# Step definitions for progress stepper
# ---------------------------------------------------------------------------

BLASTER_STEPS = [
    ("Uploading sequences", "~5s"),
    ("Searching NCBI", "varies"),
    ("Downloading hits", "1-5 min"),
    ("Combining results", "~10s"),
]

TRIM_STEPS = [
    ("Analyzing sequences", "~5s"),
    ("Applying filters", "~10s"),
    ("Writing output", "~5s"),
]

PLMCLUST_STEPS = [
    ("Preparing sequences", "~5s"),
    ("Computing ESM2 embeddings", "5-15 min"),
    ("Finding optimal clusters", "1-3 min"),
    ("Assigning clusters", "~30s"),
    ("Generating visualizations", "1-3 min"),
]


def _render_blaster_form() -> None:
    buffer = UploadBuffer()

    _page_header(
        "BLASTer",
        "Search NCBI databases for homologous sequences and download combined hits.",
    )

    # Setup card — upload + parameters
    with apple_card("Setup", "Upload query sequences and configure BLAST parameters."):
        _info_panel("About this step", [
            "<b>What it does:</b> Runs NCBI BLAST against a chosen database to find sequences similar to your query.",
            "<b>When to use:</b> First step in the discovery pipeline — start here with one or more query sequences.",
            "<b>Outputs:</b> combined_hits.fasta with all unique hits, plus per-query hit files.",
        ])
        _section_divider()
        ui.label("QUERY FILES").classes("apple-section-label")
        apple_upload(buffer, "Upload FASTA query files (.fasta, .fa, .faa)")
        _help_text("One or more FASTA files with your query protein sequences.")
        _section_divider()
        ui.label("PARAMETERS").classes("apple-section-label")
        with ui.element("div").classes("apple-form-grid"):
            num_hits = apple_number("Hits per query", value=200, min_val=1, step=1)
            blast_db = apple_select("BLAST DB", ["nr", "swissprot", "pdb", "refseq_protein"], "nr")
            evalue = apple_number("E-value", value=1e-5, step=1e-6)
            email = apple_input("NCBI email", placeholder="you@example.com")
            with ui.element("div").classes("full-width"):
                api_key = apple_input("NCBI API key (optional)")
        _help_text("An NCBI API key increases your rate limit from 3 to 10 requests/sec.")

    # Run card
    with apple_card("", ""):
        status = ui.label("Ready. Upload files and click Run.").classes("apple-status")
        stepper_container = ui.element("div")
        stepper_container.set_visibility(False)
        loading = ui.row().classes("apple-loading")
        with loading:
            ui.element("div").classes("helix-spinner")
            ui.label("Running BLASTer workflow...")
        loading.set_visibility(False)
        results_container = ui.column().classes("mt-2")

        async def run_job() -> None:
            acquired = _ACTIVE_RUN_LOCK.acquire(blocking=False)
            if not acquired:
                status.set_text("Another workflow is currently running. Please wait until it finishes.")
                ui.notify("Another workflow is running. Please wait.", color="warning")
                return
            status.set_text("Starting BLASTer...")
            stepper_container.set_visibility(True)
            stepper_container.clear()
            with stepper_container:
                _progress_stepper(BLASTER_STEPS, 0)
            loading.set_visibility(True)
            run_btn.disable()
            try:
                await run.io_bound(
                    _run_blaster,
                    buffer,
                    int(num_hits.value),
                    blast_db.value,
                    float(evalue.value),
                    email.value or "",
                    api_key.value or None,
                    status,
                    results_container,
                )
                stepper_container.clear()
                with stepper_container:
                    _progress_stepper(BLASTER_STEPS, len(BLASTER_STEPS))
                buffer.clear()
            except Exception as err:
                status.set_text(f"Run failed: {err}")
            finally:
                _ACTIVE_RUN_LOCK.release()
                loading.set_visibility(False)
                run_btn.enable()

        run_btn = apple_button("Run BLASTer", run_job)


def _render_trim_form() -> None:
    buffer = UploadBuffer()
    suggested = {"low": 100, "high": 500}

    _page_header(
        "Trim",
        "Filter sequences by length to remove outliers before clustering.",
    )

    # Setup card — upload + parameters + histogram
    with apple_card("Setup", "Upload sequences and configure trim thresholds."):
        _info_panel("About this step", [
            "<b>What it does:</b> Removes sequences outside a specified length range using statistical thresholds or manual bounds.",
            "<b>When to use:</b> After BLASTer — clean your hit set before embedding and clustering.",
            "<b>Outputs:</b> A trimmed FASTA containing only sequences within the chosen length window.",
        ])
        _section_divider()
        ui.label("INPUT FILES").classes("apple-section-label")
        hist_container = ui.column().classes("mt-2")

        def on_trim_uploaded(uploaded_path: Path) -> None:
            _render_trim_histogram(uploaded_path, hist_container, low, high, status_label)
            suggested["low"] = int(low.value)
            suggested["high"] = int(high.value)

        apple_upload(buffer, "Upload FASTA files (.fasta, .fa, .faa)", on_uploaded=on_trim_uploaded)
        _help_text("Upload the FASTA file(s) output from BLASTer or your own sequence set.")
        _section_divider()
        ui.label("THRESHOLDS").classes("apple-section-label")
        auto_mode = ui.checkbox("Auto-detect trim range (5th/95th percentile)", value=True)
        with ui.element("div").classes("apple-form-grid"):
            low = apple_number("Low threshold (aa)", value=100, step=1)
            high = apple_number("High threshold (aa)", value=500, step=1)

        with ui.row().classes("items-center gap-2 mt-1"):
            ui.button(
                "Use Suggested",
                on_click=lambda: (
                    setattr(low, "value", suggested["low"]),
                    setattr(high, "value", suggested["high"]),
                    status_label.set_text(f"Applied suggested range: {suggested['low']}-{suggested['high']}."),
                ),
            ).classes("apple-btn").props("unelevated no-caps size=sm")
            ui.button(
                "Reset Defaults",
                on_click=lambda: (
                    setattr(low, "value", 100),
                    setattr(high, "value", 500),
                    status_label.set_text("Reset to manual defaults: 100-500."),
                ),
            ).classes("apple-btn").props("unelevated no-caps size=sm")

        def _toggle_manual() -> None:
            low.set_enabled(not auto_mode.value)
            high.set_enabled(not auto_mode.value)

        auto_mode.on("update:model-value", lambda _: _toggle_manual())
        _toggle_manual()

    # Run card
    with apple_card("", ""):
        status_label = ui.label("Ready. Upload files and click Run.").classes("apple-status")
        stepper_container = ui.element("div")
        stepper_container.set_visibility(False)
        loading = ui.row().classes("apple-loading")
        with loading:
            ui.element("div").classes("helix-spinner")
            ui.label("Running Trim workflow...")
        loading.set_visibility(False)
        results_container = ui.column().classes("mt-2")

        async def run_job() -> None:
            acquired = _ACTIVE_RUN_LOCK.acquire(blocking=False)
            if not acquired:
                status_label.set_text("Another workflow is currently running. Please wait until it finishes.")
                ui.notify("Another workflow is running. Please wait.", color="warning")
                return
            status_label.set_text("Starting Trim...")
            stepper_container.set_visibility(True)
            stepper_container.clear()
            with stepper_container:
                _progress_stepper(TRIM_STEPS, 0)
            loading.set_visibility(True)
            run_btn.disable()
            try:
                manual_low = None if auto_mode.value else int(low.value)
                manual_high = None if auto_mode.value else int(high.value)
                await run.io_bound(
                    _run_trim,
                    buffer,
                    bool(auto_mode.value),
                    manual_low,
                    manual_high,
                    status_label,
                    results_container,
                )
                stepper_container.clear()
                with stepper_container:
                    _progress_stepper(TRIM_STEPS, len(TRIM_STEPS))
                buffer.clear()
            except Exception as err:
                status_label.set_text(f"Run failed: {err}")
            finally:
                _ACTIVE_RUN_LOCK.release()
                loading.set_visibility(False)
                run_btn.enable()

        run_btn = apple_button("Run Trim", run_job)


def _render_cluster_form() -> None:
    buffer = UploadBuffer()

    _page_header(
        "pLM-clust v2",
        "Cluster protein sequences using ESM2 language model embeddings.",
    )

    # Setup card — upload + parameters
    with apple_card("Setup", "Upload trimmed sequences and configure clustering."):
        _info_panel("About this step", [
            "<b>What it does:</b> Computes ESM2 embeddings for each sequence, then clusters via silhouette-optimised k-means.",
            "<b>When to use:</b> After trimming — use your cleaned FASTA to identify functional sub-groups.",
            "<b>Outputs:</b> Cluster assignments CSV, t-SNE/UMAP projections, and per-cluster FASTA files.",
        ])
        _section_divider()
        ui.label("INPUT FILES").classes("apple-section-label")
        apple_upload(buffer, "Upload FASTA files (.fasta, .fa, .faa)")
        _help_text("Upload the trimmed FASTA file(s) from the previous step.")
        _section_divider()
        ui.label("PARAMETERS").classes("apple-section-label")
        with ui.element("div").classes("apple-form-grid"):
            cluster_number = apple_input("Clusters (number or 'auto')", value="auto")
            keep_separate = ui.checkbox("Keep files separate", value=False)
            silhouette_min = apple_number("Silhouette min k", value=2, step=1)
            silhouette_max = apple_number("Silhouette max k", value=10, step=1)
        _help_text("Use 'auto' to let the algorithm find the optimal cluster count via silhouette analysis.")

    # Run card
    with apple_card("", ""):
        status = ui.label("Ready. Upload files and click Run.").classes("apple-status")
        stepper_container = ui.element("div")
        stepper_container.set_visibility(False)
        loading = ui.row().classes("apple-loading")
        with loading:
            ui.element("div").classes("helix-spinner")
            ui.label("Running pLM-clust v2 workflow...")
        loading.set_visibility(False)
        results_container = ui.column().classes("mt-2")

        async def run_job() -> None:
            acquired = _ACTIVE_RUN_LOCK.acquire(blocking=False)
            if not acquired:
                status.set_text("Another workflow is currently running. Please wait until it finishes.")
                ui.notify("Another workflow is running. Please wait.", color="warning")
                return
            status.set_text("Starting pLM-clust v2...")
            stepper_container.set_visibility(True)
            stepper_container.clear()
            with stepper_container:
                _progress_stepper(PLMCLUST_STEPS, 0)
            loading.set_visibility(True)
            run_btn.disable()
            try:
                await run.io_bound(
                    _run_plmclust,
                    buffer,
                    str(cluster_number.value),
                    int(silhouette_min.value),
                    int(silhouette_max.value),
                    bool(keep_separate.value),
                    status,
                    results_container,
                )
                stepper_container.clear()
                with stepper_container:
                    _progress_stepper(PLMCLUST_STEPS, len(PLMCLUST_STEPS))
                buffer.clear()
            except Exception as err:
                status.set_text(f"Run failed: {err}")
            finally:
                _ACTIVE_RUN_LOCK.release()
                loading.set_visibility(False)
                run_btn.enable()

        run_btn = apple_button("Run pLM-clust v2", run_job)


def register_routes() -> None:
    @ui.page("/")
    def page_blast() -> None:
        with _page_wrapper("/"):
            with ui.element("div").classes("apple-content-inner"):
                _render_blaster_form()

    @ui.page("/trim")
    def page_trim() -> None:
        with _page_wrapper("/trim"):
            with ui.element("div").classes("apple-content-inner"):
                _render_trim_form()

    @ui.page("/cluster")
    def page_cluster() -> None:
        with _page_wrapper("/cluster"):
            with ui.element("div").classes("apple-content-inner"):
                _render_cluster_form()


register_routes()


def create_fastapi_app() -> FastAPI:
    base = _ensure_workdir()
    fastapi_app = FastAPI()
    if os.getenv("NICEGUI_DISABLE_PROCESS_POOL") == "1":
        run.setup = lambda: None
    ui.run_with(app=fastapi_app, storage_secret=os.getenv("UHT_STORAGE_SECRET", "uht-discovery"))
    fastapi_app.mount("/files", StaticFiles(directory=str(base)), name="files")
    return fastapi_app


def main() -> None:
    base = _ensure_workdir()
    app.add_static_files("/files", str(base))
    if os.getenv("NICEGUI_DISABLE_PROCESS_POOL") == "1":
        run.setup = lambda: None
    ui.run(
        host=os.getenv("UHT_HOST", "0.0.0.0"),
        port=int(os.getenv("UHT_PORT", "7860")),
        reload=False,
        show=False,
        storage_secret=os.getenv("UHT_STORAGE_SECRET", "uht-discovery"),
    )


if __name__ == "__main__":
    main()
