jql 8.2.0

#![doc = include_str!("../README.md")]

mod args;
mod panic;

#[cfg(not(target_arch = "loongarch64"))]
use mimalloc::MiMalloc;

#[cfg(not(target_arch = "loongarch64"))]
#[global_allocator]
static GLOBAL: MiMalloc = MiMalloc;

use std::{
    path::Path,
    process::exit,
};

use anyhow::{
    Context,
    Result,
    anyhow,
};
use args::Args;
use clap::Parser;
use colored_json::{
    ColoredFormatter,
    CompactFormatter,
    PrettyFormatter,
};
use jql_runner::runner;
use panic::use_custom_panic_hook;
use serde::Deserialize;
use serde_json::Value;
use serde_stacker::Deserializer;
use tokio::{
    fs::File,
    io::{
        AsyncBufReadExt,
        AsyncReadExt,
        AsyncWriteExt,
        BufReader,
        stdin,
        stdout,
    },
};

/// Reads a file from `path`.
async fn read_file(path: impl AsRef<Path>) -> Result<String> {
    let display_path = path.as_ref().display();
    let mut file = File::open(&path)
        .await
        .with_context(|| format!("Failed to open file {display_path}"))?;
    let mut contents = vec![];

    file.read_to_end(&mut contents)
        .await
        .with_context(|| format!("Failed to read from file {display_path}"))?;

    Ok(String::from_utf8_lossy(&contents).into_owned())
}

/// Renders the output or the error and exits.
fn render(result: Result<String>) {
    match result {
        Ok(output) => println!("{output}"),
        Err(error) => {
            eprintln!("{error}");
            exit(1);
        }
    }
}

/// Processes the JSON content based on the arguments.
async fn process_json(json: &str, args: &Args) -> Result<String> {
    if args.validate {
        return serde_json::from_str::<Value>(json).map_or_else(
            |_| Err(anyhow!("Invalid JSON file or content")),
            |_| Ok("Valid JSON file or content".to_string()),
        );
    }

    let query = match args.query_from_file.as_deref() {
        Some(path) => read_file(path).await?,
        // We can safely unwrap since clap is taking care of the validation.
        None => args.query.as_deref().unwrap().to_string(),
    };

    let mut deserializer = serde_json::Deserializer::from_str(json);

    deserializer.disable_recursion_limit();

    let deserializer = Deserializer::new(&mut deserializer);
    let value: Value = Value::deserialize(deserializer)
        .with_context(|| "Failed to deserialize the JSON data".to_string())?;
    let mut result: Value = runner::raw(&query, &value)?;

    if args.sort_keys {
        result.sort_all_objects();
    }

    if args.inline {
        return ColoredFormatter::new(CompactFormatter {})
            .to_colored_json_auto(&result)
            .with_context(|| "Failed to inline the JSON data".to_string());
    }

    if args.raw_string && result.is_string() {
        // We can safely unwrap since the result is a string.
        return Ok(String::from(result.as_str().unwrap()));
    }

    ColoredFormatter::new(PrettyFormatter::new())
        .to_colored_json_auto(&result)
        .with_context(|| "Failed to format the JSON data".to_string())
}

#[tokio::main]
async fn main() -> Result<()> {
    // Use a custom panic hook.
    use_custom_panic_hook();

    let args = Args::parse();

    if let Some(path) = args.json_file.as_deref() {
        let contents = read_file(path).await?;

        render(process_json(&contents, &args).await);

        return Ok(());
    }

    let mut stdout = stdout();

    if args.stream {
        let mut reader = BufReader::new(stdin()).lines();

        while let Some(mut line) = reader
            .next_line()
            .await
            .with_context(|| "Failed to read stream".to_string())?
        {
            render(process_json(&line, &args).await);

            stdout
                .flush()
                .await
                .with_context(|| "Failed to flush stdout".to_string())?;

            line.clear();
        }

        return Ok(());
    }

    let mut buffer = Vec::new();
    let mut stdin = stdin();

    // By default, read the whole piped content from stdin.
    stdin
        .read_to_end(&mut buffer)
        .await
        .with_context(|| "Failed to read piped content from stdin".to_string())?;

    let lines = String::from_utf8(buffer)
        .with_context(|| "Failed to convert piped content from stdin".to_string())?;

    render(process_json(&lines, &args).await);

    Ok(())
}

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

    #[tokio::test]
    async fn sort_keys_flag_sorts_objects_recursively() {
        let json = r#"{ "root": { "d": 1, "b": { "y": 1, "x": 2 }, "a": 3 } }"#;
        let args = Args::parse_from(["jql", "--sort-keys", "--inline", r#""root""#]);

        assert_eq!(
            process_json(json, &args).await.unwrap(),
            r#"{"a":3,"b":{"x":2,"y":1},"d":1}"#
        );
    }

    #[tokio::test]
    async fn output_keeps_source_order_without_the_flag() {
        let json = r#"{ "root": { "d": 1, "a": 3 } }"#;
        let args = Args::parse_from(["jql", "--inline", r#""root""#]);

        assert_eq!(process_json(json, &args).await.unwrap(), r#"{"d":1,"a":3}"#);
    }
}