src/streamlit_app.py

"""Streamlit phase explorer (Deployment 2).

Runs the pipeline phases that exist today and renders EACH phase's INPUT and OUTPUT — for every phase the
filters/config that drive it, for the embedders the exact query the candidate text is compared against, plus
row-count, timing, score spread and a preview table — so you can watch the funnel stage by stage: preprocess
(load → validate → materialize → quality) then rank (hard_filter → potion → soft → composite → granite →
cross-encoder → ensemble → reasoning). A per-phase markdown/JSON report can be exported. HF Space entrypoint
(Docker, 8501).

Run locally:  uv run streamlit run src/streamlit_app.py
"""

from __future__ import annotations

import json
import time
from collections.abc import Callable
from dataclasses import dataclass, field
from typing import Any

import orjson
import polars as pl
import streamlit as st

from common.checkpoint import checkpoint
from common.io import atomic_write_bytes
from common.logging import get_logger
from common.paths import DEFAULT_CANDIDATES, REPO_ROOT, resolve_candidates
from common.runtime import configure_compute
from config import DebugSettings, Settings, load_settings
from jd import JDSpec, embed_query, extract_spec, read_subjective
from jd.parse import _OBJECTIVE_PROMPT, _SUBJECTIVE_PROMPT, _SYSTEM, extract_objective, extract_subjective
from jd.spec import resolve
from models.llm import index_gen_client
from preprocess import DATASET_AS_OF, load_candidates, preprocess_stages
from rank import add_explanation, bm25_prefilter, composite_narrow, ensemble, hard_filter, soft_filter, write_submission
from rank.composite import (  # reused so the explorer's filter detail matches the engine exactly
    _BM25_FIELDS,
    _drop_disallowed_countries,
    field_scoped_bm25,
)
from rank.hard_filter import _toggle_masks              # reused so per-toggle drop counts match the engine exactly
from retrieve import cross_encoder_rerank, granite_precision_score, potion_broad_retrieval


@dataclass
class PhaseOutput:
    """The captured result of running a phase — frame plus the inputs that drove it and the engine metrics."""

    name: str
    frame: pl.DataFrame
    rows_in: int
    seconds: float
    inputs: dict[str, Any] | None = None       # filters/config + embedding query that drove this phase
    metrics: dict[str, Any] = field(default_factory=dict)  # captured step.end metrics (drop counts, score_max…)
    is_funnel: bool = True                      # False for the JD-parse phases (not part of the candidate funnel)


def _display_frame(frame: pl.DataFrame, n: int = 50) -> pl.DataFrame:
    """ALL columns, renderable: scalars kept as-is; nested list/struct columns stringified (truncated).

    Streamlit can't render polars List/Struct columns directly (and pandas isn't a dependency), so nested
    columns are converted to a short text repr on the head — letting every column show with values.
    """
    head = frame.head(n)
    projected: list[pl.Expr] = []
    for name, dtype in head.schema.items():
        if isinstance(dtype, (pl.List, pl.Struct)):
            # list cells arrive as a polars Series → .to_list() for a clean python repr; struct cells are dicts.
            projected.append(
                pl.col(name).map_elements(
                    lambda value: str(value.to_list() if hasattr(value, "to_list") else value)[:200],
                    return_dtype=pl.String,
                ).alias(name)
            )
        else:
            projected.append(pl.col(name))
    return head.select(projected)


@st.cache_resource(show_spinner=False)
def _load_phase_input(path: str, limit: int) -> pl.DataFrame:
    """Load candidates (cached by path+limit so re-runs are instant).

    ``@st.cache_resource`` (not ``@st.cache_data``) — returns the SAME in-memory object on every cache hit
    instead of a fresh copy, avoiding a real copy cost on a 100K-row/119-column frame. Safe because every
    downstream stage (validate/materialize/quality/hard_filter/...) is a Polars transform that returns a NEW
    frame rather than mutating its input in place, so the cached original is never altered by a caller.
    """
    return load_candidates(path, None if limit == 0 else limit)


def _rank_stages(settings: Settings) -> list[tuple[str, Callable[[pl.DataFrame], pl.DataFrame]]]:
    """The rank-phase stages (frame→frame closures), reading the JD spec from the generated extraction."""
    spec = extract_spec(settings)
    ideal_summary = read_subjective(resolve(settings.jd.extraction_path)).get("ideal_summary", "")
    query = embed_query(spec, ideal_summary)  # ideal_summary + extracted skills/roles → the semantic query
    return [
        ("hard_filter", lambda frame: hard_filter(frame, settings)),
        ("bm25_prefilter", lambda frame: bm25_prefilter(frame, spec, settings)),
        ("potion_broad", lambda frame: potion_broad_retrieval(frame, query, settings)),
        ("soft_filter", lambda frame: soft_filter(frame, spec, settings)),
        ("composite_narrow", lambda frame: composite_narrow(frame, spec, settings)),
        ("granite_precision", lambda frame: granite_precision_score(frame, query, settings)),
        ("cross_encoder", lambda frame: cross_encoder_rerank(frame, query, settings)),
        ("ensemble", lambda frame: ensemble(frame, settings)),
        ("explain", lambda frame: add_explanation(frame, spec, settings)),
        # Final stage: write submission.csv (4-col contract) so EVERY pipeline run in the explorer refreshes the
        # deliverable, not just the headless CLI. Returns the written submission frame (shown as this phase's output).
        ("submission", lambda frame: write_submission(frame, settings)),
    ]


def _phase_inputs(name: str, settings: Settings, spec: JDSpec, ideal_summary: str) -> dict[str, Any] | None:
    """Describe what DRIVES a phase: its filter/config block, JD-derived signals, and (embedders) the query.

    Keyed by the phase name the explorer uses (``rank.<stage>`` / ``preprocess.<stage>``). The embedding
    phases all compare candidate ``text`` against ``ideal_summary`` — surfaced here as ``embed_query``.
    """
    years = {"years_min": spec.years_min, "years_max": spec.years_max}
    if name == "rank.hard_filter":
        return {"label": "Hard-filter toggles (candidate-only drops)", "config": settings.hard_filter.model_dump()}
    if name == "rank.bm25_prefilter":
        bm25_fields = {
            f"{group} ({'+'.join(columns)})": [term for attr in attrs for term in (getattr(spec, attr, []) or [])]
            for group, attrs, columns, _weight in _BM25_FIELDS
        }
        return {
            "label": "BM25 pre-filter (location drop + cheap lexical cut, BEFORE potion/soft exist)",
            "config": settings.bm25_prefilter.model_dump(),
            "jd_signals": {"acceptable_countries (location DROP; [] = accept all)": spec.acceptable_countries,
                           "bm25_query_terms_by_field": bm25_fields},
        }
    if name == "rank.potion_broad":
        return {
            "label": "Broad retrieval (static potion embedder)", "config": settings.broad_retrieval.model_dump(),
            "model": settings.models.embedder_broad, "embed_query": embed_query(spec, ideal_summary), "score_column": "potion_score",
        }
    if name == "rank.soft_filter":
        return {
            "label": "Soft-filter penalty weights", "config": settings.soft_filter.model_dump(),
            "jd_signals": {"target_industries": spec.target_industries,
                           "exclude_companies": spec.exclude_companies, **years},
            "score_column": "soft_score",
        }
    if name == "rank.composite_narrow":
        # Field-scoped BM25: each group's query terms come from a distinct JD bag, matched vs a distinct text section.
        bm25_fields = {
            f"{group} ({'+'.join(columns)})": [term for attr in attrs for term in (getattr(spec, attr, []) or [])]
            for group, attrs, columns, _weight in _BM25_FIELDS
        }
        return {
            "label": "Composite narrow (location drop + field-scoped BM25 + z-blend)",
            "config": settings.composite_narrow.model_dump(),
            "jd_signals": {"acceptable_countries (location DROP; [] = accept all)": spec.acceptable_countries,
                           "bm25_query_terms_by_field": bm25_fields},
            "score_column": "composite_score",
        }
    if name == "rank.granite_precision":
        return {
            "label": "Granite precision embedding (ONNX bi-encoder)", "config": settings.precision.model_dump(),
            "model": settings.models.embedder_precision, "embed_query": embed_query(spec, ideal_summary), "score_column": "granite_score",
        }
    if name == "rank.cross_encoder":
        return {
            "label": "Cross-encoder rerank (query+doc together)", "config": settings.rerank.model_dump(),
            "model": settings.models.reranker, "embed_query": embed_query(spec, ideal_summary), "score_column": "reranker_score",
        }
    if name == "rank.ensemble":
        return {"label": "Ensemble fusion weights", "config": settings.ensemble.model_dump(), "score_column": "final_score"}
    if name == "rank.explain":
        r = settings.reasoning
        panel: dict[str, Any] = {
            "label": "Multi-band explainability (JD-grounded why per finalist + artifact)",
            "config": r.model_dump(),
            "tiers": {
                f"T1 ranks 1-{r.tier1_top}": "Qwen3-1.7B polish (reused, only if mode=llm) · 4 sentences · no gap",
                f"T2 ranks {r.tier1_top + 1}-{r.tier2_top}": "pure NLG · 3-4 sentences · no gap",
                f"T3 ranks {r.tier2_top + 1}-{r.tier3_top}": "pure NLG · 3-4 sentences · ≤1 soft gap",
                f"T4 ranks {r.tier3_top + 1}-100": "pure NLG (static) · 3-4 sentences · ≤1 soft gap",
            },
            "variant_selection": f"{r.variant_select} (top-{r.variant_topk})",
        }
        run_summary = REPO_ROOT / "data/reports/run_summary.json"
        if run_summary.exists():  # surface the last run's aggregate (tier/signal/gap distributions)
            panel["last_run"] = orjson.loads(run_summary.read_bytes())
        return panel
    if name == "preprocess.quality":
        return {"label": "Quality pillar weights (candidate_quality blend)", "config": settings.quality.model_dump()}
    return None


def _stage_metrics(
    name: str, frame_in: pl.DataFrame, frame_out: pl.DataFrame, settings: Settings, spec: JDSpec
) -> dict[str, Any]:
    """Per-phase drop counts, computed by reusing the engine's own filter helpers (so they match exactly).

    The hard filter reports each enabled toggle's independent drop count; the composite reports how many
    rows the JD location filter removed and how many got a non-zero BM25 hit. Other phases add nothing here
    (their score spread is shown separately from the output frame).
    """
    if name == "rank.hard_filter":
        masks = _toggle_masks(frame_in, settings.hard_filter)
        counts = (
            frame_in.select([mask.alias(toggle) for toggle, mask in masks.items()]).sum().row(0, named=True)
            if masks else {}
        )
        return {**{toggle: int(count or 0) for toggle, count in counts.items()},
                "dropped (any toggle)": frame_in.height - frame_out.height}
    if name == "rank.bm25_prefilter":
        location_kept = _drop_disallowed_countries(frame_in, spec)
        return {"location_dropped": frame_in.height - location_kept.height}
    if name == "rank.composite_narrow":
        location_kept = _drop_disallowed_countries(frame_in, spec)
        metrics: dict[str, Any] = {"location_dropped": frame_in.height - location_kept.height}
        # Field-scoped lexical hits over the survivors (reuses the engine helper; ≤ top_n docs, so cheap).
        _blended, hits, _per_field = field_scoped_bm25(frame_out, spec, settings.composite_narrow)
        metrics["bm25_matched (survivors, ≥1 field)"] = int((hits > 0).sum())
        return metrics
    return {}


def _jd_kv_frame(data: dict[str, Any]) -> pl.DataFrame:
    """A 2-column field/value frame of an extracted JD dict — so a JSON phase renders like the others."""
    if not data:
        return pl.DataFrame({"field": [], "value": []}, schema={"field": pl.String, "value": pl.String})
    return pl.DataFrame({"field": [str(k) for k in data], "value": [str(v) for v in data.values()]})


def _jd_phases(settings: Settings, regenerate: bool, on_stage_start: Callable[[str], None] | None = None):
    """YIELDS the two JD-parse LLM phases (objective keyword spec + subjective ideal_summary) as soon as EACH
    completes — not batched at the end, so a slow objective call (e.g. weak CPU hardware) is visible on its
    own instead of the UI going quiet until both are done. ``on_stage_start``, if given, fires right BEFORE
    each (possibly slow) call so the caller can show "currently running: X" while it's still in flight.

    ``regenerate`` runs the live Qwen3 calls (each timed independently) and atomically rewrites
    ``jd_extraction.json``; otherwise the existing artifact is read (a near-zero file read, still shown as a
    phase). Outputs are non-funnel — their row counts are field counts, not candidate counts.

    ``regenerate=False`` with NO existing artifact is escalated to a live parse anyway: reading a missing file
    silently yields an EMPTY JDSpec, which drives the entire funnel with no filters and no query — the whole
    run "succeeds" but ranks on nothing. That's the actual root cause behind repeated
    ``jd.extract.missing ... empty spec (defaults)`` warnings on a fresh deployment where the artifact was
    never generated. "Don't force a re-parse when a good cache exists" must never mean "run with nothing."
    """
    extraction_path = resolve(settings.jd.extraction_path)
    jd_path = resolve(settings.jd.path)
    jd_text = jd_path.read_text(encoding="utf-8") if jd_path.exists() else ""
    model = settings.models.index_gen_llm_repo
    obj_prompt = _OBJECTIVE_PROMPT.replace("{jd_text}", jd_text or "(JD source missing)")
    sum_prompt = _SUBJECTIVE_PROMPT.replace("{jd_text}", jd_text or "(JD source missing)")

    if not regenerate and not extraction_path.exists() and jd_text.strip():
        get_logger("streamlit").warning(
            "streamlit.jd_extraction_missing_forcing_live_parse", path=str(extraction_path),
        )
        regenerate = True

    if regenerate and jd_text.strip():
        client = index_gen_client(settings)
        if on_stage_start:
            on_stage_start("jd.objective_extract")
        started = time.perf_counter()
        objective = extract_objective(jd_text, client) or {}
        obj_seconds = time.perf_counter() - started
        source = "live Qwen3 parse"
        yield PhaseOutput(
            "jd.objective_extract", _jd_kv_frame(objective), 0, obj_seconds,
            inputs={"label": f"JD objective extraction — {source}", "model": model, "llm_system": _SYSTEM,
                    "llm_prompt": obj_prompt},
            metrics={"fields_extracted": len(objective)}, is_funnel=False,
        )
        if on_stage_start:
            on_stage_start("jd.summary_extract")
        started = time.perf_counter()
        subjective = extract_subjective(jd_text, client) or {}
        sum_seconds = time.perf_counter() - started
        # Persist the measured parse time alongside the spec so cached runs can replay the ORIGINAL timing.
        # ``_timing`` is inert to the engine (extract_spec/read_subjective only read the keys they know).
        atomic_write_bytes(
            extraction_path,
            orjson.dumps(
                {**objective, "subjective": subjective,
                 "_timing": {"objective_seconds": obj_seconds, "summary_seconds": sum_seconds}},
                option=orjson.OPT_INDENT_2,
            ),
        )
        yield PhaseOutput(
            "jd.summary_extract", _jd_kv_frame(subjective), 0, sum_seconds,
            inputs={"label": f"JD summary extraction — {source}", "model": model, "llm_system": _SYSTEM,
                    "llm_prompt": sum_prompt},
            metrics={"has_ideal_summary": bool(subjective.get("ideal_summary"))}, is_funnel=False,
        )
    else:  # cached read — replay the original parse time recorded in the artifact (else the bare file-read time)
        started = time.perf_counter()
        raw = orjson.loads(extraction_path.read_bytes()) if extraction_path.exists() else {}
        read_seconds = time.perf_counter() - started
        raw = raw if isinstance(raw, dict) else {}
        subjective = raw.pop("subjective", {}) or {}
        timing = raw.pop("_timing", {}) or {}
        objective = raw
        obj_seconds = float(timing.get("objective_seconds", read_seconds))
        sum_seconds = float(timing.get("summary_seconds", 0.0))
        source = ("cached jd_extraction.json (original parse time)" if timing
                  else "cached jd_extraction.json (no recorded time — tick Regenerate to time the live parse)")
        yield PhaseOutput(
            "jd.objective_extract", _jd_kv_frame(objective), 0, obj_seconds,
            inputs={"label": f"JD objective extraction — {source}", "model": model, "llm_system": _SYSTEM,
                    "llm_prompt": obj_prompt},
            metrics={"fields_extracted": len(objective)}, is_funnel=False,
        )
        yield PhaseOutput(
            "jd.summary_extract", _jd_kv_frame(subjective), 0, sum_seconds,
            inputs={"label": f"JD summary extraction — {source}", "model": model, "llm_system": _SYSTEM,
                    "llm_prompt": sum_prompt},
            metrics={"has_ideal_summary": bool(subjective.get("ideal_summary"))}, is_funnel=False,
        )


def _run_pipeline(
    path: str, limit: int, save_checkpoints: bool, regenerate_jd: bool,
    on_stage_start: Callable[[str], None] | None = None,
):
    """YIELDS each phase's ``PhaseOutput`` as soon as it finishes — the JD-parse phases, then preprocess
    sub-stages, then the rank stages. Never collects everything before returning: on slow hardware a single
    phase (JD parse, granite, the LLM polish) can take minutes, and the caller should be able to render/inspect
    each phase's own time the moment it's known, not wait for the whole funnel to finish first. ``on_stage_start``,
    if given, fires right BEFORE each stage runs — so the caller can show "currently running: X" while a slow
    stage is still in flight, instead of only finding out after it finally completes.

    Stage sequences come from the same lists the engine folds (``preprocess_stages`` + ``_rank_stages``),
    so the explorer mirrors the real pipeline. ``save_checkpoints`` also writes each phase to Parquet. Drop
    counts are recomputed here via the engine's own filter helpers (``_toggle_masks`` /
    ``_drop_disallowed_countries``) so they're single-sourced rather than scraped from logs.
    """
    settings = load_settings()
    if save_checkpoints:  # override the config toggle for this run
        settings = settings.model_copy(
            update={"debug": DebugSettings(checkpoints=True, checkpoint_dir=settings.debug.checkpoint_dir)}
        )
    yield from _jd_phases(settings, regenerate_jd, on_stage_start)  # JD parse first (writes the artifact if regen)
    spec = extract_spec(settings)  # reads the (now fresh, if regen) artifact the JD phases just wrote
    ideal_summary = read_subjective(resolve(settings.jd.extraction_path)).get("ideal_summary", "")

    if on_stage_start:
        on_stage_start("preprocess.load")
    started = time.perf_counter()
    current = _load_phase_input(path, limit)
    yield PhaseOutput("preprocess.load", current, current.height, time.perf_counter() - started,
                       inputs=_phase_inputs("preprocess.load", settings, spec, ideal_summary))
    checkpoint(0, "load", current, settings)

    stages: list[tuple[str, Callable[[pl.DataFrame], pl.DataFrame]]] = [
        *((f"preprocess.{s.name}", s.run) for s in preprocess_stages(DATASET_AS_OF)),
        *((f"rank.{name}", fn) for name, fn in _rank_stages(settings)),
    ]
    for index, (name, run) in enumerate(stages, start=1):
        if on_stage_start:
            on_stage_start(name)
        rows_in = current.height
        frame_in = current
        stage_started = time.perf_counter()
        current = run(current)
        inputs = _phase_inputs(name, settings, spec, ideal_summary)
        metrics = _stage_metrics(name, frame_in, current, settings, spec)
        score_column = (inputs or {}).get("score_column")  # surface the score maximum alongside the drop counts
        if score_column and score_column in current.columns and current.height \
                and current.schema[score_column].is_numeric():
            metrics[f"{score_column}_max"] = round(float(current.get_column(score_column).max()), 4)
        yield PhaseOutput(
            name, current, rows_in, time.perf_counter() - stage_started, inputs=inputs, metrics=metrics,
        )
        checkpoint(index, name.split(".", 1)[-1], current, settings)


def _score_stats(frame: pl.DataFrame, column: str | None) -> dict[str, float] | None:
    """min / mean / max of a numeric score column, if it's present and numeric (else None)."""
    if not column or column not in frame.columns or frame.height == 0:
        return None
    if not frame.schema[column].is_numeric():
        return None
    row = frame.select(
        pl.col(column).min().alias("min"), pl.col(column).mean().alias("mean"), pl.col(column).max().alias("max")
    ).row(0, named=True)
    return {key: round(float(value), 4) for key, value in row.items() if value is not None}


def _render_inputs(output: PhaseOutput) -> None:
    """Render the INPUTS panel: filters/config, JD signals, the embedding-compare query, and engine metrics."""
    inputs = output.inputs
    st.markdown("**Inputs — filters / config that drove this phase**")
    if inputs is None:
        st.caption("No tunable filters — this phase transforms every row (no config or query).")
    else:
        if inputs.get("config"):
            st.caption(inputs.get("label", "config"))
            st.json(inputs["config"])
        if inputs.get("jd_signals"):
            st.caption("JD-derived signals (from jd_extraction.json)")
            st.json(inputs["jd_signals"])
        if inputs.get("model"):
            st.caption(f"model: `{inputs['model']}`")
        if inputs.get("llm_system"):
            st.caption("LLM system prompt (persona — pins Qwen3 to JSON, no thinking loop)")
            st.code(inputs["llm_system"], language="text")
        if inputs.get("llm_prompt"):
            st.markdown("**LLM prompt** — the structured Role-Task-Context-Constraints-Output prompt sent:")
            st.text_area(
                "llm_prompt", inputs["llm_prompt"], height=260, disabled=True,
                label_visibility="collapsed", key=f"p_{output.name}",
            )
        if "embed_query" in inputs:
            st.markdown("**Embedding-compare query** — candidate `text` is scored against this:")
            st.text_area(
                "embed_query", inputs["embed_query"] or "(none — phase passes through at score 0)",
                height=110, disabled=True, label_visibility="collapsed", key=f"q_{output.name}",
            )
    stats = _score_stats(output.frame, (inputs or {}).get("score_column"))
    if stats:
        st.caption(f"`{inputs['score_column']}` spread over the output")
        cols = st.columns(len(stats))
        for col, (label, value) in zip(cols, stats.items()):
            col.metric(label, f"{value:.4f}")
    if output.metrics:
        st.caption("engine metrics (per-rule drop counts, score maxima)")
        st.dataframe(
            pl.DataFrame({"metric": list(output.metrics.keys()),
                          "value": [str(v) for v in output.metrics.values()]}),
            width="stretch", hide_index=True,
        )


def _render_phase(output: PhaseOutput) -> None:
    """Render one phase: metrics row, the INPUTS panel, then the OUTPUT preview, inside an expander."""
    with st.expander(f"Phase — {output.name}", expanded=True):
        left, mid, right = st.columns(3)
        if output.is_funnel:
            left.metric("rows out", f"{output.frame.height:,}", delta=output.frame.height - output.rows_in)
        else:  # JD-parse phase — its frame is the extracted key/value pairs, not candidate rows
            left.metric("fields out", f"{output.frame.height:,}")
        mid.metric("columns", output.frame.width)
        right.metric("seconds", f"{output.seconds:.2f}")
        in_col, out_col = st.columns([2, 3])
        with in_col:
            _render_inputs(output)
        with out_col:
            out_label = "rows after this phase" if output.is_funnel else "extracted JD fields (key / value)"
            st.markdown(f"**Output — {out_label}**")
            st.caption(f"all {output.frame.width} columns (first 50 rows; nested columns shown as text)")
            st.dataframe(_display_frame(output.frame), width="stretch")
            with st.popover("Column dtypes"):
                schema = pl.DataFrame(
                    {"column": list(output.frame.columns),
                     "dtype": [str(dtype) for dtype in output.frame.dtypes]}
                )
                st.dataframe(schema, width="stretch", height=400)


def _render_timing_summary(outputs: list[PhaseOutput]) -> None:
    """Top-of-page overview: total wall-clock (incl. JD parse), the candidate funnel, and a per-phase table."""
    total_seconds = sum(output.seconds for output in outputs)  # includes the JD-parse phases
    slowest = max(outputs, key=lambda output: output.seconds)
    funnel = [output for output in outputs if output.is_funnel] or outputs  # JD phases excluded from the funnel
    left, mid, right = st.columns(3)
    left.metric("Total time", f"{total_seconds:.1f}s")
    mid.metric("Funnel", f"{funnel[0].frame.height:,} → {funnel[-1].frame.height:,}")
    right.metric("Slowest phase", f"{slowest.name.split('.')[-1]} ({slowest.seconds:.1f}s)")
    overview = pl.DataFrame({
        "phase": [output.name for output in outputs],
        "rows_in": [output.rows_in for output in outputs],
        "rows_out": [output.frame.height for output in outputs],
        "seconds": [round(output.seconds, 2) for output in outputs],
        "% of total": [round(100 * output.seconds / total_seconds, 1) if total_seconds else 0.0 for output in outputs],
    })
    st.dataframe(overview, width="stretch")


def _render_jd_panel(spec: JDSpec, ideal_summary: str) -> None:
    """Show the JD spec that drives the funnel (loaded from the generated jd_extraction.json)."""
    with st.expander("JD spec (from jd_extraction.json — drives filters + retrieval)", expanded=False):
        st.json({
            "job_titles": spec.job_titles, "must_have_skills": spec.must_have_skills,
            "nice_to_have": spec.nice_to_have, "target_industries": spec.target_industries,
            "years": f"{spec.years_min:.0f}–{spec.years_max:.0f}",
        })
        st.text_area("ideal_summary (potion embedding query)", ideal_summary or "(none)", height=120, disabled=True)


def _scalar_preview(frame: pl.DataFrame, n: int = 5) -> list[dict[str, Any]]:
    """First ``n`` rows over the scalar columns only (nested columns dropped) — for the exported report."""
    scalar = [name for name, dtype in frame.schema.items() if not isinstance(dtype, (pl.List, pl.Struct))]
    return frame.select(scalar).head(n).to_dicts() if scalar else []


def _report_payload(outputs: list[PhaseOutput], spec: JDSpec, ideal_summary: str) -> dict[str, Any]:
    """Assemble the exportable report — JD spec plus each phase's inputs, outputs, metrics and a small preview."""
    return {
        "jd_spec": {
            "job_titles": spec.job_titles, "must_have_skills": spec.must_have_skills,
            "nice_to_have": spec.nice_to_have, "target_industries": spec.target_industries,
            "exclude_companies": spec.exclude_companies, "acceptable_countries": spec.acceptable_countries,
            "years_min": spec.years_min, "years_max": spec.years_max,
        },
        "ideal_summary": ideal_summary,
        "phases": [
            {
                "name": output.name, "rows_in": output.rows_in, "rows_out": output.frame.height,
                "delta": output.frame.height - output.rows_in, "seconds": round(output.seconds, 3),
                "columns": output.frame.width, "inputs": output.inputs, "metrics": output.metrics,
                "score_stats": _score_stats(output.frame, (output.inputs or {}).get("score_column")),
                "preview": _scalar_preview(output.frame),
            }
            for output in outputs
        ],
    }


def _report_markdown(payload: dict[str, Any]) -> str:
    """Render the report payload as readable markdown (one section per phase: input → output)."""
    lines = ["# Reranker — per-phase input/output report", "", "## JD spec", "",
             "```json", json.dumps(payload["jd_spec"], indent=2, ensure_ascii=False), "```", "",
             "**ideal_summary (embedding-compare query):**", "", f"> {payload['ideal_summary'] or '(none)'}", "",
             "## Phases", ""]
    for phase in payload["phases"]:
        lines.append(f"### {phase['name']}  ({phase['rows_in']:,} → {phase['rows_out']:,}, "
                     f"Δ{phase['delta']:+,}, {phase['seconds']:.2f}s)")
        inputs = phase["inputs"]
        if inputs:
            if inputs.get("config"):
                lines += ["", f"_{inputs.get('label', 'config')}:_", "",
                          "```json", json.dumps(inputs["config"], indent=2, ensure_ascii=False), "```"]
            if inputs.get("jd_signals"):
                lines += ["", "_JD signals:_", "",
                          "```json", json.dumps(inputs["jd_signals"], indent=2, ensure_ascii=False), "```"]
            if inputs.get("model"):
                lines += ["", f"_model:_ `{inputs['model']}`"]
            if inputs.get("llm_system"):
                lines += ["", "_LLM system prompt:_", "", "```text", inputs["llm_system"], "```"]
            if inputs.get("llm_prompt"):
                lines += ["", "_LLM prompt:_", "", "```text", inputs["llm_prompt"], "```"]
            if "embed_query" in inputs:
                lines += ["", "_embedding-compare query:_", "", f"> {inputs['embed_query'] or '(none)'}"]
        else:
            lines += ["", "_No tunable filters — transforms every row._"]
        if phase["score_stats"]:
            lines += ["", f"_score spread:_ `{json.dumps(phase['score_stats'])}`"]
        if phase["metrics"]:
            lines += ["", f"_engine metrics:_ `{json.dumps(phase['metrics'], ensure_ascii=False)}`"]
        lines.append("")
    return "\n".join(lines)


def _render_export(outputs: list[PhaseOutput], spec: JDSpec, ideal_summary: str, persist: bool) -> None:
    """Offer the report as markdown + JSON downloads; optionally also write it under data/reports/."""
    payload = _report_payload(outputs, spec, ideal_summary)
    markdown = _report_markdown(payload)
    json_text = json.dumps(payload, indent=2, ensure_ascii=False, default=str)
    st.subheader("Export — per-phase input/output report")
    left, right = st.columns(2)
    left.download_button("Download report.md", markdown, file_name="phase_report.md",
                         mime="text/markdown", width="stretch")
    right.download_button("Download report.json", json_text, file_name="phase_report.json",
                          mime="application/json", width="stretch")
    if persist:
        reports_dir = REPO_ROOT / "data" / "reports"
        reports_dir.mkdir(parents=True, exist_ok=True)
        (reports_dir / "phase_report.md").write_text(markdown, encoding="utf-8")
        (reports_dir / "phase_report.json").write_text(json_text, encoding="utf-8")
        st.success(f"Report written to {reports_dir / 'phase_report.md'} and phase_report.json")


@st.cache_resource(show_spinner="Warming models into memory (one-time per server process)…")
def _warm_models(_settings: Settings) -> bool:
    """Eagerly load every model into the process's memory ONCE per server process — not once per session/rerun
    (``_settings`` is underscore-prefixed so Streamlit doesn't try to hash the Pydantic object; the cache key is
    just "has this ever run", which is exactly what we want here).

    ``rank.py --setup`` at Docker BUILD time only downloads weights to disk (a separate, short-lived process —
    its warm ``@lru_cache`` is gone the moment that process exits). The actual Streamlit server starts a FRESH
    process with an empty cache, so without this, the first real request pays the full cold-load cost (the LLM
    alone takes ~7s) instead of it happening once at server startup, before any user is waiting on it.
    """
    from retrieve.granite import _load_granite
    from retrieve.potion import _load_potion
    from retrieve.reranker import _load_reranker

    log = get_logger("streamlit")
    log.info("streamlit.warm_models_start")
    started = time.perf_counter()
    index_gen_client(_settings).ensure()
    _load_potion(_settings.models.embedder_broad)
    threads = _settings.compute.cpu_threads
    _load_granite(_settings.models.embedder_precision, _settings.precision.onnx_file, _settings.precision.max_length, threads)
    _load_reranker(_settings.models.reranker, _settings.rerank.onnx_file, _settings.rerank.max_length, threads)
    log.info("streamlit.warm_models_complete", seconds=round(time.perf_counter() - started, 3))
    return True


@st.cache_resource(show_spinner="Warming candidate data into memory (one-time per server process)…")
def _warm_candidates(_settings: Settings) -> bool:
    """Eagerly populate ``_load_phase_input``'s cache with the FULL candidate pool, ONCE per server process —
    same pattern as ``_warm_models`` (``_settings`` underscore-prefixed so Streamlit doesn't hash the Pydantic
    object; the cache key is just "has this run once").

    Warms ``limit=0`` (the full ~100K pool) specifically. The sidebar's own default is ``limit=1000``, which
    is already fast (reads only the first 1000 JSONL lines, never touches the rest of the 465MB file) and
    doesn't need this. ``limit=0`` is the case that matches the real production reproduce run — even on a
    genuinely cold container (no Parquet cache yet), this hides the ~13-20s parse inside server startup
    instead of the user's first "Run pipeline" click, the same way ``_warm_models`` hides the model cold-load
    cost. In a built image (``rank.py --setup`` already ran), ``load_candidates`` itself reads straight from
    its own Parquet cache (~1-2s) — this call is then near-instant, just populating the SESSION-level cache
    too so later reruns skip even that.
    """
    log = get_logger("streamlit")
    log.info("streamlit.warm_candidates_start")
    started = time.perf_counter()
    _load_phase_input(str(DEFAULT_CANDIDATES), 0)
    log.info("streamlit.warm_candidates_complete", seconds=round(time.perf_counter() - started, 3))
    return True


def main() -> None:
    """The Streamlit page."""
    settings = load_settings()  # safe to load before configure_compute() — config loading imports no compute libs
    configure_compute(settings.compute.cpu_threads)  # thread-pool budget (covers the lazy onnxruntime imports too)
    st.set_page_config(page_title="Reranker — phase explorer", layout="wide")  # MUST be the first st.* call
    _warm_models(settings)      # runs once per server process, on the FIRST page load — before "Run pipeline"
    _warm_candidates(settings)  # same — pre-parses the full candidate pool so a limit=0 run isn't cold either
    st.title("Reranker — per-phase explorer")
    st.caption(
        "JD→top-100 candidate ranker (one online run, ≤5 min, CPU). Each pipeline phase below shows its "
        "INPUT (filters/config, LLM prompt, or embedding-compare query) and OUTPUT — JD parse "
        "(objective_extract→summary_extract) then preprocess (load→validate→materialize→quality) then rank "
        "(hard_filter→potion→soft→composite→granite→cross_encoder→ensemble→reasoning). Total time includes the "
        "JD-parse phases."
    )

    with st.sidebar:
        st.header("Run controls")
        candidates_path = st.text_input("Candidates file", str(DEFAULT_CANDIDATES))
        limit = int(st.number_input("Max candidates (0 = all)", min_value=0, value=1000, step=500))
        regenerate_jd = st.checkbox(
            "Regenerate JD extraction (live LLM parse)", value=False,
            help="re-run the JD parse (Qwen3) to write a fresh data/jd_extraction.json before ranking; "
                 "off = use the existing file",
        )
        save_checkpoints = st.checkbox(
            "Save each phase to Parquet", value=False, help="writes data/reports/debug/NN_<phase>.parquet"
        )
        persist_report = st.checkbox(
            "Also write report to data/reports/", value=False,
            help="persist phase_report.md + .json on disk (downloads are always offered)",
        )
        run_clicked = st.button("Run pipeline", type="primary", width="stretch")

    # On-disk snapshot before the run; the JD-parse phases below show the extraction actually used this run.
    _render_jd_panel(
        extract_spec(settings),
        read_subjective(resolve(settings.jd.extraction_path)).get("ideal_summary", ""),
    )

    if not run_clicked:
        st.info("Set options in the sidebar and click **Run pipeline**.")
        return
    candidates = resolve_candidates(candidates_path)
    if not candidates.exists():
        st.error(f"File not found: {candidates_path} (also tried under the repo root).")
        return

    run_started = time.perf_counter()
    spinner_msg = "Parsing JD (live Qwen3) then running phases…" if regenerate_jd else "Running phases…"
    st.subheader("Timing")
    # Live-updating placeholders: rewritten after EVERY phase (not once at the end), and a "currently running"
    # line that fires BEFORE a phase starts — so a slow phase (e.g. JD parse or LLM polish on weak CPU hardware)
    # is visible while it's still in flight instead of leaving the page blank until the whole funnel finishes.
    timing_placeholder = st.empty()
    status_placeholder = st.empty()
    outputs: list[PhaseOutput] = []

    def _on_stage_start(name: str) -> None:
        status_placeholder.info(f"▶ Running **{name}** — elapsed so far: {time.perf_counter() - run_started:.1f}s")

    with st.spinner(spinner_msg):
        for output in _run_pipeline(
            str(candidates), limit, save_checkpoints, regenerate_jd, on_stage_start=_on_stage_start
        ):
            outputs.append(output)
            with timing_placeholder.container():
                _render_timing_summary(outputs)
            _render_phase(output)
    status_placeholder.empty()
    wall_seconds = time.perf_counter() - run_started
    # Spec/summary as written by the JD phases this run (fresh if regenerated) — for the export header.
    spec = extract_spec(settings)
    ideal_summary = read_subjective(resolve(settings.jd.extraction_path)).get("ideal_summary", "")
    funnel = [output for output in outputs if output.is_funnel]
    st.caption(f"wall-clock (incl. model loads): {wall_seconds:.1f}s")
    st.success(
        f"{len(outputs)} phase(s) complete (incl. 2 JD-parse): "
        f"{funnel[0].frame.height:,} → {funnel[-1].frame.height:,}."
    )
    _render_export(outputs, spec, ideal_summary, persist_report)


main()