rust_excel_core/

workbook.rs

use std::io::Cursor;
use std::path::Path;

use calamine::{open_workbook_from_rs, Reader, Xlsx};
use umya_spreadsheet::{reader, writer, Spreadsheet, Worksheet};

use crate::cell::from_calamine_data;
use crate::error::{ExcelError, ExcelResult};
use crate::types::{CellValue, RangeData, SheetInfo, WorkbookInfo};

/// A wrapper around an Excel workbook providing a clean API.
pub struct Workbook {
    inner: Spreadsheet,
    /// Cached original bytes for fast calamine reads. Cleared on mutation.
    cached_bytes: Option<Vec<u8>>,
}

impl Workbook {
    /// Create a new empty workbook with one default sheet.
    pub fn new() -> Self {
        Self {
            inner: umya_spreadsheet::new_file(),
            cached_bytes: None,
        }
    }

    /// Open a workbook from a file path.
    pub fn open(path: impl AsRef<Path>) -> ExcelResult<Self> {
        let spreadsheet =
            reader::xlsx::read(path).map_err(|e| ExcelError::Spreadsheet(e.to_string()))?;
        Ok(Self {
            inner: spreadsheet,
            cached_bytes: None,
        })
    }

    /// Open a workbook from bytes in memory.
    pub fn open_from_bytes(bytes: &[u8]) -> ExcelResult<Self> {
        let cursor = Cursor::new(bytes);
        let spreadsheet = reader::xlsx::read_reader(cursor, true)
            .map_err(|e| ExcelError::Spreadsheet(e.to_string()))?;
        Ok(Self {
            inner: spreadsheet,
            cached_bytes: None,
        })
    }

    /// Open a workbook from bytes, caching the bytes for fast calamine reads.
    ///
    /// Use this on upload: the original bytes are kept so `read_sheet_data_fast()`
    /// can use calamine directly without re-serializing through umya.
    pub fn open_from_bytes_cached(bytes: Vec<u8>) -> ExcelResult<Self> {
        let cursor = Cursor::new(&bytes);
        let spreadsheet = reader::xlsx::read_reader(cursor, true)
            .map_err(|e| ExcelError::Spreadsheet(e.to_string()))?;
        Ok(Self {
            inner: spreadsheet,
            cached_bytes: Some(bytes),
        })
    }

    /// Save the workbook to a file path.
    pub fn save(&self, path: impl AsRef<Path>) -> ExcelResult<()> {
        writer::xlsx::write(&self.inner, path)
            .map_err(|e| ExcelError::Spreadsheet(e.to_string()))?;
        Ok(())
    }

    /// Serialize the workbook to bytes.
    pub fn save_to_bytes(&self) -> ExcelResult<Vec<u8>> {
        let mut buf = Vec::new();
        writer::xlsx::write_writer(&self.inner, Cursor::new(&mut buf))
            .map_err(|e| ExcelError::Spreadsheet(e.to_string()))?;
        Ok(buf)
    }

    /// Get summary info about the workbook.
    pub fn info(&self) -> WorkbookInfo {
        let sheets: Vec<SheetInfo> = self
            .inner
            .get_sheet_collection()
            .iter()
            .enumerate()
            .map(|(i, s)| SheetInfo {
                name: s.get_name().to_string(),
                index: i,
            })
            .collect();

        WorkbookInfo {
            sheet_count: sheets.len(),
            sheets,
        }
    }

    /// Get the names of all sheets.
    pub fn sheet_names(&self) -> Vec<String> {
        self.inner
            .get_sheet_collection()
            .iter()
            .map(|s| s.get_name().to_string())
            .collect()
    }

    /// Add a new empty sheet with the given name.
    pub fn add_sheet(&mut self, name: &str) -> ExcelResult<()> {
        if self.inner.get_sheet_by_name(name).is_some() {
            return Err(ExcelError::SheetAlreadyExists(name.to_string()));
        }
        self.inner
            .new_sheet(name)
            .map_err(|e| ExcelError::Spreadsheet(e.to_string()))?;
        Ok(())
    }

    /// Remove a sheet by name.
    pub fn remove_sheet(&mut self, name: &str) -> ExcelResult<()> {
        self.inner
            .remove_sheet_by_name(name)
            .map_err(|e| ExcelError::Spreadsheet(e.to_string()))?;
        Ok(())
    }

    /// Rename a sheet.
    pub fn rename_sheet(&mut self, old_name: &str, new_name: &str) -> ExcelResult<()> {
        let sheet = self
            .inner
            .get_sheet_by_name_mut(old_name)
            .ok_or_else(|| ExcelError::SheetNotFound(old_name.to_string()))?;
        sheet.set_name(new_name);
        Ok(())
    }

    /// Get an immutable reference to a sheet by name.
    pub fn get_sheet(&self, name: &str) -> ExcelResult<&Worksheet> {
        self.inner
            .get_sheet_by_name(name)
            .ok_or_else(|| ExcelError::SheetNotFound(name.to_string()))
    }

    /// Get a mutable reference to a sheet by name.
    pub fn get_sheet_mut(&mut self, name: &str) -> ExcelResult<&mut Worksheet> {
        self.inner
            .get_sheet_by_name_mut(name)
            .ok_or_else(|| ExcelError::SheetNotFound(name.to_string()))
    }

    /// Get a reference to the inner spreadsheet.
    pub fn inner(&self) -> &Spreadsheet {
        &self.inner
    }

    /// Get a mutable reference to the inner spreadsheet.
    pub fn inner_mut(&mut self) -> &mut Spreadsheet {
        &mut self.inner
    }

    /// Mark the workbook as dirty, invalidating cached bytes.
    /// Call this after any mutation (cell write, row/col insert, style change, etc.).
    pub fn mark_dirty(&mut self) {
        self.cached_bytes = None;
    }

    /// Whether cached bytes are available for fast reads.
    pub fn has_cache(&self) -> bool {
        self.cached_bytes.is_some()
    }

    /// Fast bulk read of an entire sheet's data using calamine.
    ///
    /// If cached bytes are available (from upload), reads directly from them
    /// with calamine (~2.4x faster). Otherwise falls back to serializing
    /// through umya first.
    pub fn read_sheet_data(&self, name: &str) -> ExcelResult<RangeData> {
        // Verify the sheet exists first
        if self.inner.get_sheet_by_name(name).is_none() {
            return Err(ExcelError::SheetNotFound(name.to_string()));
        }

        // Fast path: use cached bytes directly with calamine
        if let Some(ref bytes) = self.cached_bytes {
            return Self::read_data_from_bytes(bytes, name);
        }

        // Slow path: serialize via umya then parse with calamine
        let bytes = self.save_to_bytes()?;
        Self::read_data_from_bytes(&bytes, name)
    }

    /// Create a new .xlsx file from raw data using rust_xlsxwriter.
    ///
    /// This bypasses umya-spreadsheet entirely for maximum write performance.
    /// Each tuple in the input is (sheet_name, data). Returns the .xlsx bytes.
    pub fn create_from_data(sheets: Vec<(String, RangeData)>) -> ExcelResult<Vec<u8>> {
        let mut wb = rust_xlsxwriter::Workbook::new();

        for (name, data) in &sheets {
            let ws = wb.add_worksheet();
            ws.set_name(name)
                .map_err(|e| ExcelError::Spreadsheet(e.to_string()))?;

            for (ri, row) in data.rows.iter().enumerate() {
                for (ci, val) in row.iter().enumerate() {
                    let r = ri as u32;
                    let c = ci as u16;
                    match val {
                        CellValue::String(s) => {
                            ws.write_string(r, c, s)
                                .map_err(|e| ExcelError::Spreadsheet(e.to_string()))?;
                        }
                        CellValue::Number(n) => {
                            ws.write_number(r, c, *n)
                                .map_err(|e| ExcelError::Spreadsheet(e.to_string()))?;
                        }
                        CellValue::Boolean(b) => {
                            ws.write_boolean(r, c, *b)
                                .map_err(|e| ExcelError::Spreadsheet(e.to_string()))?;
                        }
                        CellValue::Formula(f) => {
                            ws.write_formula(r, c, f.as_str())
                                .map_err(|e| ExcelError::Spreadsheet(e.to_string()))?;
                        }
                        CellValue::Empty => {}
                    }
                }
            }
        }

        wb.save_to_buffer()
            .map_err(|e| ExcelError::Spreadsheet(e.to_string()))
    }

    /// Serialize the workbook to bytes using rust_xlsxwriter for speed.
    ///
    /// Walks the umya in-memory model and writes cell data via rust_xlsxwriter
    /// (~1.4x faster than umya's native writer). Falls back to umya on error.
    /// Note: styles/merges/formatting may not be fully preserved — use
    /// `save_to_bytes()` when full fidelity is needed.
    pub fn save_to_bytes_fast(&self) -> ExcelResult<Vec<u8>> {
        let sheets: Vec<(String, RangeData)> = self
            .inner
            .get_sheet_collection()
            .iter()
            .map(|ws| {
                let name = ws.get_name().to_string();
                let (min_row, max_row, min_col, max_col) = Self::sheet_used_range(ws);
                let mut rows = Vec::new();
                for r in min_row..=max_row {
                    let mut row = Vec::new();
                    for c in min_col..=max_col {
                        let val = match ws.get_cell((c, r)) {
                            Some(cell) => crate::cell::read_cell_value(cell),
                            None => CellValue::Empty,
                        };
                        row.push(val);
                    }
                    rows.push(row);
                }
                (
                    name,
                    RangeData {
                        start_row: min_row,
                        start_col: min_col,
                        end_row: max_row,
                        end_col: max_col,
                        rows,
                    },
                )
            })
            .collect();

        match Self::create_from_data(sheets) {
            Ok(bytes) => Ok(bytes),
            Err(_) => self.save_to_bytes(), // fallback to umya
        }
    }

    /// Get the used range of a worksheet (min_row, max_row, min_col, max_col).
    fn sheet_used_range(ws: &Worksheet) -> (u32, u32, u32, u32) {
        let cells = ws.get_cell_collection();
        if cells.is_empty() {
            return (1, 1, 1, 1);
        }
        let mut min_row = u32::MAX;
        let mut max_row = 0u32;
        let mut min_col = u32::MAX;
        let mut max_col = 0u32;
        for cell in cells {
            let r = *cell.get_coordinate().get_row_num();
            let c = *cell.get_coordinate().get_col_num();
            min_row = min_row.min(r);
            max_row = max_row.max(r);
            min_col = min_col.min(c);
            max_col = max_col.max(c);
        }
        (min_row, max_row, min_col, max_col)
    }

    /// Read sheet data directly from raw .xlsx bytes using calamine only.
    ///
    /// This is the true fast path — no umya-spreadsheet involved at all.
    /// Use this when you already have raw bytes (e.g. from an upload) and
    /// just need to extract cell data without loading into the mutable model.
    pub fn read_data_from_bytes(bytes: &[u8], sheet_name: &str) -> ExcelResult<RangeData> {
        let cursor = Cursor::new(bytes);
        let mut cal_wb: Xlsx<_> = open_workbook_from_rs(cursor)
            .map_err(|e: calamine::XlsxError| ExcelError::Spreadsheet(e.to_string()))?;

        let range = cal_wb
            .worksheet_range(sheet_name)
            .map_err(|e| ExcelError::Spreadsheet(e.to_string()))?;

        let (start_row, start_col) = range.start().unwrap_or((0, 0));
        let (end_row, end_col) = range.end().unwrap_or((0, 0));

        let rows: Vec<Vec<CellValue>> = range
            .rows()
            .map(|row| row.iter().map(from_calamine_data).collect())
            .collect();

        Ok(RangeData {
            start_row: start_row + 1,
            start_col: start_col + 1,
            end_row: end_row + 1,
            end_col: end_col + 1,
            rows,
        })
    }
}

impl Default for Workbook {
    fn default() -> Self {
        Self::new()
    }
}