datui-lib 0.4.4

//! Export modal state and focus management.

use crate::CompressionFormat;
use crate::widgets::text_input::TextInput;
use polars::prelude::{LazyFrame, PolarsResult};

#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
pub enum ExportFormat {
    #[default]
    Csv,
    /// CSV with the tab as its delimiter.
    Tsv,
    /// CSV with `|` as its delimiter.
    Psv,
    Parquet,
    Json,
    Ndjson,
    /// Arrow IPC / Feather v2
    Ipc,
    Avro,
}

impl ExportFormat {
    pub const ALL: [Self; 8] = [
        Self::Csv,
        Self::Tsv,
        Self::Psv,
        Self::Parquet,
        Self::Json,
        Self::Ndjson,
        Self::Ipc,
        Self::Avro,
    ];

    pub const fn as_str(self) -> &'static str {
        match self {
            Self::Csv => "CSV",
            Self::Tsv => "TSV",
            Self::Psv => "PSV",
            Self::Parquet => "Parquet",
            Self::Json => "JSON",
            Self::Ndjson => "NDJSON",
            Self::Ipc => "Arrow",
            Self::Avro => "Avro",
        }
    }

    pub fn extension(self) -> &'static str {
        match self {
            Self::Csv => "csv",
            Self::Tsv => "tsv",
            Self::Psv => "psv",
            Self::Parquet => "parquet",
            Self::Json => "json",
            Self::Ndjson => "jsonl",
            Self::Ipc => "arrow",
            Self::Avro => "avro",
        }
    }

    pub fn from_extension(ext: &str) -> Option<Self> {
        match ext.to_lowercase().as_str() {
            "csv" => Some(Self::Csv),
            "tsv" => Some(Self::Tsv),
            "psv" => Some(Self::Psv),
            "parquet" => Some(Self::Parquet),
            "json" => Some(Self::Json),
            "ndjson" | "jsonl" => Some(Self::Ndjson),
            "arrow" | "ipc" | "feather" => Some(Self::Ipc),
            "avro" => Some(Self::Avro),
            _ => None,
        }
    }

    /// Whether the format stores list, array and struct columns as they are.
    /// The others get them as JSON text (`nested_json`).
    pub fn holds_nesting(self) -> bool {
        !self.is_delimited()
    }

    /// CSV and its presets, which write delimited text.
    pub fn is_delimited(self) -> bool {
        matches!(self, Self::Csv | Self::Tsv | Self::Psv)
    }

    /// The delimiter a preset sets; `None` for CSV, whose delimiter is the user's.
    pub fn preset_delimiter(self) -> Option<u8> {
        match self {
            Self::Tsv => Some(b'\t'),
            Self::Psv => Some(b'|'),
            _ => None,
        }
    }

    /// `lf` as this format can write it (planned, not run): binary as base64 and dates
    /// as text for CSV and JSON, nested columns as JSON and durations as ISO 8601 for
    /// CSV, and types Avro lacks cast to ones it has.
    pub fn prepare(self, lf: LazyFrame) -> PolarsResult<LazyFrame> {
        match self {
            Self::Csv | Self::Tsv | Self::Psv => crate::export::nested_json::lazy_as_json(lf),
            Self::Json | Self::Ndjson => crate::export::nested_json::lazy_for_json(lf),
            Self::Avro => crate::export::avro_types::lazy_for_avro(lf),
            Self::Parquet | Self::Ipc => Ok(lf),
        }
    }

    pub fn supports_compression(self) -> bool {
        self.is_delimited() || matches!(self, Self::Json | Self::Ndjson)
    }

    /// The format a path's extension names, looking through a trailing compression
    /// extension so `out.csv.gz` still names CSV. None for a bare or unknown extension.
    pub fn from_path(path: &str) -> Option<Self> {
        let path = std::path::Path::new(path);
        let ext = path.extension()?.to_str()?;
        if let Some(format) = Self::from_extension(ext) {
            return Some(format);
        }
        if matches!(ext.to_lowercase().as_str(), "gz" | "zst" | "bz2" | "xz") {
            let stem = path.file_stem()?.to_str()?;
            return Self::from_extension(stem.rsplit('.').next()?);
        }
        None
    }
}

/// The export form's fields.
#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
pub enum ExportFocus {
    #[default]
    FormatSelector,
    PathInput,
    CsvDelimiter,
    CsvIncludeHeader,
    /// The compression of the chosen format; offered by the formats that take one.
    Compression,
    /// A column naming each row's source file. Offered only when the files disagree,
    /// where it tells a null from an absent cell downstream.
    SourceFile,
}

/// The compressions a delimited or JSON export steps through, in order.
pub const COMPRESSION_OPTIONS: [Option<CompressionFormat>; 5] = [
    None,
    Some(CompressionFormat::Gzip),
    Some(CompressionFormat::Zstd),
    Some(CompressionFormat::Bzip2),
    Some(CompressionFormat::Xz),
];

pub struct ExportModal {
    pub focus: ExportFocus,
    pub selected_format: ExportFormat,
    pub path_input: TextInput,
    pub csv_delimiter_input: TextInput,
    pub csv_include_header: bool,
    /// Add a column naming the file each row came from. See `ExportFocus::SourceFile`.
    pub source_file: bool,
    /// Whether this dataset has files to name. Set when the modal opens.
    pub offer_source_file: bool,
    /// Whether the view has list, array or struct columns, which a format
    /// without nesting writes as JSON; the dialog says so. Set when it opens.
    pub nested_columns: bool,
    /// Whether a column or struct field of the view has a name Avro does not
    /// allow; the dialog says those are renamed. Set when it opens.
    pub avro_renames: bool,
    pub csv_compression: Option<CompressionFormat>,
    pub json_compression: Option<CompressionFormat>,
    pub ndjson_compression: Option<CompressionFormat>,
    pub history_limit: usize,
    /// Why the form cannot export, or why its last write failed, said inline on
    /// its own status line. Cleared by typing in the path.
    pub path_error: Option<String>,
    /// The counts the dialog writes, when it was opened from Value Counts. Kept after
    /// the form closes while the write runs, so a failed one reopens on them.
    pub(crate) counts: Option<polars::prelude::DataFrame>,
}

impl ExportModal {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn open(
        &mut self,
        default_format: Option<ExportFormat>,
        history_limit: usize,
        theme: &crate::config::Theme,
        file_delimiter: Option<u8>,
    ) {
        self.focus = ExportFocus::PathInput;
        self.history_limit = history_limit;
        if let Some(format) = default_format {
            self.selected_format = format;
        }
        // Ctrl+P / Ctrl+N recall the paths exported to before.
        self.path_input = TextInput::new()
            .with_history("export_path".to_string())
            .with_history_limit(history_limit)
            .with_theme(theme);
        self.path_input.clear();
        self.csv_delimiter_input = TextInput::new()
            .with_history_limit(history_limit)
            .with_theme(theme);
        // `--delimiter` if the file was read with one, else a comma.
        let delimiter_char = file_delimiter.unwrap_or(b',');
        self.csv_delimiter_input
            .suggest(format!("{}", delimiter_char as char));
        self.csv_include_header = true;
        self.source_file = false;
        self.offer_source_file = false;
        self.nested_columns = false;
        self.avro_renames = false;
        self.csv_compression = None;
        self.json_compression = None;
        self.ndjson_compression = None;
        self.path_error = None;
        self.counts = None;
    }

    pub fn close(&mut self) {
        self.focus = ExportFocus::FormatSelector;
        self.path_input.clear();
        self.path_error = None;
        self.counts = None;
    }

    /// The dataset the counts were of is left: a write still running exports them, but
    /// a failure reopens on the view, not on counts of a dataset gone.
    pub(crate) fn forget_counts(&mut self) {
        self.counts = None;
    }

    /// Follow the typed path's extension with the format choice, so `out.csv` never
    /// receives Parquet bytes. An unknown extension leaves the format alone; runs only
    /// when the path changes, so a format picked after typing stands.
    pub fn sync_format_to_path(&mut self) {
        let value = self.path_input.value().trim().to_string();
        if let Some(format) = ExportFormat::from_path(&value) {
            self.selected_format = format;
            // A trailing compression extension is part of what the path asks for:
            // `out.csv.gz` left at Compression: None writes plain bytes to a .gz name.
            if format.supports_compression()
                && let Some(comp) = CompressionFormat::from_extension(std::path::Path::new(&value))
            {
                self.set_compression_for(format, Some(comp));
            }
        }
    }

    /// Offer `stem` with the format's extension as the path: Enter takes it, typing
    /// replaces it, and stepping the format carries it along.
    pub fn suggest_path(&mut self, stem: &str) {
        let path = format!("{stem}.{}", self.selected_format.extension());
        self.path_input.suggest(path);
    }

    /// The reverse of [`Self::sync_format_to_path`]: a format picked after typing
    /// rewrites the path's format extension. A path with no format extension is left
    /// alone; a compression suffix stays only if the new format supports one.
    pub fn sync_path_to_format(&mut self) {
        let suggested = self.path_input.is_suggested();
        let value = self.path_input.value().trim().to_string();
        if value.is_empty() || ExportFormat::from_path(&value).is_none() {
            return;
        }
        let path = std::path::Path::new(&value);
        let compression = CompressionFormat::from_extension(path)
            .filter(|_| self.selected_format.supports_compression());
        // Strip the compression suffix, then the format extension, textually:
        // Path::set_extension would also eat the `v2` of `data.v2`.
        let mut base = value.as_str();
        if CompressionFormat::from_extension(path).is_some()
            && let Some((rest, _)) = base.rsplit_once('.')
        {
            base = rest;
        }
        if let Some((rest, ext)) = base.rsplit_once('.')
            && ExportFormat::from_extension(ext).is_some()
        {
            base = rest;
        }
        let new_path = match compression {
            Some(comp) => format!(
                "{base}.{}.{}",
                self.selected_format.extension(),
                comp.extension()
            ),
            None => format!("{base}.{}", self.selected_format.extension()),
        };
        if suggested {
            self.path_input.suggest(new_path);
        } else {
            self.path_input.set_value(new_path);
        }
    }

    /// Set the compression field the given format reads at export time.
    fn set_compression_for(&mut self, format: ExportFormat, comp: Option<CompressionFormat>) {
        match format {
            ExportFormat::Csv | ExportFormat::Tsv | ExportFormat::Psv => {
                self.csv_compression = comp
            }
            ExportFormat::Json => self.json_compression = comp,
            ExportFormat::Ndjson => self.ndjson_compression = comp,
            ExportFormat::Parquet | ExportFormat::Ipc | ExportFormat::Avro => {}
        }
    }

    /// The compression the chosen format writes with; `None` for one that takes none.
    pub fn compression(&self) -> Option<CompressionFormat> {
        match self.selected_format {
            ExportFormat::Csv | ExportFormat::Tsv | ExportFormat::Psv => self.csv_compression,
            ExportFormat::Json => self.json_compression,
            ExportFormat::Ndjson => self.ndjson_compression,
            ExportFormat::Parquet | ExportFormat::Ipc | ExportFormat::Avro => None,
        }
    }

    /// Step the chosen format's compression through [`COMPRESSION_OPTIONS`].
    pub fn step_compression(&mut self, delta: i8) {
        let next = crate::app::form::step_value(&COMPRESSION_OPTIONS, self.compression(), delta);
        self.set_compression_for(self.selected_format, next);
    }

    /// Step the format, carrying the typed path's extension with it.
    pub fn step_format(&mut self, delta: i8) {
        self.selected_format =
            crate::app::form::step_value(&ExportFormat::ALL, self.selected_format, delta);
        self.sync_path_to_format();
        // A format without a delimiter row or compression takes focus off it.
        crate::app::form::Form::settle_focus(self);
    }

    /// The fields in Tab order: a list, since options vary by format and dataset.
    pub fn focus_order(&self) -> Vec<ExportFocus> {
        crate::app::form::Form::fields(self)
            .into_iter()
            .map(|(field, _)| field)
            .collect()
    }
}

impl crate::app::form::Form for ExportModal {
    type Field = ExportFocus;

    fn fields(&self) -> Vec<(ExportFocus, crate::app::form::FieldKind)> {
        use crate::app::form::FieldKind::{Checkbox, Choice, Text};
        let mut fields = vec![
            (ExportFocus::FormatSelector, Choice),
            (ExportFocus::PathInput, Text),
        ];
        if self.selected_format == ExportFormat::Csv {
            // The presets say the delimiter.
            fields.push((ExportFocus::CsvDelimiter, Text));
        }
        if self.selected_format.is_delimited() {
            fields.push((ExportFocus::CsvIncludeHeader, Checkbox));
        }
        if self.selected_format.supports_compression() {
            fields.push((ExportFocus::Compression, Choice));
        }
        if self.offer_source_file {
            fields.push((ExportFocus::SourceFile, Checkbox));
        }
        fields
    }

    fn focused(&self) -> ExportFocus {
        self.focus
    }

    fn set_focused(&mut self, field: ExportFocus) {
        self.focus = field;
        self.path_input.set_focused(field == ExportFocus::PathInput);
        self.csv_delimiter_input
            .set_focused(field == ExportFocus::CsvDelimiter);
    }
}

impl Default for ExportModal {
    fn default() -> Self {
        Self {
            focus: ExportFocus::FormatSelector,
            selected_format: ExportFormat::Csv,
            path_input: TextInput::new(),
            csv_delimiter_input: TextInput::new(),
            csv_include_header: true,
            source_file: false,
            offer_source_file: false,
            nested_columns: false,
            avro_renames: false,
            csv_compression: None,
            json_compression: None,
            ndjson_compression: None,
            history_limit: 1000,
            path_error: None,
            counts: None,
        }
    }
}

#[cfg(test)]
mod tests;