sheeter-engine 1.2.0

use crate::{CanonicalValue, Transforms};
use std::error::Error;
use std::io::{BufRead, BufReader, Read, Write};

const SCHEMA_SAMPLE_ROWS: usize = 10_000;

/// Streams NDJSON into CSV/TSV, sampling the first `SCHEMA_SAMPLE_ROWS`
/// rows to discover the header union, then streaming the rest.
///
/// Keys that appear only after the sample window are silently dropped.
/// Use `--select` to lock the schema explicitly if the input is known to
/// be heterogeneous.
///
/// # Errors
/// Returns an error if any line is malformed JSON or I/O fails.
///
/// # Examples
///
/// ```
/// use sheeter_engine::{stream_ndjson_to_delimited, Transforms};
///
/// let input = b"{\"a\":1}\n{\"a\":2}\n";
/// let mut out = Vec::new();
/// stream_ndjson_to_delimited(input.as_slice(), &mut out, b',', &Transforms::default()).unwrap();
/// assert_eq!(out, b"a\n1\n2\n");
/// ```
pub fn stream_ndjson_to_delimited(
    reader: impl Read,
    writer: impl Write,
    to_delim: u8,
    transforms: &Transforms,
) -> Result<(), Box<dyn Error>> {
    let mut buf = BufReader::new(reader);
    let mut headers: Vec<String> = Vec::new();
    let mut sample: Vec<CanonicalValue> = Vec::new();
    let mut line = String::new();

    loop {
        line.clear();
        if buf.read_line(&mut line)? == 0 {
            break;
        }
        let trimmed = line.trim();
        if trimmed.is_empty() {
            continue;
        }
        let v: serde_json::Value = serde_json::from_str(trimmed)?;
        let canon = CanonicalValue::from_json_value(v);
        if let CanonicalValue::Object(o) = &canon {
            for (k, _) in o {
                if !headers.contains(k) {
                    headers.push(k.clone());
                }
            }
        }
        sample.push(canon);
        if sample.len() >= SCHEMA_SAMPLE_ROWS {
            break;
        }
    }

    let mut renamed = headers.clone();
    for (from, to) in &transforms.rename {
        for h in renamed.iter_mut() {
            if h == from {
                *h = to.clone();
                break;
            }
        }
    }
    let final_headers: Vec<String> = match &transforms.select {
        Some(sel) => sel.clone(),
        None => renamed,
    };

    let mut wtr = csv::WriterBuilder::new()
        .delimiter(to_delim)
        .from_writer(writer);
    wtr.write_record(&final_headers)?;

    let limit = transforms.limit.unwrap_or(usize::MAX);
    let mut emitted = 0usize;
    let mut seen: std::collections::HashSet<String> = std::collections::HashSet::new();
    let dedup_cols = transforms.dedup.clone();

    for canon in sample {
        if emitted >= limit {
            break;
        }
        if let CanonicalValue::Object(row) = canon {
            if let Some(new_row) = transforms.apply_to_row(row) {
                write_row(
                    &mut wtr,
                    &new_row,
                    &final_headers,
                    &dedup_cols,
                    &mut seen,
                    &mut emitted,
                )?;
            }
        }
    }

    while emitted < limit {
        line.clear();
        if buf.read_line(&mut line)? == 0 {
            break;
        }
        let trimmed = line.trim();
        if trimmed.is_empty() {
            continue;
        }
        let v: serde_json::Value = serde_json::from_str(trimmed)?;
        if let CanonicalValue::Object(row) = CanonicalValue::from_json_value(v) {
            if let Some(new_row) = transforms.apply_to_row(row) {
                write_row(
                    &mut wtr,
                    &new_row,
                    &final_headers,
                    &dedup_cols,
                    &mut seen,
                    &mut emitted,
                )?;
            }
        }
    }

    wtr.flush()?;
    Ok(())
}

/// Streams CSV/TSV into NDJSON, one line per row.
///
/// # Errors
/// Returns an error if any row is malformed or I/O fails.
///
/// # Examples
///
/// ```
/// use sheeter_engine::{stream_delimited_to_ndjson, Transforms};
///
/// let input = b"a\n1\n2\n";
/// let mut out = Vec::new();
/// stream_delimited_to_ndjson(input.as_slice(), b',', &mut out, &Transforms::default()).unwrap();
/// assert_eq!(out, b"{\"a\":\"1\"}\n{\"a\":\"2\"}\n");
/// ```
pub fn stream_delimited_to_ndjson(
    reader: impl Read,
    from_delim: u8,
    writer: impl Write,
    transforms: &Transforms,
) -> Result<(), Box<dyn Error>> {
    let mut rdr = csv::ReaderBuilder::new()
        .delimiter(from_delim)
        .from_reader(reader);
    let raw_headers: Vec<String> = rdr.headers()?.iter().map(|s| s.to_string()).collect();
    let mut out = std::io::BufWriter::new(writer);
    let limit = transforms.limit.unwrap_or(usize::MAX);
    let mut emitted = 0usize;
    let mut seen: std::collections::HashSet<String> = std::collections::HashSet::new();
    let dedup_cols = transforms.dedup.clone();

    for result in rdr.records() {
        if emitted >= limit {
            break;
        }
        let record = result?;
        let mut row: Vec<(String, CanonicalValue)> = Vec::with_capacity(raw_headers.len());
        for (i, h) in raw_headers.iter().enumerate() {
            row.push((
                h.clone(),
                CanonicalValue::String(record.get(i).unwrap_or("").to_string()),
            ));
        }
        let new_row = match transforms.apply_to_row(row) {
            Some(r) => r,
            None => continue,
        };
        if let Some(cols) = &dedup_cols {
            let key = dedup_key(&new_row, cols);
            if !seen.insert(key) {
                continue;
            }
        }
        let obj = CanonicalValue::Object(new_row);
        let json = serde_json::to_string(&obj.to_json_value())?;
        writeln!(out, "{json}")?;
        emitted += 1;
    }
    out.flush()?;
    Ok(())
}

/// Streams NDJSON to NDJSON, applying transforms per line.
///
/// # Errors
/// Returns an error if any line is malformed or I/O fails.
///
/// # Examples
///
/// ```
/// use sheeter_engine::{stream_ndjson_to_ndjson, Transforms, Filter};
///
/// let input = b"{\"a\":1}\n{\"a\":2}\n";
/// let t = Transforms { filter: Some(Filter::parse("a > 1").unwrap()), ..Default::default() };
/// let mut out = Vec::new();
/// stream_ndjson_to_ndjson(input.as_slice(), &mut out, &t).unwrap();
/// assert_eq!(out, b"{\"a\":2}\n");
/// ```
pub fn stream_ndjson_to_ndjson(
    reader: impl Read,
    writer: impl Write,
    transforms: &Transforms,
) -> Result<(), Box<dyn Error>> {
    let buf = BufReader::new(reader);
    let mut out = std::io::BufWriter::new(writer);
    let limit = transforms.limit.unwrap_or(usize::MAX);
    let mut emitted = 0usize;
    let mut seen: std::collections::HashSet<String> = std::collections::HashSet::new();
    let dedup_cols = transforms.dedup.clone();

    for line in buf.lines() {
        if emitted >= limit {
            break;
        }
        let line = line?;
        let trimmed = line.trim();
        if trimmed.is_empty() {
            continue;
        }
        let v: serde_json::Value = serde_json::from_str(trimmed)?;
        let row = match CanonicalValue::from_json_value(v) {
            CanonicalValue::Object(o) => o,
            _ => continue,
        };
        let new_row = match transforms.apply_to_row(row) {
            Some(r) => r,
            None => continue,
        };
        if let Some(cols) = &dedup_cols {
            let key = dedup_key(&new_row, cols);
            if !seen.insert(key) {
                continue;
            }
        }
        let obj = CanonicalValue::Object(new_row);
        let json = serde_json::to_string(&obj.to_json_value())?;
        writeln!(out, "{json}")?;
        emitted += 1;
    }
    out.flush()?;
    Ok(())
}

fn write_row<W: Write>(
    wtr: &mut csv::Writer<W>,
    row: &[(String, CanonicalValue)],
    final_headers: &[String],
    dedup_cols: &Option<Vec<String>>,
    seen: &mut std::collections::HashSet<String>,
    emitted: &mut usize,
) -> Result<(), Box<dyn Error>> {
    if let Some(cols) = dedup_cols {
        let key = dedup_key(row, cols);
        if !seen.insert(key) {
            return Ok(());
        }
    }
    let values: Vec<String> = final_headers
        .iter()
        .map(|h| {
            row.iter()
                .find(|(k, _)| k == h)
                .map(|(_, v)| canon_str(v))
                .unwrap_or_default()
        })
        .collect();
    wtr.write_record(&values)?;
    *emitted += 1;
    Ok(())
}

fn dedup_key(row: &[(String, CanonicalValue)], cols: &[String]) -> String {
    if cols.is_empty() {
        row.iter()
            .map(|(_, v)| canon_str(v))
            .collect::<Vec<_>>()
            .join("\x1f")
    } else {
        cols.iter()
            .map(|c| {
                row.iter()
                    .find(|(k, _)| k == c)
                    .map(|(_, v)| canon_str(v))
                    .unwrap_or_default()
            })
            .collect::<Vec<_>>()
            .join("\x1f")
    }
}

fn canon_str(v: &CanonicalValue) -> String {
    match v {
        CanonicalValue::Null => String::new(),
        CanonicalValue::Bool(b) => b.to_string(),
        CanonicalValue::Number(n) => n.to_string(),
        CanonicalValue::String(s) => s.clone(),
        CanonicalValue::Bytes(b) => {
            use base64::Engine;
            base64::engine::general_purpose::STANDARD.encode(b)
        }
        _ => String::new(),
    }
}