kibble 0.1.0

//! Detect an SFT dataset's schema and map each record to a `{messages}` `Row`.

use serde_json::Value;

#[derive(Debug, Clone, PartialEq)]
pub enum Schema {
    Messages,
    ShareGpt,
    Alpaca,
    PromptResponse,
    QuestionAnswer,
    Pair(String, String),
    TextOnly,
    RawFiles,
    Unknown(Vec<String>),
}

impl Schema {
    pub fn label(&self) -> String {
        match self {
            Schema::Messages => "messages".to_string(),
            Schema::ShareGpt => "sharegpt".to_string(),
            Schema::Alpaca => "alpaca".to_string(),
            Schema::PromptResponse => "prompt-response".to_string(),
            Schema::QuestionAnswer => "question-answer".to_string(),
            Schema::Pair(u, a) => format!("pair:{u}/{a}"),
            Schema::TextOnly => "text-only".to_string(),
            Schema::RawFiles => "raw-files".to_string(),
            Schema::Unknown(_) => "unknown".to_string(),
        }
    }
}

#[derive(Debug, Clone)]
pub struct ColumnMap {
    pub user: String,
    pub assistant: String,
    pub system: Option<String>,
}

fn has_str_key(v: &Value, key: &str) -> bool {
    v.get(key).and_then(|x| x.as_str()).is_some()
}

fn has_array_of(v: &Value, key: &str, a: &str, b: &str) -> bool {
    v.get(key)
        .and_then(|x| x.as_array())
        .and_then(|arr| arr.first())
        .map(|first| first.get(a).is_some() && first.get(b).is_some())
        .unwrap_or(false)
}

fn column_names(v: &Value) -> Vec<String> {
    v.as_object().map(|o| o.keys().cloned().collect()).unwrap_or_default()
}

/// Inspect the first record of `sample` and decide the schema. An `override_map` short-circuits
/// detection to the generic user/assistant shape (`PromptResponse`), letting the caller map exotic
/// columns explicitly.
pub fn detect_schema(sample: &[Value], override_map: Option<&ColumnMap>) -> Schema {
    if override_map.is_some() {
        return Schema::PromptResponse;
    }
    let Some(first) = sample.first() else {
        return Schema::Unknown(vec![]);
    };
    if has_array_of(first, "messages", "role", "content") {
        return Schema::Messages;
    }
    if has_array_of(first, "conversations", "from", "value") {
        return Schema::ShareGpt;
    }
    if has_str_key(first, "instruction") && has_str_key(first, "output") {
        return Schema::Alpaca;
    }
    if has_str_key(first, "prompt") && has_str_key(first, "response") {
        return Schema::PromptResponse;
    }
    if has_str_key(first, "question") && has_str_key(first, "answer") {
        return Schema::QuestionAnswer;
    }
    // Synonym-detected user/assistant pair (covers prompt+full_answer, query+completion, etc.),
    // after the named schemas so it never shadows them.
    const USER_SYNS: [&str; 5] = ["prompt", "instruction", "question", "query", "input"];
    const ASSISTANT_SYNS: [&str; 6] = ["response", "full_answer", "completion", "output", "answer", "assistant"];
    let user_col = USER_SYNS.iter().find(|k| has_str_key(first, k));
    let asst_col = ASSISTANT_SYNS.iter().find(|k| has_str_key(first, k));
    if let (Some(u), Some(a)) = (user_col, asst_col) {
        if u != a {
            return Schema::Pair(u.to_string(), a.to_string());
        }
    }
    // A lone text/content string with no chat shape is a completion blob, not chat.
    let cols = column_names(first);
    let text_only = (has_str_key(first, "text") || has_str_key(first, "content"))
        && cols.iter().all(|c| c == "text" || c == "content");
    if text_only {
        return Schema::TextOnly;
    }
    Schema::Unknown(cols)
}

use crate::build::{Msg, Row};
use crate::clean::clean_text;
use std::path::{Path, PathBuf};

#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub enum DropReason {
    Empty,
    MissingField(String),
    TooFewTurns,
    NotChat,
}

pub struct NormalizeReport {
    pub schema: Schema,
    pub rows_in: usize,
    pub rows_out: usize,
    pub wrote_output: bool,
    pub dropped: Vec<(DropReason, usize)>,
    pub files: Vec<String>,
}

/// Parse `--map user=<col>,assistant=<col>[,system=<col>]`. `user` and `assistant` are required.
pub fn parse_column_map(spec: &str) -> Result<ColumnMap, String> {
    let mut user = None;
    let mut assistant = None;
    let mut system = None;
    for part in spec.split(',') {
        let (k, v) = part.split_once('=').ok_or_else(|| format!("bad --map segment: {part}"))?;
        match k.trim() {
            "user" => user = Some(v.trim().to_string()),
            "assistant" => assistant = Some(v.trim().to_string()),
            "system" => system = Some(v.trim().to_string()),
            other => return Err(format!("unknown --map key: {other}")),
        }
    }
    Ok(ColumnMap {
        user: user.ok_or("--map requires user=<col>")?,
        assistant: assistant.ok_or("--map requires assistant=<col>")?,
        system,
    })
}

fn str_field(v: &Value, key: &str) -> Result<String, DropReason> {
    v.get(key)
        .and_then(|x| x.as_str())
        .map(|s| s.to_string())
        .ok_or_else(|| DropReason::MissingField(key.to_string()))
}

fn turn(role: &str, content: &str) -> Msg {
    Msg { role: role.to_string(), content: clean_text(content) }
}

/// Build a 2-message Row from a user/assistant pair, dropping if either side is empty after cleaning.
fn user_assistant_row(user: &str, assistant: &str, system: Option<&str>) -> Result<Row, DropReason> {
    let u = clean_text(user);
    let a = clean_text(assistant);
    if u.trim().is_empty() || a.trim().is_empty() {
        return Err(DropReason::Empty);
    }
    let mut messages = Vec::new();
    if let Some(s) = system {
        let s = clean_text(s);
        if !s.trim().is_empty() {
            messages.push(Msg { role: "system".to_string(), content: s });
        }
    }
    messages.push(Msg { role: "user".to_string(), content: u });
    messages.push(Msg { role: "assistant".to_string(), content: a });
    Ok(Row { messages })
}

pub fn map_record(value: &Value, schema: &Schema, map: Option<&ColumnMap>) -> Result<Row, DropReason> {
    if let Some(m) = map {
        let user = str_field(value, &m.user)?;
        let assistant = str_field(value, &m.assistant)?;
        let system = m.system.as_ref().and_then(|s| value.get(s)).and_then(|x| x.as_str());
        return user_assistant_row(&user, &assistant, system);
    }
    match schema {
        Schema::Messages => {
            let arr = value.get("messages").and_then(|x| x.as_array()).ok_or(DropReason::NotChat)?;
            let mut kept: Vec<Msg> = Vec::new();
            for m in arr {
                let role = m.get("role").and_then(|r| r.as_str()).unwrap_or("");
                if !matches!(role, "system" | "user" | "assistant") {
                    continue;
                }
                let Some(content) = m.get("content").and_then(|c| c.as_str()) else { continue };
                kept.push(turn(role, content));
            }
            if kept.len() < 2 {
                return Err(DropReason::TooFewTurns);
            }
            Ok(Row { messages: kept })
        }
        Schema::ShareGpt => {
            let arr = value.get("conversations").and_then(|x| x.as_array()).ok_or(DropReason::NotChat)?;
            let mut kept: Vec<Msg> = Vec::new();
            for m in arr {
                let from = m.get("from").and_then(|r| r.as_str()).unwrap_or("");
                let role = match from {
                    "human" | "user" => "user",
                    "gpt" | "assistant" | "chatgpt" | "bard" => "assistant",
                    "system" => "system",
                    _ => continue,
                };
                let Some(content) = m.get("value").and_then(|c| c.as_str()) else { continue };
                kept.push(turn(role, content));
            }
            if kept.len() < 2 {
                return Err(DropReason::TooFewTurns);
            }
            Ok(Row { messages: kept })
        }
        Schema::Alpaca => {
            let instruction = str_field(value, "instruction")?;
            let output = str_field(value, "output")?;
            let input = value.get("input").and_then(|x| x.as_str()).unwrap_or("");
            let user = if input.trim().is_empty() {
                instruction
            } else {
                format!("{instruction}\n{input}")
            };
            user_assistant_row(&user, &output, None)
        }
        Schema::PromptResponse => {
            let user = str_field(value, "prompt")?;
            let assistant = str_field(value, "response")?;
            user_assistant_row(&user, &assistant, None)
        }
        Schema::QuestionAnswer => {
            let user = str_field(value, "question")?;
            let assistant = str_field(value, "answer")?;
            user_assistant_row(&user, &assistant, None)
        }
        Schema::Pair(u, a) => {
            let user = str_field(value, u)?;
            let assistant = str_field(value, a)?;
            user_assistant_row(&user, &assistant, None)
        }
        Schema::TextOnly | Schema::RawFiles | Schema::Unknown(_) => Err(DropReason::NotChat),
    }
}

/// All jsonl/parquet files under `raw_dir`, recursively. Metadata files (README, .gitattributes,
/// dataset_infos.json, and other bare .json) are excluded by only accepting .jsonl and .parquet.
pub fn collect_data_files(raw_dir: &Path) -> Vec<PathBuf> {
    let mut out = Vec::new();
    let mut stack = vec![raw_dir.to_path_buf()];
    while let Some(dir) = stack.pop() {
        let Ok(rd) = std::fs::read_dir(&dir) else { continue };
        for entry in rd.flatten() {
            let p = entry.path();
            if p.is_dir() {
                stack.push(p);
            } else if matches!(
                p.extension().and_then(|e| e.to_str()).map(|s| s.to_lowercase()).as_deref(),
                Some("jsonl") | Some("parquet")
            ) {
                out.push(p);
            }
        }
    }
    out.sort();
    out
}

fn tally(drops: Vec<DropReason>) -> Vec<(DropReason, usize)> {
    use std::collections::HashMap;
    let mut counts: HashMap<DropReason, usize> = HashMap::new();
    for d in drops {
        *counts.entry(d).or_insert(0) += 1;
    }
    let mut v: Vec<(DropReason, usize)> = counts.into_iter().collect();
    v.sort_by_key(|b| std::cmp::Reverse(b.1));
    v
}

/// Read `raw_dir`'s data files, detect the schema, map each record, and (when chat-shaped) write
/// normalized `{messages}` JSONL to `out_file`. Always returns a report. For Unknown/TextOnly the
/// report has `wrote_output = false`, `rows_out = 0`, and no file is created.
pub fn normalize_dataset(
    raw_dir: &Path,
    out_file: &Path,
    override_map: Option<&ColumnMap>,
    max_rows: Option<usize>,
) -> std::io::Result<NormalizeReport> {
    let data_files = collect_data_files(raw_dir);
    let files: Vec<String> = data_files
        .iter()
        .map(|p| p.strip_prefix(raw_dir).unwrap_or(p).to_string_lossy().replace('\\', "/"))
        .collect();

    // No structured data files, but the raw dir holds other files (e.g. .txt / code): this is a
    // raw-files dataset — recommend a `files` source rather than reporting "unknown, columns []".
    if data_files.is_empty() {
        let has_any_file = std::fs::read_dir(raw_dir)
            .map(|rd| rd.flatten().any(|e| e.path().is_file()))
            .unwrap_or(false);
        if has_any_file {
            return Ok(NormalizeReport {
                schema: Schema::RawFiles, rows_in: 0, rows_out: 0, wrote_output: false, dropped: vec![], files,
            });
        }
    }

    // Detect schema from the first READABLE, non-empty data file (skip unreadable/empty leading files).
    let sample = data_files
        .iter()
        .find_map(|f| {
            let recs = crate::records::read_records(f, Some(50)).ok()?;
            if recs.is_empty() { None } else { Some(recs) }
        })
        .unwrap_or_default();
    let schema = detect_schema(&sample, override_map);

    if matches!(schema, Schema::Unknown(_) | Schema::TextOnly | Schema::RawFiles) {
        return Ok(NormalizeReport { schema, rows_in: 0, rows_out: 0, wrote_output: false, dropped: vec![], files });
    }

    if let Some(parent) = out_file.parent() {
        std::fs::create_dir_all(parent)?;
    }
    use std::io::Write;
    let mut writer = std::io::BufWriter::new(std::fs::File::create(out_file)?);
    let mut rows_in = 0usize;
    let mut rows_out = 0usize;
    let mut drops: Vec<DropReason> = Vec::new();
    'files: for f in &data_files {
        let records = match crate::records::read_records(f, None) {
            Ok(r) => r,
            Err(e) => {
                eprintln!("kibble: skipping unreadable data file {}: {e}", f.display());
                continue;
            }
        };
        for rec in records {
            if let Some(cap) = max_rows {
                if rows_out >= cap {
                    break 'files;
                }
            }
            rows_in += 1;
            match map_record(&rec, &schema, override_map) {
                Ok(row) => {
                    let line = serde_json::to_string(&serde_json::json!({ "messages": row.messages.iter().map(|m| serde_json::json!({"role": m.role, "content": m.content})).collect::<Vec<_>>() }))?;
                    writer.write_all(line.as_bytes())?;
                    writer.write_all(b"\n")?;
                    rows_out += 1;
                }
                Err(reason) => drops.push(reason),
            }
        }
    }
    writer.flush()?;
    let wrote_output = rows_out > 0;
    if !wrote_output {
        // No usable rows — remove the empty file so a source never points at nothing.
        drop(writer);
        let _ = std::fs::remove_file(out_file);
    }
    Ok(NormalizeReport { schema, rows_in, rows_out, wrote_output, dropped: tally(drops), files })
}

#[cfg(test)]
mod tests {
    use super::*;
    use serde_json::json;

    #[test]
    fn detects_messages() {
        let s = vec![json!({"messages":[{"role":"user","content":"hi"},{"role":"assistant","content":"yo"}]})];
        assert_eq!(detect_schema(&s, None), Schema::Messages);
    }

    #[test]
    fn detects_sharegpt() {
        let s = vec![json!({"conversations":[{"from":"human","value":"hi"},{"from":"gpt","value":"yo"}]})];
        assert_eq!(detect_schema(&s, None), Schema::ShareGpt);
    }

    #[test]
    fn detects_alpaca() {
        let s = vec![json!({"instruction":"do x","input":"","output":"done"})];
        assert_eq!(detect_schema(&s, None), Schema::Alpaca);
    }

    #[test]
    fn detects_prompt_response() {
        let s = vec![json!({"prompt":"q","response":"a"})];
        assert_eq!(detect_schema(&s, None), Schema::PromptResponse);
    }

    #[test]
    fn detects_question_answer() {
        let s = vec![json!({"question":"q","answer":"a"})];
        assert_eq!(detect_schema(&s, None), Schema::QuestionAnswer);
    }

    #[test]
    fn detects_text_only() {
        let s = vec![json!({"text":"just a blob of prose"})];
        assert_eq!(detect_schema(&s, None), Schema::TextOnly);
    }

    #[test]
    fn detects_unknown_carries_columns() {
        let s = vec![json!({"foo":1,"bar":2})];
        match detect_schema(&s, None) {
            Schema::Unknown(cols) => {
                assert!(cols.contains(&"foo".to_string()));
                assert!(cols.contains(&"bar".to_string()));
            }
            other => panic!("expected Unknown, got {other:?}"),
        }
    }

    #[test]
    fn override_forces_prompt_response_shape() {
        let s = vec![json!({"q":"x","a":"y"})];
        let m = ColumnMap { user: "q".into(), assistant: "a".into(), system: None };
        // With an override, detection returns PromptResponse (the generic user/assistant mapping).
        assert_eq!(detect_schema(&s, Some(&m)), Schema::PromptResponse);
    }

    #[test]
    fn label_strings() {
        assert_eq!(Schema::Messages.label(), "messages");
        assert_eq!(Schema::ShareGpt.label(), "sharegpt");
        assert_eq!(Schema::Unknown(vec![]).label(), "unknown");
        assert_eq!(Schema::TextOnly.label(), "text-only");
    }

    #[test]
    fn maps_messages_passthrough() {
        let v = json!({"messages":[{"role":"user","content":"hi"},{"role":"assistant","content":"yo"}]});
        let row = map_record(&v, &Schema::Messages, None).unwrap();
        assert_eq!(row.messages.len(), 2);
        assert_eq!(row.messages[0].role, "user");
        assert_eq!(row.messages[1].role, "assistant");
    }

    #[test]
    fn maps_sharegpt_roles() {
        let v = json!({"conversations":[{"from":"human","value":"hi"},{"from":"gpt","value":"yo"}]});
        let row = map_record(&v, &Schema::ShareGpt, None).unwrap();
        assert_eq!(row.messages[0].role, "user");
        assert_eq!(row.messages[1].role, "assistant");
        assert_eq!(row.messages[1].content, "yo");
    }

    #[test]
    fn maps_alpaca_with_and_without_input() {
        let with = json!({"instruction":"sum","input":"1 2","output":"3"});
        let row = map_record(&with, &Schema::Alpaca, None).unwrap();
        assert_eq!(row.messages[0].role, "user");
        assert!(row.messages[0].content.contains("sum"));
        assert!(row.messages[0].content.contains("1 2"));
        assert_eq!(row.messages[1].content, "3");
        let without = json!({"instruction":"hi","output":"yo"});
        let row2 = map_record(&without, &Schema::Alpaca, None).unwrap();
        assert_eq!(row2.messages[0].content, "hi");
    }

    #[test]
    fn maps_prompt_response_and_qa() {
        let pr = json!({"prompt":"q","response":"a"});
        assert_eq!(map_record(&pr, &Schema::PromptResponse, None).unwrap().messages.len(), 2);
        let qa = json!({"question":"q","answer":"a"});
        assert_eq!(map_record(&qa, &Schema::QuestionAnswer, None).unwrap().messages.len(), 2);
    }

    #[test]
    fn override_map_pulls_named_columns() {
        let v = json!({"q":"hello","a":"world"});
        let m = ColumnMap { user: "q".into(), assistant: "a".into(), system: None };
        let row = map_record(&v, &Schema::PromptResponse, Some(&m)).unwrap();
        assert_eq!(row.messages[0].content, "hello");
        assert_eq!(row.messages[1].content, "world");
    }

    #[test]
    fn missing_field_drops_with_reason() {
        let v = json!({"prompt":"q"});
        assert_eq!(map_record(&v, &Schema::PromptResponse, None), Err(DropReason::MissingField("response".into())));
    }

    #[test]
    fn empty_content_drops() {
        let v = json!({"prompt":"   ","response":"a"});
        assert_eq!(map_record(&v, &Schema::PromptResponse, None), Err(DropReason::Empty));
    }

    #[test]
    fn text_only_is_not_chat() {
        let v = json!({"text":"blob"});
        assert_eq!(map_record(&v, &Schema::TextOnly, None), Err(DropReason::NotChat));
    }

    #[test]
    fn parse_column_map_forms() {
        let m = parse_column_map("user=q,assistant=a").unwrap();
        assert_eq!(m.user, "q");
        assert_eq!(m.assistant, "a");
        assert!(m.system.is_none());
        let m2 = parse_column_map("user=q,assistant=a,system=s").unwrap();
        assert_eq!(m2.system.as_deref(), Some("s"));
        assert!(parse_column_map("user=q").is_err()); // assistant required
    }

    fn tmp_dir(tag: &str) -> std::path::PathBuf {
        let d = std::env::temp_dir().join(format!("kibble_norm_{}_{}", tag, std::process::id()));
        let _ = std::fs::remove_dir_all(&d);
        std::fs::create_dir_all(&d).unwrap();
        d
    }

    #[test]
    fn normalize_dataset_writes_messages_and_report() {
        let raw = tmp_dir("ok");
        std::fs::write(raw.join("train.jsonl"), "{\"prompt\":\"q1\",\"response\":\"a1\"}\n{\"prompt\":\"q2\",\"response\":\"a2\"}\n").unwrap();
        let out = raw.join("normalized").join("data.jsonl");
        let rep = normalize_dataset(&raw, &out, None, None).unwrap();
        assert_eq!(rep.schema, Schema::PromptResponse);
        assert_eq!(rep.rows_in, 2);
        assert_eq!(rep.rows_out, 2);
        assert!(rep.wrote_output);
        assert!(out.exists());
        let written = std::fs::read_to_string(&out).unwrap();
        assert_eq!(written.lines().count(), 2);
        let first: serde_json::Value = serde_json::from_str(written.lines().next().unwrap()).unwrap();
        assert_eq!(first["messages"][0]["content"], "q1");
    }

    #[test]
    fn normalize_dataset_unknown_writes_no_output() {
        let raw = tmp_dir("unk");
        std::fs::write(raw.join("train.jsonl"), "{\"foo\":1,\"bar\":2}\n").unwrap();
        let out = raw.join("normalized").join("data.jsonl");
        let rep = normalize_dataset(&raw, &out, None, None).unwrap();
        assert!(matches!(rep.schema, Schema::Unknown(_)));
        assert!(!rep.wrote_output);
        assert_eq!(rep.rows_out, 0);
        assert!(!out.exists());
    }

    #[test]
    fn normalize_dataset_skips_unreadable_file_and_uses_the_good_one() {
        let raw = tmp_dir("badfile");
        // an unreadable/corrupt parquet (garbage bytes) sorts before the good jsonl
        std::fs::write(raw.join("aaa_corrupt.parquet"), b"not a real parquet file").unwrap();
        std::fs::write(raw.join("train.jsonl"), "{\"prompt\":\"q1\",\"response\":\"a1\"}\n{\"prompt\":\"q2\",\"response\":\"a2\"}\n").unwrap();
        let out = raw.join("normalized").join("data.jsonl");
        let rep = normalize_dataset(&raw, &out, None, None).unwrap();
        assert_eq!(rep.schema, Schema::PromptResponse); // detected from the good file, not the corrupt one
        assert_eq!(rep.rows_out, 2);
        assert!(rep.wrote_output);
        assert!(out.exists());
    }

    #[test]
    fn collect_data_files_finds_jsonl_and_parquet_recursively() {
        let raw = tmp_dir("collect");
        std::fs::write(raw.join("train.jsonl"), "{}\n").unwrap();
        std::fs::create_dir_all(raw.join("data")).unwrap();
        std::fs::write(raw.join("data").join("t.parquet"), "x").unwrap();
        std::fs::write(raw.join("README.md"), "hi").unwrap();
        std::fs::write(raw.join(".gitattributes"), "x").unwrap();
        let files = collect_data_files(&raw);
        assert_eq!(files.len(), 2);
    }

    #[test]
    fn detects_pair_synonyms() {
        let s = vec![json!({"prompt":"q","full_answer":"a"})];
        assert_eq!(detect_schema(&s, None), Schema::Pair("prompt".into(), "full_answer".into()));
        let s2 = vec![json!({"query":"q","completion":"a"})];
        assert_eq!(detect_schema(&s2, None), Schema::Pair("query".into(), "completion".into()));
    }

    #[test]
    fn pair_does_not_shadow_named_schemas() {
        // instruction+output is Alpaca (named), not Pair
        let s = vec![json!({"instruction":"do x","output":"done"})];
        assert_eq!(detect_schema(&s, None), Schema::Alpaca);
        // prompt+response stays PromptResponse (named)
        let s2 = vec![json!({"prompt":"q","response":"a"})];
        assert_eq!(detect_schema(&s2, None), Schema::PromptResponse);
    }

    #[test]
    fn classification_columns_stay_unknown() {
        let s = vec![json!({"text":"blah","label":3})];
        assert!(matches!(detect_schema(&s, None), Schema::Unknown(_)));
    }

    #[test]
    fn maps_pair_columns() {
        let v = json!({"prompt":"hello","full_answer":"world"});
        let row = map_record(&v, &Schema::Pair("prompt".into(), "full_answer".into()), None).unwrap();
        assert_eq!(row.messages[0].content, "hello");
        assert_eq!(row.messages[1].content, "world");
        // missing assistant column drops with MissingField
        let v2 = json!({"prompt":"hello"});
        assert_eq!(map_record(&v2, &Schema::Pair("prompt".into(), "full_answer".into()), None),
                   Err(DropReason::MissingField("full_answer".into())));
    }

    #[test]
    fn label_pair_and_owned() {
        assert_eq!(Schema::Messages.label(), "messages".to_string());
        assert_eq!(Schema::Pair("prompt".into(), "full_answer".into()).label(), "pair:prompt/full_answer".to_string());
    }

    #[test]
    fn raw_text_only_dir_is_rawfiles() {
        let raw = tmp_dir("rawtxt");
        std::fs::write(raw.join("train.txt"), "just prose, no structured data\n").unwrap();
        std::fs::write(raw.join("README.md"), "readme").unwrap();
        let out = raw.join("normalized").join("data.jsonl");
        let rep = normalize_dataset(&raw, &out, None, None).unwrap();
        assert_eq!(rep.schema, Schema::RawFiles);
        assert!(!rep.wrote_output);
        assert!(!out.exists());
    }

    #[test]
    fn empty_raw_dir_is_unknown_not_rawfiles() {
        let raw = tmp_dir("emptyraw");
        let out = raw.join("normalized").join("data.jsonl");
        let rep = normalize_dataset(&raw, &out, None, None).unwrap();
        assert!(matches!(rep.schema, Schema::Unknown(_)));
    }

    #[test]
    fn max_rows_caps_output() {
        let raw = tmp_dir("cap");
        std::fs::write(raw.join("train.jsonl"),
            "{\"prompt\":\"q1\",\"response\":\"a1\"}\n{\"prompt\":\"q2\",\"response\":\"a2\"}\n{\"prompt\":\"q3\",\"response\":\"a3\"}\n").unwrap();
        let out = raw.join("normalized").join("data.jsonl");
        let rep = normalize_dataset(&raw, &out, None, Some(2)).unwrap();
        assert_eq!(rep.rows_out, 2);
        assert_eq!(std::fs::read_to_string(&out).unwrap().lines().count(), 2);
    }

    #[test]
    fn rawfiles_label() {
        assert_eq!(Schema::RawFiles.label(), "raw-files".to_string());
    }
}