plv 0.27.0

use anyhow::Result;
use polars::prelude::*;
use std::io::{BufRead, BufReader, Read};
use std::path::Path;

pub enum FileFormat {
    Csv,
    /// Extension names tabs: `.tsv`, `.tab`.
    Tsv,
    /// Extension names nothing: `.txt`. The delimiter is sniffed from the file.
    Text,
    Parquet,
    /// One JSON object per line: `.jsonl`, `.ndjson`. Read by
    /// [`crate::data::jsonl`] rather than by Polars, so that a nested value
    /// reaches the cell as the file's own bytes.
    JsonLines,
    Unknown,
}

pub fn detect_format(path: &Path) -> FileFormat {
    match path.extension().and_then(|e| e.to_str()) {
        Some("csv") => FileFormat::Csv,
        Some("tsv" | "tab") => FileFormat::Tsv,
        Some("txt") => FileFormat::Text,
        Some("parquet") => FileFormat::Parquet,
        Some("jsonl" | "ndjson") => FileFormat::JsonLines,
        _ => FileFormat::Unknown,
    }
}

pub fn load(path: &Path) -> Result<LazyFrame> {
    let pl_path = PlRefPath::try_from_path(path)?;
    match detect_format(path) {
        FileFormat::Csv | FileFormat::Tsv | FileFormat::Text => {
            let separator = separator(path)?.expect("delimited text has a separator");
            let skip = preamble(path, separator)?.lines;
            Ok(delimited(pl_path, separator, skip)?)
        }
        FileFormat::Parquet => Ok(LazyFrame::scan_parquet(pl_path, Default::default())?),
        // Never reached: `Store::open_file` sends JSONL down its own path
        // before asking for a frame.
        FileFormat::JsonLines => anyhow::bail!("jsonl is not read through polars"),
        FileFormat::Unknown => anyhow::bail!(
            "unsupported file format (use .csv, .tsv, .tab, .txt, .jsonl or .parquet)"
        ),
    }
}

/// The field separator plv reads `path` with, or `None` when the format is not
/// delimited text and so cannot be edited in place.
///
/// Resolved the same way as [`load`], sniffing included, so the write path
/// splits records exactly where the read path did.
pub fn separator(path: &Path) -> Result<Option<u8>> {
    match detect_format(path) {
        FileFormat::Csv => Ok(Some(b',')),
        FileFormat::Tsv => Ok(Some(b'\t')),
        FileFormat::Text => Ok(Some(sniff_delimiter(&read_sample(path)?))),
        FileFormat::Parquet | FileFormat::JsonLines | FileFormat::Unknown => Ok(None),
    }
}

fn delimited(path: PlRefPath, separator: u8, skip_lines: usize) -> PolarsResult<LazyFrame> {
    LazyCsvReader::new(path)
        .with_separator(separator)
        .with_skip_lines(skip_lines)
        .finish()
}

/// The comment lines a delimited file opens with, ahead of its header.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub struct Preamble {
    /// Lines, for Polars, which skips them by counting newlines.
    pub lines: usize,
    /// Bytes, terminators included, for the index and the writer, which skip
    /// them by offset.
    pub bytes: u64,
}

/// Find the comment lines at the top of `path`.
///
/// Only a run at the very start counts, never a `#` line further down: past
/// the header a line is a record, and one whose first field happens to begin
/// with `#` is data. The reader, the row index and the writer all ask this
/// one function, so the three cannot disagree about where the header is.
pub fn preamble(path: &Path, separator: u8) -> Result<Preamble> {
    let file = std::fs::File::open(path)?;
    Ok(preamble_of(BufReader::new(file), separator)?)
}

/// The preamble is every line starting with `#` before the first line that
/// does not — with one exception. When the last of them splits into as many
/// fields as the line after it, it is the header written with a `#` in front,
/// the way VCF and BED write `#chrom`, and it stays the header. A file that
/// is nothing but `#` lines has no preamble: taking it all would leave no
/// header at all.
fn preamble_of(mut reader: impl BufRead, separator: u8) -> std::io::Result<Preamble> {
    let mut found = Preamble::default();
    let mut last = Vec::new();
    let mut line = Vec::new();
    loop {
        line.clear();
        let read = reader.read_until(b'\n', &mut line)?;
        if read == 0 {
            return Ok(Preamble::default());
        }
        if line.starts_with(b"#") {
            found.lines += 1;
            found.bytes += read as u64;
            std::mem::swap(&mut last, &mut line);
            continue;
        }
        if found.lines > 0 {
            let fields = count_unquoted(trim_eol(&line), separator);
            if fields > 0 && count_unquoted(trim_eol(&last), separator) == fields {
                found.lines -= 1;
                found.bytes -= last.len() as u64;
            }
        }
        return Ok(found);
    }
}

fn trim_eol(line: &[u8]) -> &[u8] {
    let line = line.strip_suffix(b"\n").unwrap_or(line);
    line.strip_suffix(b"\r").unwrap_or(line)
}

/// How much of a file to look at when sniffing its delimiter.
const SAMPLE_BYTES: u64 = 64 * 1024;

/// Candidates, in the order ties are broken.
const DELIMITERS: [u8; 4] = *b"\t,;|";

/// Lines to compare before deciding. The header alone can be misleading.
const SAMPLE_LINES: usize = 5;

fn read_sample(path: &Path) -> Result<Vec<u8>> {
    let mut sample = Vec::new();
    std::fs::File::open(path)?
        .take(SAMPLE_BYTES)
        .read_to_end(&mut sample)?;
    Ok(sample)
}

/// Pick the field separator for a file whose extension doesn't name one.
///
/// Counts each candidate outside double-quoted spans on the first few lines and
/// keeps the one that occurs at least once and the *same* number of times on
/// every line — that consistency is what separates a real delimiter from a
/// character the text merely happens to contain. A sample that settles nothing
/// (one line, no repeats, a stray delimiter inside a quoted field spanning a
/// newline) falls back to a tab, which is what `.txt` usually means.
fn sniff_delimiter(sample: &[u8]) -> u8 {
    let lines: Vec<&[u8]> = sample
        .split(|&b| b == b'\n')
        .map(|line| line.strip_suffix(b"\r").unwrap_or(line))
        // A comment says nothing about the delimiter, and prose in one would
        // only spoil the count.
        .filter(|line| !line.is_empty() && !line.starts_with(b"#"))
        .take(SAMPLE_LINES)
        .collect();

    DELIMITERS
        .into_iter()
        .find(|&delim| {
            let mut counts = lines.iter().map(|line| count_unquoted(line, delim));
            match counts.next() {
                Some(first) => first > 0 && counts.all(|n| n == first),
                None => false,
            }
        })
        .unwrap_or(b'\t')
}

/// Occurrences of `delim` in `line` that fall outside a quoted field. A doubled
/// `""` inside a quoted field toggles twice, which lands back where it started.
fn count_unquoted(line: &[u8], delim: u8) -> usize {
    let mut quoted = false;
    line.iter()
        .filter(|&&b| {
            if b == b'"' {
                quoted = !quoted;
            }
            !quoted && b == delim
        })
        .count()
}

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

    #[test]
    fn tab_separated_extensions_split_on_tabs() {
        for ext in ["tsv", "tab", "txt"] {
            let df = load(&write_temp(ext, "name\tcount\na\t1\nb\t2\n"))
                .unwrap()
                .collect()
                .unwrap();
            assert_eq!(df.shape(), (2, 2), "{ext}");
            assert_eq!(df.get_column_names(), ["name", "count"], "{ext}");
        }
    }

    #[test]
    fn comma_separated_txt_is_sniffed() {
        let df = load(&write_temp("txt", "name,count\na,1\nb,2\n"))
            .unwrap()
            .collect()
            .unwrap();
        assert_eq!(df.shape(), (2, 2));
        assert_eq!(df.get_column_names(), ["name", "count"]);
    }

    #[test]
    fn sniffs_each_candidate() {
        assert_eq!(sniff_delimiter(b"a\tb\n1\t2\n"), b'\t');
        assert_eq!(sniff_delimiter(b"a,b\n1,2\n"), b',');
        assert_eq!(sniff_delimiter(b"a;b\n1;2\n"), b';');
        assert_eq!(sniff_delimiter(b"a|b\n1|2\n"), b'|');
    }

    #[test]
    fn ignores_delimiters_inside_quotes() {
        // Every line has one real semicolon; the commas are all quoted away.
        let sample = b"name;note\n\"a,b\";\"x,y,z\"\n\"c,d\";\"p,q,r\"\n";
        assert_eq!(sniff_delimiter(sample), b';');
    }

    #[test]
    fn prose_and_empty_samples_fall_back_to_tab() {
        // Commas appear, but not the same number of times on every line.
        assert_eq!(sniff_delimiter(b"one, two, three\nfour five\n"), b'\t');
        assert_eq!(sniff_delimiter(b""), b'\t');
    }

    #[test]
    fn windows_line_endings_do_not_break_counting() {
        assert_eq!(sniff_delimiter(b"a,b\r\n1,2\r\n"), b',');
    }

    #[test]
    fn leading_comment_lines_are_skipped() {
        let tsv = "# what this is\n# and why\n#\nname\tcount\na\t1\nb\t2\n";
        for ext in ["tsv", "txt"] {
            let df = load(&write_temp(ext, tsv)).unwrap().collect().unwrap();
            assert_eq!(df.shape(), (2, 2), "{ext}");
            assert_eq!(df.get_column_names(), ["name", "count"], "{ext}");
        }
    }

    #[test]
    fn the_preamble_is_the_leading_run_of_comments() {
        let of = |text: &str| preamble_of(text.as_bytes(), b',').unwrap();
        assert_eq!(of("a,b\n1,2\n"), Preamble::default(), "none");
        assert_eq!(of("# note\na,b\n"), Preamble { lines: 1, bytes: 7 });
        assert_eq!(
            of("# one\r\n# two\r\na,b\r\n"),
            Preamble {
                lines: 2,
                bytes: 14
            },
            "crlf counted in bytes"
        );
        assert_eq!(
            of("a,b\n# 1,2\n"),
            Preamble::default(),
            "a # line after the header is a record"
        );
    }

    #[test]
    fn a_header_written_with_a_hash_stays_the_header() {
        let of = |text: &str| preamble_of(text.as_bytes(), b'\t').unwrap();
        // VCF and BED: the header is `#chrom`, and the comments above it are
        // still comments.
        assert_eq!(
            of("## format 4\n#chrom\tpos\nchr1\t10\n"),
            Preamble {
                lines: 1,
                bytes: 12
            }
        );
        assert_eq!(of("#chrom\tpos\nchr1\t10\n"), Preamble::default());
        assert_eq!(
            of("# just\n# comments\n"),
            Preamble::default(),
            "nothing left"
        );
    }

    #[test]
    fn comments_do_not_sway_the_sniffer() {
        assert_eq!(sniff_delimiter(b"# a, b, c\n# d\na;b\n1;2\n"), b';');
    }

    fn write_temp(ext: &str, contents: &str) -> std::path::PathBuf {
        let dir = std::env::temp_dir().join("plv-loader-tests");
        std::fs::create_dir_all(&dir).unwrap();
        let path = dir.join(format!("sample-{ext}-{:x}.{ext}", contents.len()));
        std::fs::write(&path, contents).unwrap();
        path
    }
}