src/redstack/pipelines/offline/compose.py

from __future__ import annotations

import csv
from collections.abc import Mapping
from pathlib import Path
from typing import TYPE_CHECKING, Final, final

from redstack.adapters.candidate_jsonl import JsonlCandidateSourceAdapter
from redstack.adapters.entropy import OfflineEntropy
from redstack.adapters.st_embedder import SentenceTransformerEmbeddingAdapter
from redstack.config.loader import (
    load_eligibility_rules,
    load_jd_anchors,
    load_lexicon_seed,
)
from redstack.domain.errors import ArtifactContractError
from redstack.features.registry import FEATURE_REGISTRY
from redstack.pipelines.offline.build_artifact_store import BuildArtifactStore
from redstack.pipelines.offline.context import OfflinePipelineContext
from redstack.pipelines.offline.graph import (
    OFFLINE_EXECUTION_GRAPH,
    OfflineExecutionGraph,
    StageNode,
)
from redstack.pipelines.offline.pipeline import OfflinePipeline, OfflinePipelineReport
from redstack.pipelines.offline.registry import (
    OFFLINE_ARTIFACT_REGISTRY,
    OfflineArtifactRegistry,
)
from redstack.pipelines.offline.runner import StageCallable, StageReceipt, StageResult
from redstack.pipelines.offline.stages import OfflineStage
from redstack.pipelines.offline.stages._labeling_seed import GoldLabelSeed, ReviewTag
from redstack.pipelines.offline.stages.archetype_discovery import (
    ArchetypeDiscoveryStage,
)
from redstack.pipelines.offline.stages.behavioral_calib import (
    BehavioralCalibrationStage,
)
from redstack.pipelines.offline.stages.census import CensusStage
from redstack.pipelines.offline.stages.embedding_gen import (
    AnchorEmbeddingStage,
    CandidateEmbeddingStage,
    CareerEmbeddingStage,
    EmbeddingManifestStage,
)
from redstack.pipelines.offline.stages.feature_importance import (
    FeatureImportanceStage,
)
from redstack.pipelines.offline.stages.feature_snapshot import FeatureSnapshotStage
from redstack.pipelines.offline.stages.honeypot_discovery import (
    HoneypotDiscoveryStage,
)
from redstack.pipelines.offline.stages.jd_concepts import JdConceptStage
from redstack.pipelines.offline.stages.labeling import LabelingStage
from redstack.pipelines.offline.stages.lexicon_discovery import LexiconDiscoveryStage
from redstack.pipelines.offline.stages.normalization import NormalizationStage
from redstack.pipelines.offline.stages.packaging import PackagingStage
from redstack.pipelines.offline.stages.ranking_calib import RankingCalibrationStage
from redstack.pipelines.offline.stages.reasoning_templates import (
    ReasoningTemplateStage,
)
from redstack.pipelines.offline.stages.reproducibility import ReproducibilityStage
from redstack.pipelines.offline.stages.risk_calib import RiskCalibrationStage
from redstack.pipelines.offline.stages.validation import ValidationStage
from redstack.pipelines.offline.stages.vocab_expansion import VocabExpansionStage
from redstack.pipelines.offline.stages.weight_search import WeightSearchStage

if TYPE_CHECKING:
    from redstack.config.schema import RedstackConfig

__all__: tuple[str, ...] = (
    "run_offline_build",
    "run_offline_build_with_locked_heuristics",
)

#: Default offline encode batch size handed to the sentence-transformers adapter.
_ST_BATCH_SIZE: Final[int] = 32


class GoldLabelSeedMissingError(ArtifactContractError):
    """O8 needs the human-curated review workspace output; none is committed."""


#: Required ``golden_labels.csv`` header. ``cited_features`` is a
#: ``;``-separated list within its cell (CSV has no native list type).
_REVIEW_CSV_FIELDS: Final[tuple[str, ...]] = (
    "candidate_id",
    "tier",
    "reasoning",
    "reviewer",
    "archetype_id",
    "is_honeypot_suspect",
    "is_borderline",
    "cited_features",
)
_CSV_TRUE: Final[frozenset[str]] = frozenset({"true", "1", "yes"})


def _parse_review_row(row: Mapping[str, str], *, line_no: int) -> ReviewTag:
    """Parse one ``golden_labels.csv`` data row into a ``ReviewTag``.

    Raises:
        GoldLabelSeedMissingError: the row is malformed (bad int, empty
            required field) — surfaced with its 1-based data line for the
            reviewer to fix.
    """
    try:
        archetype_raw = row.get("archetype_id", "").strip()
        return ReviewTag(
            candidate_id=row["candidate_id"].strip(),
            tier=int(row["tier"]),
            reasoning=row["reasoning"],
            reviewer=row["reviewer"].strip(),
            archetype_id=int(archetype_raw) if archetype_raw else None,
            is_honeypot_suspect=row.get("is_honeypot_suspect", "").strip().lower()
            in _CSV_TRUE,
            is_borderline=row.get("is_borderline", "").strip().lower() in _CSV_TRUE,
            cited_features=tuple(
                f.strip()
                for f in row.get("cited_features", "").split(";")
                if f.strip()
            ),
        )
    except (KeyError, ValueError) as exc:
        raise GoldLabelSeedMissingError(
            f"golden_labels.csv data row {line_no} is malformed: {exc}"
        ) from exc


def _load_gold_label_seed(path: Path) -> GoldLabelSeed:
    """Load the committed Offline Pipeline Part 7 review tags from ``path``.

    ``path`` is a CSV with header :data:`_REVIEW_CSV_FIELDS` (one row per
    reviewer tag); ``archetype_id`` empty means ``None``, ``is_honeypot_suspect``
    / ``is_borderline`` are ``true``/``false``-style strings, and
    ``cited_features`` is a ``;``-separated list.

    Raises:
        GoldLabelSeedMissingError: ``path`` does not exist, or a data row is
            malformed. A missing file is *not* a code bug — O8 is
            human-in-the-loop labeling; the seed is a workspace output
            authored by reviewers, never synthesized here.
    """
    if not path.is_file():
        raise GoldLabelSeedMissingError(
            f"O8 requires committed gold labels at {path} (Offline Pipeline "
            "Part 7: the human-in-the-loop labeling workspace's output) — "
            "none found. This is human-curated ground truth and cannot be "
            "generated by the pipeline; a reviewer must commit it first."
        )
    with path.open(encoding="utf-8", newline="") as handle:
        reader = csv.DictReader(handle)
        tags = tuple(
            _parse_review_row(row, line_no=line_no)
            for line_no, row in enumerate(reader, start=2)
        )
    if not tags:
        raise GoldLabelSeedMissingError(f"golden_labels.csv at {path} has no data rows")
    return GoldLabelSeed(tags=tags)


@final
class _LazyLabelingStage:
    """Defers O8's :class:`GoldLabelSeed` load until the stage actually runs.

    Lets ``plan() → run()`` reach every stage that does not depend on O8
    (O0-O7, O13*, O14) and fail only at O8 itself, with a clear message,
    instead of refusing to even wire the pipeline when no seed is committed.
    """

    stage_id: Final[str] = "O8"
    stage_version: Final[str] = LabelingStage.stage_version

    def __init__(
        self,
        golden_labels_path: Path,
        registry: OfflineArtifactRegistry = OFFLINE_ARTIFACT_REGISTRY,
    ) -> None:
        self._golden_labels_path = golden_labels_path
        self._registry = registry

    def __call__(
        self,
        ctx: OfflinePipelineContext,
        upstream: Mapping[str, StageReceipt],
    ) -> StageResult:
        seed = _load_gold_label_seed(self._golden_labels_path)
        stage = LabelingStage(seed=seed, registry=self._registry)
        return stage(ctx, upstream)


def _build_stages(
    *, configs_root: Path, golden_labels_path: Path
) -> dict[str, StageCallable]:
    """Wire every O0-O18 stage callable, loading the authoring seeds it needs."""
    lexicon_seed = load_lexicon_seed(configs_root)
    jd_anchors = load_jd_anchors(configs_root)
    eligibility_rules = load_eligibility_rules(configs_root)

    stages: tuple[StageCallable, ...] = (
        CensusStage(),
        NormalizationStage(),
        ValidationStage(),
        HoneypotDiscoveryStage(),
        LexiconDiscoveryStage(seed=lexicon_seed),
        VocabExpansionStage(),
        JdConceptStage(anchors=jd_anchors, eligibility=eligibility_rules),
        ArchetypeDiscoveryStage(),
        _LazyLabelingStage(golden_labels_path=golden_labels_path),
        WeightSearchStage(),
        FeatureImportanceStage(),
        BehavioralCalibrationStage(),
        RiskCalibrationStage(),
        CandidateEmbeddingStage(),
        AnchorEmbeddingStage(),
        CareerEmbeddingStage(),
        EmbeddingManifestStage(),
        FeatureSnapshotStage(),
        RankingCalibrationStage(),
        ReasoningTemplateStage(),
        PackagingStage(),
        ReproducibilityStage(),
    )
    return {stage.stage_id: stage for stage in stages}


def _build_context(
    config: RedstackConfig, *, code_version: str
) -> OfflinePipelineContext:
    """Bind every adapter and build the immutable offline build context.

    Raises:
        ValueError: ``config.offline`` is absent (wrong run mode).
    """
    offline = config.offline
    if offline is None:
        msg = "an offline build requires a config with an 'offline' runtime block"
        raise ValueError(msg)

    candidates_path = Path(config.paths.candidates_path).resolve()
    artifacts_root = Path(config.paths.artifacts_root).resolve()

    candidate_source = JsonlCandidateSourceAdapter(candidates_path)
    embedding_model = SentenceTransformerEmbeddingAdapter(
        offline.st_model_id,
        revision=offline.st_model_revision,
        batch_size=_ST_BATCH_SIZE,
    )
    entropy = OfflineEntropy(seed=offline.seed, as_of=offline.as_of.date())
    artifact_store = BuildArtifactStore(artifacts_root, OFFLINE_ARTIFACT_REGISTRY)

    return OfflinePipelineContext.build(
        config=config,
        candidate_source=candidate_source,
        embedding_model=embedding_model,
        artifact_store=artifact_store,
        entropy=entropy,
        feature_registry=FEATURE_REGISTRY,
        code_version=code_version,
    )


def run_offline_build(
    config: RedstackConfig,
    *,
    configs_root: Path,
    code_version: str,
    force: tuple[str, ...] | None = None,
) -> OfflinePipelineReport:
    """Bind adapters, build the context, wire O0-O18, and execute the build.

    Args:
        config: The fully-composed, validated offline ``RedstackConfig``.
        configs_root: Path to the ``configs/`` directory (authoring seeds).
        code_version: The build's code provenance, recorded into the report.
        force: Stage ids to recompute regardless of checkpoint freshness.

    Returns:
        The terminal :class:`OfflinePipelineReport`.

    Raises:
        ValueError: ``config.offline`` is absent (wrong run mode).
        GoldLabelSeedMissingError: O8 is reached with no committed gold labels.
    """
    ctx = _build_context(config, code_version=code_version)
    golden_labels_path = Path(config.paths.golden_labels_path).resolve()
    stages = _build_stages(
        configs_root=configs_root, golden_labels_path=golden_labels_path
    )
    pipeline = OfflinePipeline(stages=stages)
    return pipeline.execute(ctx, force=force)


# --------------------------------------------------------------------------- #
# Locked-heuristics bypass: no gold labels, fixed expert-authored weights.    #
# --------------------------------------------------------------------------- #
@final
class _FixedScoringWeightsStage(OfflineStage):
    """O9 substitute — package given component weights, no calibration search."""

    stage_id = "O9"
    stage_version = "locked-1.0"

    def __init__(
        self,
        weights: Mapping[str, float],
        neutral_prior: float,
        registry: OfflineArtifactRegistry = OFFLINE_ARTIFACT_REGISTRY,
    ) -> None:
        super().__init__(registry)
        self._weights = dict(weights)
        self._neutral_prior = neutral_prior

    def _run(
        self,
        ctx: OfflinePipelineContext,
        upstream: Mapping[str, StageReceipt],
    ) -> StageResult:
        payload: dict[str, object] = {
            "layout_version": ctx.layout_version,
            "weights": dict(sorted(self._weights.items())),
            "neutral_prior": self._neutral_prior,
            "calibrated_by": "locked-heuristics (no gold-label search)",
        }
        artifact = self.emit_yaml(ctx, "scoring_weights", payload)
        metrics: dict[str, object] = {
            "component_count": len(self._weights),
            "neutral_prior": self._neutral_prior,
        }
        return StageResult(artifacts=(artifact,), metrics=metrics)


@final
class _UniformFeatureImportanceStage(OfflineStage):
    """O10 substitute — flat per-feature importance, no permutation search."""

    stage_id = "O10"
    stage_version = "locked-1.0"

    def _run(
        self,
        ctx: OfflinePipelineContext,
        upstream: Mapping[str, StageReceipt],
    ) -> StageResult:
        importances = {str(d.feature_id): 1.0 for d in ctx.feature_registry.definitions}
        payload: dict[str, object] = {
            "importances": importances,
            "calibrated_by": "locked-heuristics (uniform, no permutation search)",
        }
        artifact = self.emit_json(ctx, "feature_importance", payload)
        metrics: dict[str, object] = {"features_scored": len(importances)}
        return StageResult(artifacts=(artifact,), metrics=metrics)


def _locked_heuristics_graph() -> OfflineExecutionGraph:
    """The full Part 11 DAG with O8 dropped and its dependents' edges pruned.

    O9/O16/O17 are the only nodes whose ``depends_on`` names "O8"; every other
    edge is untouched, so O11/O12 (already O8-independent) and O15/O16 (already
    gold-label-optional in their real implementations) run exactly as designed.
    """
    nodes: list[StageNode] = []
    for node in OFFLINE_EXECUTION_GRAPH.nodes:
        if node.stage_id == "O8":
            continue
        if node.stage_id == "O9":
            nodes.append(StageNode("O9", (), critical=node.critical))
            continue
        if node.stage_id in ("O16", "O17"):
            pruned = tuple(dep for dep in node.depends_on if dep != "O8")
            nodes.append(StageNode(node.stage_id, pruned, critical=node.critical))
            continue
        nodes.append(node)
    return OfflineExecutionGraph(nodes=tuple(nodes))


def run_offline_build_with_locked_heuristics(
    config: RedstackConfig,
    *,
    configs_root: Path,
    code_version: str,
    component_weights: Mapping[str, float],
    neutral_prior: float,
    force: tuple[str, ...] | None = None,
) -> OfflinePipelineReport:
    """Run the build on a reduced graph that drops O8 and its O9/O10 dependents.

    Substitutes :class:`_FixedScoringWeightsStage` /
    :class:`_UniformFeatureImportanceStage` for O9/O10 so O11/O12/O15/O16/O17/O18
    see a complete, real artifact set and O17 produces a genuine
    ``MANIFEST.json`` — without any committed gold labels.

    Args:
        config: The fully-composed, validated offline ``RedstackConfig``.
        configs_root: Path to the ``configs/`` directory (authoring seeds).
        code_version: The build's code provenance, recorded into the report.
        component_weights: One weight per ``domain.enums.ScoreComponent`` value.
        neutral_prior: The online ``ScoringPolicy.neutral_prior`` fallback.
        force: Stage ids to recompute regardless of checkpoint freshness.

    Returns:
        The terminal :class:`OfflinePipelineReport`.

    Raises:
        ValueError: ``config.offline`` is absent (wrong run mode).
    """
    ctx = _build_context(config, code_version=code_version)
    lexicon_seed = load_lexicon_seed(configs_root)
    jd_anchors = load_jd_anchors(configs_root)
    eligibility_rules = load_eligibility_rules(configs_root)

    stages: tuple[StageCallable, ...] = (
        CensusStage(),
        NormalizationStage(),
        ValidationStage(),
        HoneypotDiscoveryStage(),
        LexiconDiscoveryStage(seed=lexicon_seed),
        VocabExpansionStage(),
        JdConceptStage(anchors=jd_anchors, eligibility=eligibility_rules),
        ArchetypeDiscoveryStage(),
        _FixedScoringWeightsStage(component_weights, neutral_prior),
        _UniformFeatureImportanceStage(),
        BehavioralCalibrationStage(),
        RiskCalibrationStage(),
        CandidateEmbeddingStage(),
        AnchorEmbeddingStage(),
        CareerEmbeddingStage(),
        EmbeddingManifestStage(),
        FeatureSnapshotStage(),
        RankingCalibrationStage(),
        ReasoningTemplateStage(),
        PackagingStage(),
        ReproducibilityStage(),
    )
    pipeline = OfflinePipeline(
        stages={stage.stage_id: stage for stage in stages},
        graph=_locked_heuristics_graph(),
    )
    return pipeline.execute(ctx, force=force)