tuitab 0.8.0

use crate::data::column::ColumnMeta;
use crate::data::dataframe::DataFrame;
use crate::types::ColumnType;
use color_eyre::{eyre::eyre, Result};
use polars::prelude::*;
use std::fs::File;
use std::path::Path;

mod arrow;
pub mod db_write;
mod directory;
pub mod doc_io;
mod duckdb;
mod excel;

mod parquet;
mod sqlite;
mod txt;

pub use directory::{load_directory, load_files_list};
pub use duckdb::{
    duckdb_table_names, load_duckdb_overview, load_duckdb_table_by_name, load_duckdb_table_full,
};
pub use excel::{excel_sheet_names, load_excel_overview, load_excel_sheet_by_name};
pub use sqlite::{
    load_sqlite_overview, load_sqlite_table_by_name, load_sqlite_table_full, sqlite_table_names,
};

pub use directory::format_file_size_pub;

/// One table or view of a database, as the catalogue describes it.
///
/// The same query answers the terminal's overview sheet and the MCP server's container
/// listing, so the two cannot disagree about what a file holds.
#[derive(Clone, Debug)]
pub struct ContainerInfo {
    pub name: String,
    pub view: bool,
    /// `None` for a view: counting its rows means running it, and listing what a file
    /// holds should not execute somebody's ten-million-row query.
    pub rows: Option<i64>,
    pub columns: usize,
    /// The `CREATE TABLE` / `CREATE VIEW` the database keeps.
    pub sql: Option<String>,
}

/// Tables and views of a database, whichever engine it is.
///
/// The engine comes from the file's header rather than its name — see
/// [`db_write::kind_for_path`].
pub fn db_containers(path: &Path) -> Result<Vec<ContainerInfo>> {
    match db_write::kind_for_path(path) {
        db_write::DbKind::Sqlite => sqlite::sqlite_containers(path),
        db_write::DbKind::DuckDb => duckdb::duckdb_containers(path),
    }
}

pub fn load_file(path: &Path, delimiter: Option<u8>) -> Result<DataFrame> {
    load_file_with_doc(path, delimiter).map(|(df, _)| df)
}

/// Load `path`, additionally returning the document tree when the file is one of the
/// structured formats.  Sheets that keep the [`doc_io::DocState`] can edit and re-save
/// the real structure; callers that drop it get a read-only projection.
///
/// `forced` overrides the extension, which is how `--type` opens a file whose name says
/// nothing useful (`deploy.conf` as YAML).
pub fn load_file_with_doc(
    path: &Path,
    delimiter: Option<u8>,
) -> Result<(DataFrame, Option<doc_io::DocState>)> {
    load_file_as(path, delimiter, None)
}

pub fn load_file_as(
    path: &Path,
    delimiter: Option<u8>,
    forced: Option<crate::data::doc::Format>,
) -> Result<(DataFrame, Option<doc_io::DocState>)> {
    let ext = path
        .extension()
        .and_then(|e| e.to_str())
        .unwrap_or("csv")
        .to_lowercase();

    if let Some(fmt) = forced.or_else(|| crate::data::doc::Format::from_ext(&ext)) {
        let (df, state) = doc_io::DocState::open(path, fmt)?;
        return Ok((df, Some(state)));
    }

    // The extension said nothing useful.  Look at the contents before giving up — or,
    // for a file with no extension, before falling back to the CSV default.
    let has_ext = path.extension().is_some();
    let known_tabular = matches!(
        ext.as_str(),
        "csv"
            | "tsv"
            | "txt"
            | "parquet"
            | "arrow"
            | "feather"
            | "ipc"
            | "xlsx"
            | "xls"
            | "db"
            | "sqlite"
            | "sqlite3"
            | "duckdb"
            | "ddb"
    );
    if !has_ext || !known_tabular {
        if let Ok(text) = std::fs::read_to_string(path) {
            if let Some(fmt) = crate::data::doc::sniff(&text, !has_ext) {
                let mut doc = crate::data::doc::Doc::from_str(&text, fmt)?;
                doc.path = Some(path.to_path_buf());
                let (df, state) = doc_io::DocState::from_doc(doc)?;
                return Ok((df, Some(state)));
            }
        }
    }

    load_tabular(path, delimiter, &ext).map(|df| (df, None))
}

/// Read a file as `ext` says, whatever it happens to be called.
pub(crate) fn load_tabular(path: &Path, delimiter: Option<u8>, ext: &str) -> Result<DataFrame> {
    match ext {
        "csv" | "tsv" => crate::data::loader::load_csv(path, delimiter),
        "txt" => txt::load_txt(path),
        "parquet" => parquet::load_parquet(path),
        "arrow" | "feather" | "ipc" => arrow::load_arrow(path),
        "xlsx" | "xls" => excel::load_excel(path),
        // Not the extension: `.db` names no engine at all, and a name is only ever a
        // claim.  `kind_for_path` reads the file's own header and falls back to the
        // extension for one that does not exist yet — the same answer the writer uses.
        "db" | "sqlite" | "sqlite3" | "duckdb" | "ddb" => match db_write::kind_for_path(path) {
            db_write::DbKind::DuckDb => duckdb::load_duckdb_overview(path),
            db_write::DbKind::Sqlite => sqlite::load_sqlite_overview(path),
        },
        _ => Err(eyre!("Unsupported file format: .{}", ext)),
    }
}

/// Save to `path`, choosing the writer from its extension.
///
/// For the structured formats the rule is: a sheet that carries a document tree is
/// written by re-serialising that tree, so the original structure survives and
/// converting between formats is just picking a different extension.  A sheet with no
/// tree behind it (CSV, Parquet, SQL, pivot) is first turned into a tree using `shape`.
///
/// For a database target `sheet_name` is the name of the table to create.  The TUI asks
/// for it and never reaches here; the MCP server has nowhere to ask, so it passes its
/// own.
pub fn save_file_as(
    df: &DataFrame,
    doc: Option<&doc_io::DocState>,
    path: &Path,
    shape: doc_io::Shape,
    sheet_name: &str,
) -> Result<()> {
    let ext = path
        .extension()
        .and_then(|e| e.to_str())
        .unwrap_or("csv")
        .to_lowercase();

    if let Some(fmt) = crate::data::doc::Format::from_ext(&ext) {
        let opts = crate::data::doc::SaveOpts::default();
        return match doc {
            Some(state) => state.save_wrapped(path, fmt, &opts, sheet_name),
            None => doc_io::table_to_doc(df, shape, fmt, sheet_name)?.save_as(path, fmt, &opts),
        };
    }

    match ext.as_str() {
        "csv" => save_csv(df, path, b','),
        "tsv" => save_csv(df, path, b'\t'),
        "parquet" => parquet::save_parquet(df, path),
        "arrow" | "feather" | "ipc" => arrow::save_arrow(df, path),
        // The engine comes from the file when there is one — adding a table to an
        // existing database has to speak that database's dialect, whatever it is called.
        "db" | "sqlite" | "sqlite3" | "duckdb" | "ddb" => {
            db_write::create_table(db_write::kind_for_path(path), path, sheet_name, df)
        }
        "xlsx" | "xls" => excel::save_xlsx(df, path),
        _ => Err(eyre!("Unsupported save format: .{}", ext)),
    }
}

/// Whether [`save_file_as`] has a writer for this extension.
///
/// The same list, so a caller can ask *before* doing the work that produces the rows —
/// finding out the target is unwritable after a join and a group-by has run is a waste
/// the answer was always going to be able to prevent.
pub fn writable_ext(ext: &str) -> bool {
    let ext = ext.to_lowercase();
    crate::data::doc::Format::from_ext(&ext).is_some()
        || matches!(
            ext.as_str(),
            "csv"
                | "tsv"
                | "parquet"
                | "arrow"
                | "feather"
                | "ipc"
                | "db"
                | "sqlite"
                | "sqlite3"
                | "duckdb"
                | "ddb"
                | "xlsx"
                | "xls"
        )
}

pub fn load_from_stdin_typed(data_type: &str, delimiter: Option<u8>) -> Result<DataFrame> {
    load_from_stdin_with_doc(data_type, delimiter).map(|(df, _)| df)
}

pub fn load_from_stdin_with_doc(
    data_type: &str,
    delimiter: Option<u8>,
) -> Result<(DataFrame, Option<doc_io::DocState>)> {
    use std::io::{Read, Write};
    use tempfile::NamedTempFile;

    let mut buf = Vec::new();
    std::io::stdin().read_to_end(&mut buf)?;

    if let Some(fmt) = crate::data::doc::Format::from_name(data_type) {
        let text = String::from_utf8(buf)?;
        let (df, state) = doc_io::DocState::from_doc(crate::data::doc::Doc::from_str(&text, fmt)?)?;
        return Ok((df, Some(state)));
    }

    let mut temp_file = NamedTempFile::new()?;
    temp_file.write_all(&buf)?;
    let temp_path = temp_file.path().to_path_buf();

    let pdf = match data_type.to_lowercase().as_str() {
        "csv" | "txt" => {
            let sep = delimiter.unwrap_or(b',');
            polars::prelude::CsvReadOptions::default()
                .with_has_header(true)
                .map_parse_options(|o| o.with_separator(sep))
                .try_into_reader_with_file_path(Some(temp_path))?
                .finish()?
        }
        "tsv" => polars::prelude::CsvReadOptions::default()
            .with_has_header(true)
            .map_parse_options(|o| o.with_separator(b'\t'))
            .try_into_reader_with_file_path(Some(temp_path))?
            .finish()?,
        _ => return Err(eyre!("Unsupported stdin data type: {}", data_type)),
    };

    drop(temp_file);
    wrap_polars_df(pdf).map(|df| (df, None))
}

pub(crate) fn wrap_polars_df(pdf: polars::prelude::DataFrame) -> Result<DataFrame> {
    let col_count = pdf.width();
    let mut columns = Vec::with_capacity(col_count);

    for series in pdf.columns() {
        let name = series.name().to_string();
        let mut col_meta = ColumnMeta::new(name);

        col_meta.col_type = match series.dtype() {
            DataType::Int8
            | DataType::Int16
            | DataType::Int32
            | DataType::Int64
            | DataType::UInt8
            | DataType::UInt16
            | DataType::UInt32
            | DataType::UInt64 => ColumnType::Integer,
            DataType::Float32 | DataType::Float64 => ColumnType::Float,
            DataType::Date => ColumnType::Date,
            DataType::Datetime(_, _) => ColumnType::Datetime,
            _ => ColumnType::String,
        };

        columns.push(col_meta);
    }

    let mut df = DataFrame::from_parts(pdf, columns);
    df.calc_widths(40, 1000);
    Ok(df)
}

fn save_csv(df: &DataFrame, path: &Path, delimiter: u8) -> Result<()> {
    let mut out_df = df.to_display_polars_df();
    let mut file = File::create(path)?;
    CsvWriter::new(&mut file)
        .include_header(true)
        .with_separator(delimiter)
        .finish(&mut out_df)?;
    Ok(())
}