kglite 0.17.2

//! `aggregate` primitive: group source nodes by a composite key,
//! evaluate per-group aggregate expressions, emit one summary
//! node per group plus optional FK edges to the group-key targets.
//!
//! Aggregate functions supported (only valid as the OUTERMOST call
//! in an `agg:` expression — nested aggregates are rejected at the
//! AST walker):
//!
//! - `count(*)` — group cardinality
//! - `count_distinct(expr)` — distinct value count
//! - `sum(expr)`, `avg(expr)`, `min(expr)`, `max(expr)` — eval
//!   `expr` per row, accumulate
//! - `first(expr, by=order_expr)`, `last(expr, by=order_expr)` —
//!   order-sensitive picks
//!
//! Outputs:
//! - An allocated CSV under `computed/` — one row per group with
//!   pk = `group:` followed by the JSON string tuple of group_by values
//!
//! Declared FK edges reference group-key columns in this same summary CSV.
//!
//! Plus a synthesised `NodeSpec[into]` registered in the
//! blueprint so the standard Phase 2/3 loader picks up the
//! summary nodes, and `JunctionEdge` entries for each `edges[]`
//! entry registered on `NodeSpec[into].connections.junction_edges`
//! so Phase 5 wires the FK edges.

use super::super::schema::ComputeOp;
use super::paths::{ComputePaths, Output};
use std::collections::{HashMap, HashSet};
use std::path::Path;

use indexmap::IndexMap;

use super::super::expr::{self, value_cmp, Bindings, Expr, Value};
use super::super::schema::{AggregateEdge, Blueprint, JunctionEdge, NodeSpec};
use super::output::StagedCsv;
use super::values::{group_id, group_key, ComputedType, GroupKey};
use super::{csv_cell_to_value, resolve_input_path, resolve_source_spec, value_to_csv_cell};

struct RowBindings<'a> {
    headers: &'a [String],
    values: &'a [Value],
}

impl<'a> Bindings for RowBindings<'a> {
    fn get(&self, name: &str) -> Option<Value> {
        self.headers
            .iter()
            .position(|h| h == name)
            .map(|i| self.values[i].clone())
    }
}

#[derive(Clone)]
enum AggKind {
    Count,               // count(*)
    CountDistinct(Expr), // count_distinct(expr)
    Sum(Expr),
    Avg(Expr),
    Min(Expr),
    Max(Expr),
    First {
        value: Expr,
        order: Option<Expr>,
    },
    Last {
        value: Expr,
        order: Option<Expr>,
    },
    /// Row-level expression — no aggregation; the outermost call
    /// is e.g. `if(...)` or a literal. Evaluates against the FIRST
    /// row of the group. Useful for aggregating constants per group.
    RowLevel(Expr),
}

#[derive(Default)]
struct AggState {
    count: i64,
    sum: f64,
    integer_sum: crate::graph::core::numeric_sum::IntegerSum,
    saw_float: bool,
    n_for_avg: u64,
    min: Option<Value>,
    max: Option<Value>,
    first_order: Option<Value>,
    first_value: Option<Value>,
    last_order: Option<Value>,
    last_value: Option<Value>,
    distinct: HashSet<String>,
}

// Preserve standalone allocation entry; pipeline dispatch shares its allocated paths.
#[allow(dead_code)]
pub fn run_aggregate(
    blueprint: &mut Blueprint,
    input_root: &Path,
    from: &str,
    group_by: &[String],
    into: &str,
    agg: &IndexMap<String, String>,
    edges: &[AggregateEdge],
) -> Result<(), String> {
    run_aggregate_allocated(
        blueprint, input_root, from, group_by, into, agg, edges, None,
    )
}

// The allocated entry mirrors the existing compute primitive arguments.
#[allow(clippy::too_many_arguments)]
pub(super) fn run_aggregate_allocated(
    blueprint: &mut Blueprint,
    input_root: &Path,
    from: &str,
    group_by: &[String],
    into: &str,
    agg: &IndexMap<String, String>,
    edges: &[AggregateEdge],
    paths: Option<&ComputePaths>,
) -> Result<(), String> {
    let spec = resolve_source_spec(blueprint, from)
        .ok_or_else(|| format!("aggregate: source type '{}' not declared", from))?;
    let csv_rel = spec.csv.clone().ok_or_else(|| {
        format!(
            "aggregate: source type '{}' has no csv: declared (aggregate \
             on synthesised types is deferred)",
            from
        )
    })?;
    let csv_path = resolve_input_path(input_root, &csv_rel);

    // Mirror the rest of the blueprint loader: a missing source CSV
    // is not an error — the type simply contributes zero rows. The
    // compute primitive degrades to a no-op so partial datasets
    // (e.g. SEC without XBRL extracts) build cleanly.
    if !csv_path.exists() {
        return Ok(());
    }

    validate_group_edges(group_by, edges)?;

    // 1. Parse + classify each agg expression.
    let mut classified: Vec<(String, AggKind)> = Vec::with_capacity(agg.len());
    for (prop, src) in agg {
        let ast = expr::parse(src)
            .map_err(|e| format!("aggregate '{}': expression parse: {}", prop, e))?;
        let kind = classify_aggregate(ast).map_err(|e| format!("aggregate '{}': {}", prop, e))?;
        classified.push((prop.clone(), kind));
    }

    // 2. Read CSV headers + property type hints for typed cell→Value.
    let mut reader = csv::ReaderBuilder::new()
        .has_headers(true)
        .from_path(&csv_path)
        .map_err(|e| format!("aggregate: open {}: {}", csv_path.display(), e))?;
    let headers: Vec<String> = reader
        .headers()
        .map_err(|e| format!("aggregate: header: {}", e))?
        .iter()
        .map(|s| s.to_string())
        .collect();
    let mut declared_types: HashMap<String, String> = HashMap::new();
    for (col, ty) in &spec.properties {
        declared_types.insert(col.clone(), ty.clone());
    }
    let group_indices: Vec<usize> = group_by
        .iter()
        .map(|g| {
            headers
                .iter()
                .position(|h| h == g)
                .ok_or_else(|| format!("aggregate: group_by '{}' not in headers", g))
        })
        .collect::<Result<_, _>>()?;

    // 3. Single pass over CSV: build per-group state via HashMap
    //    for O(1) avg per-row lookup. Output is sorted at the end
    //    so test/CSV determinism survives.
    let mut groups: HashMap<GroupKey, Vec<AggState>> = HashMap::new();
    let mut row_values: Vec<Value> = Vec::with_capacity(headers.len());
    let n_aggs = classified.len();

    for record_result in reader.records() {
        let record = record_result.map_err(|e| format!("aggregate: row: {}", e))?;
        row_values.clear();
        for (i, h) in headers.iter().enumerate() {
            let cell = record.get(i).unwrap_or("");
            row_values.push(csv_cell_to_value(
                cell,
                declared_types.get(h).map(|s| s.as_str()),
            ));
        }

        let states = groups
            .entry(group_key(&record, &group_indices))
            .or_insert_with(|| (0..n_aggs).map(|_| AggState::default()).collect());

        let bindings = RowBindings {
            headers: &headers,
            values: &row_values,
        };

        for (i, (_prop, kind)) in classified.iter().enumerate() {
            update_state(&mut states[i], kind, &bindings)
                .map_err(|e| format!("aggregate: {}", e))?;
        }
    }

    // 4. Emit summary CSV.
    let owned_paths;
    let paths = if let Some(paths) = paths {
        paths
    } else {
        let operation = ComputeOp::Aggregate {
            from: from.to_string(),
            group_by: group_by.to_vec(),
            into: into.to_string(),
            agg: agg.clone(),
            edges: edges.to_vec(),
        };
        owned_paths = ComputePaths::new(blueprint, input_root, std::slice::from_ref(&operation))?;
        &owned_paths
    };
    let computed_rel = paths.relative(&Output::Aggregate(into.to_string()));
    let out_path = input_root.join(&computed_rel);
    let mut staged = StagedCsv::new(&out_path, "aggregate")?;
    let writer = staged.writer();

    // Header: composite-pk + group_by cols + each agg property
    let pk_col = format!("{}_id", sanitize(into).to_lowercase());
    let mut hdr: Vec<String> = vec![pk_col.clone()];
    for g in group_by {
        hdr.push(g.clone());
    }
    for (prop, _) in &classified {
        hdr.push(prop.clone());
    }
    writer
        .write_record(&hdr)
        .map_err(|e| format!("aggregate: write header: {}", e))?;

    let mut inferred_types: HashMap<String, ComputedType> = HashMap::new();

    // Sort keys for deterministic CSV order (HashMap iteration is
    // randomised). Sort cost is O(g log g) for g = group count —
    // dominated by the per-row work for large inputs.
    let mut sorted_keys: Vec<&GroupKey> = groups.keys().collect();
    sorted_keys.sort();

    for key in sorted_keys {
        let states = &groups[key];
        let pk_value = group_id(key)?;
        let mut row: Vec<String> = Vec::with_capacity(hdr.len());
        row.push(pk_value);
        for c in key {
            row.push(c.clone());
        }
        append_aggregate_cells(states, &classified, &mut inferred_types, &mut row)?;
        writer
            .write_record(&row)
            .map_err(|e| format!("aggregate: write row: {}", e))?;
    }

    // 5. Prepare the summary NodeSpec before publishing output.
    let mut into_spec = NodeSpec {
        csv: Some(computed_rel.clone()),
        pk: Some(pk_col.clone()),
        title: Some(pk_col.clone()),
        ..NodeSpec::default()
    };
    // Group-by columns are properties of the summary too.
    for g in group_by {
        let ty = declared_types
            .get(g)
            .cloned()
            .unwrap_or_else(|| "string".to_string());
        into_spec.properties.insert(g.clone(), ty);
    }
    for (prop, _) in &classified {
        let ty = inferred_types
            .get(prop)
            .map_or("string", ComputedType::resolve);
        into_spec.properties.insert(prop.clone(), ty.to_string());
    }

    // 6. FK edges from summary → group-key targets.
    for edge in edges {
        into_spec.connections.junction_edges.insert(
            edge.edge.clone(),
            JunctionEdge::computed(
                computed_rel.clone(),
                pk_col.clone(),
                edge.to.clone(),
                edge.fk.clone(),
            ),
        );
    }

    drop(reader);
    staged.publish()?;
    blueprint.nodes.insert(into.to_string(), into_spec);

    Ok(())
}

/// A summary FK must address a group column. Refuse before opening output.
fn validate_group_edges(group_by: &[String], edges: &[AggregateEdge]) -> Result<(), String> {
    for edge in edges {
        if !group_by.iter().any(|group| group == &edge.fk) {
            return Err(format!(
                "aggregate edge '{}': fk '{}' must be one of group_by {:?}",
                edge.edge, edge.fk, group_by
            ));
        }
    }
    Ok(())
}

/// Inspect an `agg:` expression's AST. The outermost call must be
/// an aggregate function or the expression is treated as row-level
/// (against the first row of each group).
fn classify_aggregate(ast: Expr) -> Result<AggKind, String> {
    if let Expr::Call(name, args) = &ast {
        match name.as_str() {
            "count" => {
                // Either count(*) or count(expr). count(expr) treats
                // each non-null evaluation as a +1.
                if args.len() == 1 {
                    if let Expr::Ident(s) = &args[0].1 {
                        if s == "*" {
                            return Ok(AggKind::Count);
                        }
                    }
                }
                // count(expr) — treat as Sum(if(expr is not null, 1, 0))
                // semantic. This site accepts only count(*);
                // count_distinct handles the column case.
                return Err(
                    "count: only count(*) supported here (use count_distinct for column counts)"
                        .to_string(),
                );
            }
            "count_distinct" => {
                if args.len() != 1 {
                    return Err("count_distinct: expected 1 argument".to_string());
                }
                return Ok(AggKind::CountDistinct(args[0].1.clone()));
            }
            "sum" => {
                if args.len() != 1 {
                    return Err("sum: expected 1 argument".to_string());
                }
                return Ok(AggKind::Sum(args[0].1.clone()));
            }
            "avg" => {
                if args.len() != 1 {
                    return Err("avg: expected 1 argument".to_string());
                }
                return Ok(AggKind::Avg(args[0].1.clone()));
            }
            "min" if args.len() == 1 && args[0].0.is_none() => {
                return Ok(AggKind::Min(args[0].1.clone()));
            }
            "max" if args.len() == 1 && args[0].0.is_none() => {
                return Ok(AggKind::Max(args[0].1.clone()));
            }
            "first" | "last" => {
                if args.is_empty() {
                    return Err(format!("{}: expected at least 1 argument", name));
                }
                let mut value: Option<Expr> = None;
                let mut order: Option<Expr> = None;
                for (kw, e) in args {
                    match kw {
                        None if value.is_none() => value = Some(e.clone()),
                        Some(k) if k == "by" => order = Some(e.clone()),
                        Some(k) => return Err(format!("{}: unknown named arg '{}'", name, k)),
                        None => return Err(format!("{}: too many positional args", name)),
                    }
                }
                let value = value.ok_or_else(|| format!("{}: missing value argument", name))?;
                return Ok(if name == "first" {
                    AggKind::First { value, order }
                } else {
                    AggKind::Last { value, order }
                });
            }
            _ => {}
        }
    }
    // Not an aggregate call — treat as row-level (eval against
    // first row of the group). Useful when the value is constant
    // per group.
    Ok(AggKind::RowLevel(ast))
}

fn update_state(state: &mut AggState, kind: &AggKind, ctx: &dyn Bindings) -> Result<(), String> {
    match kind {
        AggKind::Count => {
            state.count += 1;
        }
        AggKind::CountDistinct(expr) => {
            let v = expr::eval(expr, ctx).map_err(|e| format!("count_distinct: {}", e))?;
            if !matches!(v, Value::Null) {
                state.distinct.insert(format!("{}", v));
            }
        }
        AggKind::Sum(expr) | AggKind::Avg(expr) => {
            let v = expr::eval(expr, ctx).map_err(|e| format!("sum/avg: {}", e))?;
            match v {
                Value::Int(i) => {
                    state.sum += i as f64;
                    state.integer_sum.add(i);
                    state.n_for_avg += 1;
                }
                Value::Float(f) if f.is_finite() => {
                    state.sum += f;
                    state.saw_float = true;
                    state.n_for_avg += 1;
                }
                Value::Null => {}
                Value::Bool(b) => {
                    state.sum += if b { 1.0 } else { 0.0 };
                    state.integer_sum.add(i64::from(b));
                    state.n_for_avg += 1;
                }
                _ => {} // Non-numeric values skipped.
            }
        }
        AggKind::Min(expr) => {
            let v = expr::eval(expr, ctx).map_err(|e| format!("min: {}", e))?;
            if matches!(v, Value::Null) {
                return Ok(());
            }
            match &state.min {
                None => state.min = Some(v),
                Some(cur) if value_cmp(&v, cur) == std::cmp::Ordering::Less => state.min = Some(v),
                _ => {}
            }
        }
        AggKind::Max(expr) => {
            let v = expr::eval(expr, ctx).map_err(|e| format!("max: {}", e))?;
            if matches!(v, Value::Null) {
                return Ok(());
            }
            match &state.max {
                None => state.max = Some(v),
                Some(cur) if value_cmp(&v, cur) == std::cmp::Ordering::Greater => {
                    state.max = Some(v)
                }
                _ => {}
            }
        }
        AggKind::First { value, order } => {
            let v = expr::eval(value, ctx).map_err(|e| format!("first: value: {}", e))?;
            let o = match order {
                Some(o) => Some(expr::eval(o, ctx).map_err(|e| format!("first: by: {}", e))?),
                None => None,
            };
            let take = match (&state.first_order, &o) {
                (None, _) => true,
                (Some(cur), Some(new)) => value_cmp(new, cur) == std::cmp::Ordering::Less,
                _ => false,
            };
            if take {
                state.first_value = Some(v);
                state.first_order = o;
            }
        }
        AggKind::Last { value, order } => {
            let v = expr::eval(value, ctx).map_err(|e| format!("last: value: {}", e))?;
            let o = match order {
                Some(o) => Some(expr::eval(o, ctx).map_err(|e| format!("last: by: {}", e))?),
                None => None,
            };
            let take = match (&state.last_order, &o) {
                (None, _) => true,
                (Some(cur), Some(new)) => value_cmp(new, cur) == std::cmp::Ordering::Greater,
                _ => false,
            };
            if take {
                state.last_value = Some(v);
                state.last_order = o;
            }
        }
        AggKind::RowLevel(expr) => {
            // Evaluate on the first row only — subsequent rows
            // re-evaluate but we ignore (constant-per-group
            // semantics). Stash on first_value for finalize.
            if state.first_value.is_none() {
                let v = expr::eval(expr, ctx).map_err(|e| format!("row-level agg: {}", e))?;
                state.first_value = Some(v);
            }
        }
    }
    Ok(())
}

fn append_aggregate_cells(
    states: &[AggState],
    classified: &[(String, AggKind)],
    inferred_types: &mut HashMap<String, ComputedType>,
    row: &mut Vec<String>,
) -> Result<(), String> {
    for (i, (prop, kind)) in classified.iter().enumerate() {
        let value = finalize_state(&states[i], kind)?;
        inferred_types
            .entry(prop.clone())
            .or_default()
            .observe(&value);
        row.push(value_to_csv_cell(&value));
    }
    Ok(())
}

fn finalize_state(state: &AggState, kind: &AggKind) -> Result<Value, String> {
    Ok(match kind {
        AggKind::Count => Value::Int(state.count),
        AggKind::CountDistinct(_) => Value::Int(state.distinct.len() as i64),
        AggKind::Sum(_) => {
            if state.n_for_avg == 0 {
                Value::Null
            } else if !state.saw_float {
                Value::Int(state.integer_sum.finish()?)
            } else if state.sum.fract() == 0.0
                && state.sum >= i64::MIN as f64
                && state.sum < -(i64::MIN as f64)
            {
                // Preserve integral-float output typing only when representable.
                Value::Int(state.sum as i64)
            } else {
                Value::Float(state.sum)
            }
        }
        AggKind::Avg(_) => {
            if state.n_for_avg == 0 {
                Value::Null
            } else {
                Value::Float(state.sum / state.n_for_avg as f64)
            }
        }
        AggKind::Min(_) => state.min.clone().unwrap_or(Value::Null),
        AggKind::Max(_) => state.max.clone().unwrap_or(Value::Null),
        AggKind::First { .. } => state.first_value.clone().unwrap_or(Value::Null),
        AggKind::Last { .. } => state.last_value.clone().unwrap_or(Value::Null),
        AggKind::RowLevel(_) => state.first_value.clone().unwrap_or(Value::Null),
    })
}

fn sanitize(s: &str) -> String {
    s.chars()
        .map(|c| {
            if c.is_ascii_alphanumeric() || c == '_' {
                c
            } else {
                '_'
            }
        })
        .collect()
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::fs;

    fn write_csv(path: &Path, content: &str) {
        if let Some(parent) = path.parent() {
            fs::create_dir_all(parent).unwrap();
        }
        fs::write(path, content).unwrap();
    }

    fn make_bp(csv_rel: &str, pk: &str, props: &[(&str, &str)]) -> Blueprint {
        let mut spec = NodeSpec {
            csv: Some(csv_rel.to_string()),
            pk: Some(pk.to_string()),
            ..Default::default()
        };
        for (k, v) in props {
            spec.properties.insert(k.to_string(), v.to_string());
        }
        let mut bp = Blueprint::default();
        bp.nodes.insert("T".to_string(), spec);
        bp
    }

    #[test]
    fn aggregate_count_and_sum() {
        let tmp = tempfile::tempdir().unwrap();
        write_csv(
            &tmp.path().join("t.csv"),
            "id,group,value\n1,A,10\n2,A,20\n3,B,5\n4,A,30\n5,B,15\n",
        );
        let mut bp = make_bp("t.csv", "id", &[("group", "string"), ("value", "int")]);

        let mut agg = IndexMap::new();
        agg.insert("n".to_string(), "count(*)".to_string());
        agg.insert("total".to_string(), "sum(value)".to_string());

        run_aggregate(
            &mut bp,
            tmp.path(),
            "T",
            &["group".to_string()],
            "Summary",
            &agg,
            &[],
        )
        .unwrap();

        let out = fs::read_to_string(tmp.path().join("computed/aggregate_Summary.csv")).unwrap();
        let lines: Vec<&str> = out.lines().collect();
        // header + 2 groups
        assert_eq!(lines.len(), 3);
        // Group A: count=3 sum=60
        assert!(lines
            .iter()
            .any(|l| l.contains(",A,") && l.contains(",3,") && l.contains(",60")));
        // Group B: count=2 sum=20
        assert!(lines
            .iter()
            .any(|l| l.contains(",B,") && l.contains(",2,") && l.contains(",20")));

        // Summary type registered.
        assert!(bp.nodes.contains_key("Summary"));
    }

    #[test]
    fn aggregate_min_max_avg() {
        let tmp = tempfile::tempdir().unwrap();
        write_csv(
            &tmp.path().join("t.csv"),
            "id,group,value\n1,A,3.0\n2,A,5.0\n3,A,7.0\n",
        );
        let mut bp = make_bp("t.csv", "id", &[("group", "string"), ("value", "float")]);

        let mut agg = IndexMap::new();
        agg.insert("lo".to_string(), "min(value)".to_string());
        agg.insert("hi".to_string(), "max(value)".to_string());
        agg.insert("mean".to_string(), "avg(value)".to_string());

        run_aggregate(
            &mut bp,
            tmp.path(),
            "T",
            &["group".to_string()],
            "Stats",
            &agg,
            &[],
        )
        .unwrap();

        let out = fs::read_to_string(tmp.path().join("computed/aggregate_Stats.csv")).unwrap();
        // Only one group, so one data line.
        let lines: Vec<&str> = out.lines().collect();
        assert_eq!(lines.len(), 2);
        // lo=3.0 hi=7.0 mean=5.0
        assert!(lines[1].contains("3.0"));
        assert!(lines[1].contains("7.0"));
        assert!(lines[1].contains("5.0"));
    }

    #[test]
    fn aggregate_last_by_ordering() {
        let tmp = tempfile::tempdir().unwrap();
        // Latest transaction per person — `last(value, by=date)`
        write_csv(
            &tmp.path().join("t.csv"),
            "id,person,date,balance\n\
             1,Alice,2025-01-01,100\n\
             2,Alice,2025-02-01,150\n\
             3,Alice,2025-03-01,200\n\
             4,Bob,2025-01-15,50\n\
             5,Bob,2025-02-15,75\n",
        );
        let mut bp = make_bp(
            "t.csv",
            "id",
            &[("person", "string"), ("date", "string"), ("balance", "int")],
        );

        let mut agg = IndexMap::new();
        agg.insert(
            "latest_balance".to_string(),
            "last(balance, by=date)".to_string(),
        );

        run_aggregate(
            &mut bp,
            tmp.path(),
            "T",
            &["person".to_string()],
            "Position",
            &agg,
            &[],
        )
        .unwrap();

        let out = fs::read_to_string(tmp.path().join("computed/aggregate_Position.csv")).unwrap();
        // Alice latest balance = 200 (date 2025-03-01)
        assert!(
            out.contains(",Alice,200"),
            "expected Alice latest=200, got {}",
            out
        );
        // Bob latest balance = 75 (date 2025-02-15)
        assert!(
            out.contains(",Bob,75"),
            "expected Bob latest=75, got {}",
            out
        );
    }

    #[test]
    fn aggregate_emits_fk_edges() {
        let tmp = tempfile::tempdir().unwrap();
        write_csv(
            &tmp.path().join("t.csv"),
            "id,person,issuer,value\n1,Alice,Apple,100\n2,Alice,Apple,200\n",
        );
        let mut bp = make_bp(
            "t.csv",
            "id",
            &[("person", "string"), ("issuer", "string"), ("value", "int")],
        );

        let mut agg = IndexMap::new();
        agg.insert("total".to_string(), "sum(value)".to_string());

        let edges = vec![
            AggregateEdge {
                to: "Person".to_string(),
                fk: "person".to_string(),
                edge: "OF_PERSON".to_string(),
            },
            AggregateEdge {
                to: "Company".to_string(),
                fk: "issuer".to_string(),
                edge: "AT_COMPANY".to_string(),
            },
        ];

        run_aggregate(
            &mut bp,
            tmp.path(),
            "T",
            &["person".to_string(), "issuer".to_string()],
            "Position",
            &agg,
            &edges,
        )
        .unwrap();

        // FK junctions registered on the Position spec.
        let pos = &bp.nodes["Position"];
        assert!(pos.connections.junction_edges.contains_key("OF_PERSON"));
        assert!(pos.connections.junction_edges.contains_key("AT_COMPANY"));
        assert_eq!(
            pos.connections.junction_edges["OF_PERSON"].target,
            vec!["Person"]
        );
        assert_eq!(
            pos.connections.junction_edges["AT_COMPANY"].target_fk,
            "issuer"
        );
    }

    #[test]
    fn aggregate_count_distinct() {
        let tmp = tempfile::tempdir().unwrap();
        write_csv(
            &tmp.path().join("t.csv"),
            "id,group,tag\n1,A,foo\n2,A,bar\n3,A,foo\n4,B,baz\n",
        );
        let mut bp = make_bp("t.csv", "id", &[("group", "string"), ("tag", "string")]);

        let mut agg = IndexMap::new();
        agg.insert("n_tags".to_string(), "count_distinct(tag)".to_string());

        run_aggregate(
            &mut bp,
            tmp.path(),
            "T",
            &["group".to_string()],
            "Out",
            &agg,
            &[],
        )
        .unwrap();
        let out = fs::read_to_string(tmp.path().join("computed/aggregate_Out.csv")).unwrap();
        // Group A: distinct tags = {foo, bar} = 2
        assert!(
            out.lines()
                .any(|l| l.starts_with("A_,") || l.contains(",A,2")),
            "{}",
            out
        );
        // Group B: distinct = {baz} = 1
        assert!(out.contains(",B,1"));
    }

    #[test]
    fn aggregate_sum_of_conditional_expression() {
        let tmp = tempfile::tempdir().unwrap();
        // Sum of total_value only for buy-side transactions.
        write_csv(
            &tmp.path().join("t.csv"),
            "id,person,code,shares,price\n\
             1,A,P,10,5.0\n\
             2,A,S,5,5.0\n\
             3,A,P,20,5.0\n",
        );
        let mut bp = make_bp(
            "t.csv",
            "id",
            &[
                ("person", "string"),
                ("code", "string"),
                ("shares", "int"),
                ("price", "float"),
            ],
        );

        let mut agg = IndexMap::new();
        // Sum the shares*price only when code='P' (buy)
        agg.insert(
            "buy_value".to_string(),
            "sum(if(code == 'P', shares * price, 0))".to_string(),
        );

        run_aggregate(
            &mut bp,
            tmp.path(),
            "T",
            &["person".to_string()],
            "Buys",
            &agg,
            &[],
        )
        .unwrap();

        // Total buy value: 10*5 + 20*5 = 50 + 100 = 150 (sells skipped)
        let out = fs::read_to_string(tmp.path().join("computed/aggregate_Buys.csv")).unwrap();
        assert!(out.contains(",A,150"), "{}", out);
    }

    fn sum_state(values: &[Value]) -> Result<Value, String> {
        let kind = AggKind::Sum(expr::parse("v").unwrap());
        let headers = vec!["v".to_string()];
        let mut state = AggState::default();
        for value in values {
            let cells = [value.clone()];
            update_state(
                &mut state,
                &kind,
                &RowBindings {
                    headers: &headers,
                    values: &cells,
                },
            )?;
        }
        finalize_state(&state, &kind)
    }

    #[test]
    fn blueprint_sum_preserves_empty_bool_float_and_exact_integer_policies() {
        assert!(matches!(sum_state(&[]).unwrap(), Value::Null));
        assert!(matches!(
            sum_state(&[Value::Bool(true), Value::Bool(false)]).unwrap(),
            Value::Int(1)
        ));
        assert!(matches!(
            sum_state(&[Value::Float(1.5), Value::Float(0.5)]).unwrap(),
            Value::Int(2)
        ));
        assert!(matches!(
            sum_state(&[Value::Float(f64::INFINITY)]).unwrap(),
            Value::Null
        ));
        assert!(matches!(
            sum_state(&[
                Value::Int(9_007_199_254_740_993),
                Value::Int(-9_007_199_254_740_992)
            ])
            .unwrap(),
            Value::Int(1)
        ));
        assert!(matches!(
            sum_state(&[Value::Int(i64::MAX), Value::Int(1), Value::Int(-i64::MAX)]).unwrap(),
            Value::Int(1)
        ));
        assert!(sum_state(&[Value::Int(i64::MAX), Value::Int(1)])
            .unwrap_err()
            .contains("Integer overflow in sum"));
        assert!(matches!(
            sum_state(&[Value::Float(9_223_372_036_854_775_808.0)]).unwrap(),
            Value::Float(_)
        ));
    }

    #[test]
    fn blueprint_sum_overflow_propagates_from_csv_emission() {
        let tmp = tempfile::tempdir().unwrap();
        write_csv(
            &tmp.path().join("t.csv"),
            "id,group,value\n1,A,9223372036854775807\n2,A,1\n",
        );
        let mut blueprint = make_bp("t.csv", "id", &[("group", "string"), ("value", "int")]);
        let agg = IndexMap::from([("total".to_string(), "sum(value)".to_string())]);
        let result = run_aggregate(
            &mut blueprint,
            tmp.path(),
            "T",
            &["group".to_string()],
            "Summary",
            &agg,
            &[],
        );
        assert!(result.unwrap_err().contains("Integer overflow in sum"));
        assert!(!blueprint.nodes.contains_key("Summary"));
    }

    fn assert_failed_repeat_preserves_summary(invalid_fk: bool) {
        let tmp = tempfile::tempdir().unwrap();
        let source = tmp.path().join("t.csv");
        let original = "id,g,bucket,v\n1,a,x,9223372036854775807\n2,b,x,1\n";
        write_csv(&source, original);
        let mut bp = make_bp(
            "t.csv",
            "id",
            &[("g", "string"), ("bucket", "string"), ("v", "int")],
        );
        run_aggregate(
            &mut bp,
            tmp.path(),
            "T",
            &["g".into(), "bucket".into()],
            "Totals",
            &IndexMap::from([("s".into(), "sum(v)".into())]),
            &[],
        )
        .unwrap();
        let summary = tmp.path().join("computed/aggregate_Totals.csv");
        let before = fs::read(&summary).unwrap();
        let schema = format!("{bp:?}");
        let edges = if invalid_fk {
            vec![AggregateEdge {
                to: "T".into(),
                fk: "missing".into(),
                edge: "BAD".into(),
            }]
        } else {
            vec![]
        };
        let expression = if invalid_fk { "count(*)" } else { "sum(s)" };
        let error = run_aggregate(
            &mut bp,
            tmp.path(),
            "Totals",
            &["bucket".into()],
            "Totals",
            &IndexMap::from([("s".into(), expression.into())]),
            &edges,
        )
        .unwrap_err();
        assert!(
            error.contains(if invalid_fk {
                "must be one of group_by"
            } else {
                "Integer overflow in sum"
            }),
            "{error}"
        );
        assert_eq!(
            fs::read(&summary).unwrap(),
            before,
            "failed step replaced its active source"
        );
        assert_eq!(format!("{bp:?}"), schema);
        assert_eq!(fs::read_to_string(source).unwrap(), original);
        let files: Vec<_> = fs::read_dir(tmp.path().join("computed"))
            .unwrap()
            .map(|entry| entry.unwrap().file_name())
            .collect();
        assert_eq!(
            files,
            vec![std::ffi::OsString::from("aggregate_Totals.csv")]
        );
    }

    #[test]
    fn repeated_aggregate_overflow_preserves_completed_input_and_schema() {
        assert_failed_repeat_preserves_summary(false);
    }

    #[test]
    fn invalid_aggregate_fk_preserves_completed_input_and_schema() {
        assert_failed_repeat_preserves_summary(true);
    }
}