query-forge 0.5.0

use std::{fs, path::Path};

use anyhow::{Result, anyhow, bail};

use crate::cli::OutputFormat;

pub(crate) fn write_query_result(
    result: &query_forge::QueryResult,
    output_path: Option<&Path>,
    explicit_format: Option<OutputFormat>,
) -> Result<()> {
    match resolve_output_format(output_path, explicit_format)? {
        OutputFormat::Text => {
            let rendered = query_forge::render_text(result);
            write_or_print(&rendered, output_path)?;
        }
        OutputFormat::Csv => {
            let rendered = query_forge::render_csv(result);
            write_or_print(&rendered, output_path)?;
        }
        OutputFormat::Json => {
            let rendered = query_forge::render_json(result);
            write_or_print(&rendered, output_path)?;
        }
        OutputFormat::Jsonl => {
            let rendered = query_forge::render_jsonl(result);
            write_or_print(&rendered, output_path)?;
        }
        OutputFormat::Markdown => {
            let rendered = query_forge::render_markdown(result);
            write_or_print(&rendered, output_path)?;
        }
        OutputFormat::Html => {
            let rendered = query_forge::render_html(result);
            write_or_print(&rendered, output_path)?;
        }
        OutputFormat::Xml => {
            let rendered = query_forge::render_xml(result);
            write_or_print(&rendered, output_path)?;
        }
        OutputFormat::Xlsx => {
            let output_path = output_path
                .ok_or_else(|| anyhow!("--output is required when --format xlsx is selected"))?;
            query_forge::write_xlsx(result, output_path)?;
        }
        OutputFormat::Parquet => {
            let output_path = output_path.ok_or_else(|| {
                anyhow!("--output is required when --format parquet is selected")
            })?;
            query_forge::write_parquet(result, output_path)?;
        }
    }

    Ok(())
}

fn write_or_print(rendered: &str, output_path: Option<&Path>) -> Result<()> {
    if let Some(output_path) = output_path {
        fs::write(output_path, rendered)?;
    } else {
        println!("{rendered}");
    }

    Ok(())
}

pub(crate) fn resolve_output_format(
    output: Option<&Path>,
    format: Option<OutputFormat>,
) -> Result<OutputFormat> {
    if let Some(format) = format {
        if matches!(format, OutputFormat::Xlsx | OutputFormat::Parquet) && output.is_none() {
            let name = if matches!(format, OutputFormat::Xlsx) {
                "xlsx"
            } else {
                "parquet"
            };
            bail!("--output is required when --format {name} is selected");
        }
        return Ok(format);
    }

    if output
        .and_then(|path| path.extension())
        .map(|extension| extension.to_string_lossy())
        .is_some_and(|extension| extension.eq_ignore_ascii_case("xlsx"))
    {
        return Ok(OutputFormat::Xlsx);
    }

    if output
        .and_then(|path| path.extension())
        .map(|extension| extension.to_string_lossy())
        .is_some_and(|extension| extension.eq_ignore_ascii_case("parquet"))
    {
        return Ok(OutputFormat::Parquet);
    }

    if output
        .and_then(|path| path.extension())
        .map(|extension| extension.to_string_lossy())
        .is_some_and(|extension| extension.eq_ignore_ascii_case("csv"))
    {
        return Ok(OutputFormat::Csv);
    }

    if output
        .and_then(|path| path.extension())
        .map(|extension| extension.to_string_lossy())
        .is_some_and(|extension| extension.eq_ignore_ascii_case("json"))
    {
        return Ok(OutputFormat::Json);
    }

    if output
        .and_then(|path| path.extension())
        .map(|extension| extension.to_string_lossy())
        .is_some_and(|extension| extension.eq_ignore_ascii_case("jsonl"))
    {
        return Ok(OutputFormat::Jsonl);
    }

    if output
        .and_then(|path| path.extension())
        .map(|extension| extension.to_string_lossy())
        .is_some_and(|extension| {
            extension.eq_ignore_ascii_case("md") || extension.eq_ignore_ascii_case("markdown")
        })
    {
        return Ok(OutputFormat::Markdown);
    }

    if output
        .and_then(|path| path.extension())
        .map(|extension| extension.to_string_lossy())
        .is_some_and(|extension| extension.eq_ignore_ascii_case("xml"))
    {
        return Ok(OutputFormat::Xml);
    }

    if output
        .and_then(|path| path.extension())
        .map(|extension| extension.to_string_lossy())
        .is_some_and(|extension| {
            extension.eq_ignore_ascii_case("html") || extension.eq_ignore_ascii_case("htm")
        })
    {
        return Ok(OutputFormat::Html);
    }

    Ok(OutputFormat::Text)
}

#[cfg(test)]
mod tests {
    use std::path::Path;

    use crate::cli::OutputFormat;

    use super::resolve_output_format;

    #[test]
    fn infers_output_format_from_extension() {
        assert!(matches!(
            resolve_output_format(Some(Path::new("result.xlsx")), None).expect("xlsx format"),
            OutputFormat::Xlsx
        ));
        assert!(matches!(
            resolve_output_format(Some(Path::new("result.parquet")), None).expect("parquet format"),
            OutputFormat::Parquet
        ));
        assert!(matches!(
            resolve_output_format(Some(Path::new("result.csv")), None).expect("csv format"),
            OutputFormat::Csv
        ));
        assert!(matches!(
            resolve_output_format(Some(Path::new("result.jsonl")), None).expect("jsonl format"),
            OutputFormat::Jsonl
        ));
        assert!(matches!(
            resolve_output_format(Some(Path::new("result.json")), None).expect("json format"),
            OutputFormat::Json
        ));
        assert!(matches!(
            resolve_output_format(Some(Path::new("result.md")), None).expect("md format"),
            OutputFormat::Markdown
        ));
        assert!(matches!(
            resolve_output_format(Some(Path::new("result.html")), None).expect("html format"),
            OutputFormat::Html
        ));
        assert!(matches!(
            resolve_output_format(Some(Path::new("result.htm")), None).expect("htm format"),
            OutputFormat::Html
        ));
        assert!(matches!(
            resolve_output_format(Some(Path::new("result.txt")), None).expect("text format"),
            OutputFormat::Text
        ));
    }
}