Aura/crates/aura-registry/src/lineage.rs

//! Registry lineage for orchestration families (C18/C21): a family — a sweep,
//! Monte-Carlo, or walk-forward run — is a set of run records sharing one
//! parent identity — the `(family, run)` pair, from which the user-facing
//! `family_id` handle (`"{family}-{run}"`) is derived. This module adds the
//! persisted lineage
//! unit ([`FamilyRunRecord`]), the family-store accessors on [`Registry`]
//! ([`Registry::append_family`] / [`Registry::load_family_members`]), and the
//! round-trip that re-derives a [`Family`] from the stored links
//! ([`group_families`]). The family store is a sibling JSONL of the flat runs
//! store (`families.jsonl`), so the flat `runs.jsonl` path and API are untouched.
//!
//! Campaign realizations (cycle 0107) follow the same growth pattern: a thin
//! [`CampaignRunRecord`] linking untouched family records, one JSONL line per
//! campaign run in a second sibling store (`campaign_runs.jsonl`), written by
//! [`Registry::append_campaign_run`] and read by [`Registry::load_campaign_runs`].

use std::collections::HashMap;
use std::fs;
use std::io::Write;

use aura_core::Scalar;
use aura_engine::{
    FamilySelection, McFamily, RBootstrap, RunReport, SweepFamily, WalkForwardResult,
};
use serde::{Deserialize, Serialize};

use crate::{Generalization, Registry, RegistryError};

/// Which C12 orchestration axis produced a family.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub enum FamilyKind {
    Sweep,
    MonteCarlo,
    WalkForward,
    CrossInstrument,
}

/// One persisted family member: a run record stamped with its lineage link.
/// The parent reference is split into its two factors — `family` (the parent
/// name) and `run` (the per-name run index) — from which the user-facing
/// `family_id` handle (`"{family}-{run}"`) is derived rather than stored.
/// `ordinal` is the member's position in the family's canonical order
/// (odometer / seed-input / roll).
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct FamilyRunRecord {
    pub family: String,
    pub run: usize,
    pub kind: FamilyKind,
    pub ordinal: usize,
    pub report: RunReport,
}

impl FamilyRunRecord {
    /// The user-facing family handle — `"{family}-{run}"`, derived, never stored or parsed.
    pub fn family_id(&self) -> String {
        format!("{}-{}", self.family, self.run)
    }
}

/// A re-derived family: the members sharing one `(family, run)` key (and thus one
/// derived `family_id`), ordinal-sorted — the round-trip result of
/// [`group_families`].
#[derive(Clone, Debug, PartialEq)]
pub struct Family {
    pub id: String,
    pub kind: FamilyKind,
    pub members: Vec<FamilyRunRecord>,
}

/// Campaign realization: a THIN linking record over untouched family records
/// (#198). One JSONL line per campaign run in `campaign_runs.jsonl`, a sibling
/// of `runs.jsonl`/`families.jsonl` — the registry's growth pattern. `run` is
/// the per-campaign counter assigned by [`Registry::append_campaign_run`] (the
/// family store's per-name `run` pattern).
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct CampaignRunRecord {
    /// Campaign document content id.
    pub campaign: String,
    /// Process document content id.
    pub process: String,
    /// Per-campaign run counter — assigned on append, never caller-supplied.
    pub run: usize,
    pub seed: u64,
    pub cells: Vec<CellRealization>,
    /// Campaign-scope generalize annotations (cycle 0108), one per
    /// (strategy, window) across instruments. Absent pre-0108 — the
    /// serde-default widening keeps existing stored lines parseable.
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub generalizations: Vec<CampaignGeneralization>,
    /// The TraceStore family name this realization claims when the document
    /// requests persist_taps — a pure derivation ("{campaign8}-{run}"): the
    /// executor hands a claim sentinel (`Some`, any content) and
    /// [`Registry::append_campaign_run`] composes the derived name onto the
    /// stored line; the consumer persists the bytes (#201 d5). `None` when
    /// the document requests no taps. Absent pre-0109 — the serde-default
    /// widening keeps existing stored lines parseable.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub trace_name: Option<String>,
}

/// The recorded fault of a failed cell (#272): the pipeline stage it fired in,
/// a closed kind tag aggregations count on, and the prose detail. `no_data` is
/// one kind among others, not a special path.
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct CellFault {
    pub stage: usize,
    pub kind: CellFaultKind,
    pub detail: String,
}

/// Closed fault-kind vocabulary (#272).
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum CellFaultKind {
    NoData,
    Bind,
    Run,
    Panic,
    Window,
}

/// One failed walk-forward fold (#272): the roll ordinal that produced no OOS
/// report, with the fault that killed it. The surviving folds still pool; an
/// aggregate over this stage must name the ratio.
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct WindowFault {
    pub window_ordinal: usize,
    pub kind: CellFaultKind,
    pub detail: String,
}

/// A cell whose requested window is only partially covered by the archive
/// (#272): the effective evaluated bounds and any whole months missing inside
/// them. Recorded only when it deviates from the requested window.
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct CellCoverage {
    pub effective_from_ms: i64,
    pub effective_to_ms: i64,
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub gap_months: Vec<String>,
}

/// One realized (strategy, instrument, window) cell: the pipeline prefix that
/// actually ran (`stages` stops after an empty gate).
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct CellRealization {
    /// Blueprint content id.
    pub strategy: String,
    pub instrument: String,
    /// Inclusive epoch-ms window.
    pub window_ms: (i64, i64),
    pub stages: Vec<StageRealization>,
    /// The cell's risk regime (#219/#212) — mirrors `CellSpec::regime`: `None`
    /// is the member runner's baked default stop, `Some` an explicit document
    /// regime. Threaded through so a downstream re-run (the trace-persist C1
    /// drift alarm) binds the SAME stop the cell recorded under, instead of
    /// hardcoding the default. `#[serde(default, skip_serializing_if =
    /// "Option::is_none")]` keeps pre-#219 stored records byte-identical on
    /// round-trip (they predate the risk axis, so their implicit regime was
    /// always the default / absent).
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub regime: Option<aura_research::RiskRegime>,
    /// The regime's ordinal in the resolved regime list (0 for the default) —
    /// mirrors `CellSpec::regime_ordinal`; discriminates trace dirs for two
    /// regimes over the same (strategy, instrument, window) cell.
    /// `#[serde(default, skip_serializing_if = "is_zero_ordinal")]` for the
    /// same pre-#219 byte-identical round-trip reason.
    #[serde(default, skip_serializing_if = "is_zero_ordinal")]
    pub regime_ordinal: usize,
    /// Set when the cell failed (#272): the pipeline prefix in `stages` ran,
    /// the faulting stage did not complete, no nominee was produced.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub fault: Option<CellFault>,
    /// Set when the archive only partially covers the requested window (#272).
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub coverage: Option<CellCoverage>,
}

/// `CellRealization::regime_ordinal`'s serialize-skip predicate: the default
/// regime (ordinal 0) is omitted, keeping pre-#219 stored lines byte-stable.
fn is_zero_ordinal(n: &usize) -> bool {
    *n == 0
}

/// One realized pipeline stage. `family_id` is set for family-producing stages
/// (sweep / walk_forward); `survivor_ordinals` for gate stages (ordinals index
/// the nearest preceding `family_id`-bearing stage's family); `selection` for
/// sweep stages (walk-forward selections live in the wf family members'
/// manifests).
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct StageRealization {
    pub block: String,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub family_id: Option<String>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub survivor_ordinals: Option<Vec<usize>>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub selection: Option<StageSelection>,
    /// std::monte_carlo stages only (cycle 0108). Absent pre-0108 — the
    /// serde-default widening keeps existing stored lines parseable.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub bootstrap: Option<StageBootstrap>,
    /// std::walk_forward stages only (#272): folds that produced no OOS report;
    /// the surviving folds pooled, and aggregates must name the ratio.
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub window_faults: Vec<WindowFault>,
}

/// The recorded winner of a selection-bearing stage: its ordinal in the stage's
/// family, its winning param coordinate, and the selection provenance.
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct StageSelection {
    pub winner_ordinal: usize,
    pub params: Vec<(String, Scalar)>,
    pub selection: FamilySelection,
}

/// The recorded result of a `std::monte_carlo` stage (cycle 0108, #200 d1/d4):
/// the payload mirrors the stage's input shape.
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum StageBootstrap {
    /// After sweep/gates: one bootstrap per surviving member,
    /// `(ordinal into the population family, result)`.
    PerSurvivor(Vec<(usize, RBootstrap)>),
    /// After a walk_forward: one bootstrap over the pooled per-window
    /// OOS trade-R series (roll order).
    PooledOos(RBootstrap),
}

/// The campaign-scope generalize annotation (#200 d2): recorded at the scope
/// it is computed, keyed (strategy, window) across instruments.
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct CampaignGeneralization {
    pub strategy_ordinal: usize,
    pub window_ordinal: usize,
    /// The regime ordinal (0 for the default regime) — the third structural
    /// coordinate, so a generalize grade is unambiguous per (strategy, window,
    /// regime). `#[serde(default)]` keeps pre-feature stored records parsing.
    #[serde(default)]
    pub regime_ordinal: usize,
    /// None when fewer than 2 instrument cells produced a winner (gate
    /// truncation) — the shortfall is recorded, never computed around.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub generalization: Option<Generalization>,
    /// Per-instrument winning params — divergent winners are exposed,
    /// never averaged away. Present whenever >= 1 winner exists.
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub winners: Vec<(String, Vec<(String, Scalar)>)>,
    /// Instrument cells that contributed no winner (gate-truncated), by name.
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub missing: Vec<String>,
}

impl Registry {
    /// Assign a fresh per-name `run` index (one past the highest `run` already
    /// stored for `name`, or `0` if unseen), stamp `reports` as ordinal-ordered
    /// members carrying `name` + that `run`, and append them as N JSONL lines to
    /// the family store (a sibling of the flat runs store, `families.jsonl`,
    /// leaving the flat path untouched). Returns the derived id `"{name}-{run}"`.
    /// Reads the store once to pick the run index, then writes, the whole
    /// read-then-write critical section serialized by `Registry`'s internal
    /// `append_write_lock` (#276) — safe for concurrent writers sharing one
    /// `&Registry`.
    pub fn append_family(
        &self,
        name: &str,
        kind: FamilyKind,
        reports: &[RunReport],
    ) -> Result<String, RegistryError> {
        let _guard = self.append_write_lock.lock().unwrap_or_else(std::sync::PoisonError::into_inner);
        let run = next_run(name, &self.load_family_members()?);

        let path = self.path.with_file_name("families.jsonl");
        if let Some(parent) = path.parent().filter(|p| !p.as_os_str().is_empty()) {
            fs::create_dir_all(parent)?;
        }
        let mut file = fs::OpenOptions::new().create(true).append(true).open(&path)?;
        for (ordinal, report) in reports.iter().enumerate() {
            let record = FamilyRunRecord {
                family: name.to_string(),
                run,
                kind,
                ordinal,
                report: report.clone(),
            };
            let line = serde_json::to_string(&record).expect("a finite FamilyRunRecord serializes");
            writeln!(file, "{line}")?;
        }
        Ok(format!("{name}-{run}"))
    }

    /// Parse every stored family member, in file order. A missing file is an empty
    /// family store (`Ok(vec![])`), exactly as [`Registry::load`] treats a missing
    /// flat store.
    pub fn load_family_members(&self) -> Result<Vec<FamilyRunRecord>, RegistryError> {
        let path = self.path.with_file_name("families.jsonl");
        let text = match fs::read_to_string(&path) {
            Ok(t) => t,
            Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok(Vec::new()),
            Err(e) => return Err(RegistryError::Io(e)),
        };
        let mut records = Vec::new();
        for (i, raw) in text.lines().enumerate() {
            if raw.trim().is_empty() {
                continue;
            }
            let record = serde_json::from_str(raw)
                .map_err(|source| RegistryError::Parse { line: i + 1, source })?;
            records.push(record);
        }
        Ok(records)
    }

    /// Assign a fresh per-campaign `run` index (one past the highest `run`
    /// already stored for `record.campaign`, or `0` if unseen — the
    /// [`Registry::append_family`] counter pattern), write the record carrying
    /// that assigned run as ONE JSONL line to the campaign-run store (a sibling
    /// of the flat runs store, `campaign_runs.jsonl`), and return the assigned
    /// run. The input record's own `run` field is ignored; a `Some`
    /// `trace_name` (the executor's claim sentinel, cycle 0109) is replaced
    /// with the derived `"{campaign8}-{run}"` on the stored line, `None`
    /// stays `None`. Reads the store once to pick the run index, then writes,
    /// the whole read-then-write critical section serialized by the same
    /// `Registry` `append_write_lock` that `append_family` uses (#276) — safe
    /// for concurrent writers sharing one `&Registry`.
    pub fn append_campaign_run(
        &self,
        record: &CampaignRunRecord,
    ) -> Result<usize, RegistryError> {
        let _guard = self.append_write_lock.lock().unwrap_or_else(std::sync::PoisonError::into_inner);
        let run = next_campaign_run(&record.campaign, &self.load_campaign_runs()?);

        let path = self.path.with_file_name("campaign_runs.jsonl");
        if let Some(parent) = path.parent().filter(|p| !p.as_os_str().is_empty()) {
            fs::create_dir_all(parent)?;
        }
        let mut file = fs::OpenOptions::new().create(true).append(true).open(&path)?;
        // A `Some` trace_name (any content — the executor's claim sentinel,
        // #201 d5) is replaced with the derived name; `None` stays `None`.
        let trace_name = record.trace_name.as_ref().map(|_| derive_trace_name(&record.campaign, run));
        let stored = CampaignRunRecord { run, trace_name, ..record.clone() };
        let line = serde_json::to_string(&stored).expect("a finite CampaignRunRecord serializes");
        writeln!(file, "{line}")?;
        Ok(run)
    }

    /// Parse every stored campaign-run record, in file order. A missing file is
    /// an empty campaign-run store (`Ok(vec![])`), exactly as
    /// [`Registry::load_family_members`] treats a missing family store.
    pub fn load_campaign_runs(&self) -> Result<Vec<CampaignRunRecord>, RegistryError> {
        let path = self.path.with_file_name("campaign_runs.jsonl");
        let text = match fs::read_to_string(&path) {
            Ok(t) => t,
            Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok(Vec::new()),
            Err(e) => return Err(RegistryError::Io(e)),
        };
        let mut records = Vec::new();
        for (i, raw) in text.lines().enumerate() {
            if raw.trim().is_empty() {
                continue;
            }
            let record = serde_json::from_str(raw)
                .map_err(|source| RegistryError::Parse { line: i + 1, source })?;
            records.push(record);
        }
        Ok(records)
    }
}

/// The next per-name run index: one past the highest `run` already stored for
/// `name`, or `0` if the name is unseen. A clean numeric max over the `run`
/// field — no string parsing (the old fused-id form had to generate-and-check to
/// stay robust to `-` inside a name; the split removes that concern).
fn next_run(name: &str, records: &[FamilyRunRecord]) -> usize {
    records
        .iter()
        .filter(|r| r.family == name)
        .map(|r| r.run + 1)
        .max()
        .unwrap_or(0)
}

/// The next per-campaign run index: one past the highest `run` already stored
/// for `campaign`, or `0` if the campaign is unseen — [`next_run`]'s parallel
/// for the campaign-run store. Deliberately a separate fn, not a
/// generalization of `next_run`: the two stores' key fields stay independently
/// named and typed.
fn next_campaign_run(campaign: &str, records: &[CampaignRunRecord]) -> usize {
    records
        .iter()
        .filter(|r| r.campaign == campaign)
        .map(|r| r.run + 1)
        .max()
        .unwrap_or(0)
}

/// The TraceStore family name a campaign realization claims:
/// `"{campaign8}-{run}"` — the safe 8-char prefix of `campaign` (or the whole
/// id if shorter; the registry seam must not panic on a shorter id even
/// though the executor guarantees a 64-hex one) joined with `run` (#201 d5).
/// The single source of truth for the composition: [`Registry::
/// append_campaign_run`] calls it for the stored line, and
/// `aura-campaign::execute` calls it to mirror the same name onto the
/// returned copy.
pub fn derive_trace_name(campaign: &str, run: usize) -> String {
    let prefix = campaign.get(..8).unwrap_or(campaign);
    format!("{prefix}-{run}")
}

/// Group member records into families by their `(family, run)` key (first-seen
/// family order; each family's members ordinal-sorted), deriving the family's
/// `id` once at construction. The round-trip: `append_family(...)` then
/// `group_families(load_family_members())` recovers the families.
pub fn group_families(records: Vec<FamilyRunRecord>) -> Vec<Family> {
    let mut order: Vec<(String, usize)> = Vec::new();
    let mut by_key: HashMap<(String, usize), Family> = HashMap::new();
    for rec in records {
        let key = (rec.family.clone(), rec.run);
        let fam = by_key.entry(key.clone()).or_insert_with(|| {
            order.push(key.clone());
            Family { id: rec.family_id(), kind: rec.kind, members: Vec::new() }
        });
        fam.members.push(rec);
    }
    let mut families: Vec<Family> =
        order.into_iter().map(|k| by_key.remove(&k).expect("key was inserted")).collect();
    for fam in &mut families {
        fam.members.sort_by_key(|m| m.ordinal);
    }
    families
}

/// The member reports of a sweep family, in enumeration order — the trivial
/// extractor the CLI hands to [`Registry::append_family`].
pub fn sweep_member_reports(family: &SweepFamily) -> Vec<RunReport> {
    family.points.iter().map(|p| p.report.clone()).collect()
}

/// The member reports of a Monte-Carlo family, in seed-input order.
pub fn mc_member_reports(family: &McFamily) -> Vec<RunReport> {
    family.draws.iter().map(|d| d.report.clone()).collect()
}

/// The per-window OOS reports of a walk-forward result, in roll order.
pub fn walkforward_member_reports(result: &WalkForwardResult) -> Vec<RunReport> {
    result.windows.iter().map(|w| w.run.oos_report.clone()).collect()
}