tomq 0.1.2

use crate::dociter::TomqIter;
use crate::{output_to_stdout, DocumentIterator, TomlOpts};
use serde_json::{Map, Value};
use std::io::{Cursor, ErrorKind, Read};
use toml_edit::visit_mut::VisitMut;
use toml_edit::{ArrayOfTables, Document, Item};

use crate::error::Result;
use crate::error::TomqError;
use crate::opts::RootKey;
use crate::output::{output_to_bat, CopyFromReader};
use crate::toml_pp::Pretty;

/// Parses TOML document lazily from [`R` reader][Read] as it is consumed from
/// the returned iterator.
pub(crate) fn parse_toml_document<R: Read>(
    reader: R,
    opts: &TomlOpts,
) -> impl Iterator<Item = Result<Document>> {
    DocumentIterator::new(
        reader,
        format!("\n{}", opts.multi_doc_separator),
        opts.skip_invalid,
    )
}

/// Converts Json Document to Toml Document lazily from [`R` reader][Read]
/// as it is consumed from the returned iterator.
pub(crate) fn create_toml_doc_from_json<R: Read>(
    reader: R,
    ignore_error: bool,
    null_to_empty: bool,
    fallback_key: Option<String>,
) -> impl TomqIter<Document> {
    let des = serde_json::Deserializer::from_reader(reader);
    create_toml_doc_from_json_values(
        des.into_iter::<Value>()
            .map(|r| r.map_err(|e| TomqError::from(e))),
        ignore_error,
        null_to_empty,
        fallback_key,
    )
}

/// Converts Json Document to Toml Document lazily from [`TomqIter<serde_json::Value>` iterator][TomqIter]
/// as it is consumed from the returned iterator.
pub(crate) fn create_toml_doc_from_json_values(
    values: impl TomqIter<Value>,
    ignore_error: bool,
    null_to_empty: bool,
    fallback_key: Option<String>,
) -> impl TomqIter<Document> {
    values
        .map(move |v| {
            v.and_then(|value| match value {
                Value::Null if null_to_empty => Ok(Document::new()),
                value @ _ if !value.is_object() && fallback_key.is_some() => {
                    let value = Value::Object(Map::from_iter(
                        Some((fallback_key.clone().unwrap(), value)).into_iter(),
                    ));
                    toml_edit::ser::to_document(&value)
                        .map_err(|e| TomqError::TomlFromJson(e, value))
                }
                _ => toml_edit::ser::to_document(&value)
                    .map_err(|e| TomqError::TomlFromJson(e, value)),
            })
        })
        .filter(move |f| !ignore_error || f.is_ok())
}

/// Converts Json Document (or documents) to Toml Documents from [`R` reader][Read] and
/// outputs to stdout.
///
/// This is the equivalent to calling [`create_toml_doc_from_json`] and passing it to [`copy_toml`]
/// as the [`TomqIter<Document>`][`TomqIter`] with a selected [writer][`CopyFromReader`] based on whether it should
/// pipe the output to `bat` or stdout.
///
/// If you want to copy the toml document from a [`TomqIter<Document>`][`TomqIter`] to a [writer][`CopyFromReader`],
/// use [`copy_toml`].
pub(crate) fn print_toml_from_json_to_stdout<R: Read>(
    read: R,
    pretty: bool,
    root_key: &RootKey,
    fallback_key: Option<String>,
    multi_doc_separator: &str,
    skip_invalid: bool,
    null_to_empty_doc: bool,
    bat: bool,
) -> Result<()> {
    let toml = create_toml_doc_from_json(read, skip_invalid, null_to_empty_doc, fallback_key);
    print_toml(toml, pretty, root_key, &multi_doc_separator, bat)?;
    Ok(())
}

/// Converts Json Document (or documents) to Toml Documents from [`TomqIter<serde_json::Value>` iterator][TomqIter] and
/// outputs to stdout.
///
/// This is the equivalent to calling [`create_toml_doc_from_json_values`] and passing it to [`copy_toml`]
/// as the [`TomqIter<Document>`][`TomqIter`] with a selected [writer][`CopyFromReader`] based on whether it should
/// pipe the output to `bat` or stdout.
///
/// If you want to copy the toml document from a [`TomqIter<Document>`][`TomqIter`] to a [writer][`CopyFromReader`],
/// use [`copy_toml`].
pub(crate) fn print_toml_from_json_values(
    values: impl TomqIter<Value>,
    pretty: bool,
    root_key: &RootKey,
    fallback_key: Option<String>,
    multi_doc_separator: &str,
    skip_invalid: bool,
    null_to_empty_doc: bool,
    bat: bool,
) -> Result<()> {
    let toml =
        create_toml_doc_from_json_values(values, skip_invalid, null_to_empty_doc, fallback_key);
    print_toml(toml, pretty, root_key, &multi_doc_separator, bat)?;
    Ok(())
}

/// Converts Json Document (or documents) to Toml Documents from [`R` reader][Read] and
/// outputs to stdout.
///
/// This is the equivalent to calling [`create_toml_doc_from_json`] and passing it to [`copy_toml`]
/// as the [`TomqIter<Document>`][`TomqIter`] with a selected [writer][`CopyFromReader`] based on whether it should
/// pipe the output to `bat` or stdout.
///
/// If you want to copy the toml document from a [`TomqIter<Document>`][`TomqIter`] to a [writer][`CopyFromReader`],
/// use [`copy_toml`].
pub(crate) fn print_toml(
    toml: impl TomqIter<Document>,
    pretty: bool,
    root_key: &RootKey,
    multi_doc_separator: &str,
    bat: bool,
) -> Result<()> {
    if bat {
        copy_toml(
            toml,
            output_to_bat("toml"),
            pretty,
            root_key,
            &multi_doc_separator,
        )?;
    } else {
        copy_toml(
            toml,
            output_to_stdout(),
            pretty,
            root_key,
            &multi_doc_separator,
        )?;
    }
    Ok(())
}

/// Copies the toml from the [reader][DocIter] to the [writer][CopyFromReader].
pub(crate) fn copy_toml(
    reader: impl TomqIter<Document>,
    mut writer: impl CopyFromReader,
    pretty: bool,
    root_key: &RootKey,
    multi_doc_separator: &str,
) -> Result<()> {
    match root_key {
        RootKey::Flatten => {
            let txt = stream_toml_with_multi_doc(reader, pretty, multi_doc_separator.to_string())
                .map(|s| s.map(|s| Cursor::new(s.into_bytes())));

            writer.copy_from_reader(txt.iter_as_read())
        }
        RootKey::Key(key) => {
            let txt = copy_toml_with_root_key(reader, pretty, key)
                .map(|s| s.map(|s| Cursor::new(s.into_bytes())));
            writer.copy_from_reader(txt.iter_as_read())
        }
    }
}

/// Takes Toml documents from the [reader][DocIter] and lazily transforms those documents into
/// multi-documents using the [Tomq Multi Document Extension]() (non-Toml spec compliant)
/// as the returned iterator is consumed.
///
/// This creates a **lazy** iterator that only walks the input [`iter`][`TomqIter`] when
/// the returned iterator is advanced, this way it avoids pre-allocating memory until it is
/// really needed.
fn stream_toml_with_multi_doc(
    iter: impl TomqIter<Document>,
    pretty: bool,
    multi_doc_separator: String,
) -> impl TomqIter<String> {
    iter.enumerate().map(move |(idx, v)| {
        v.and_then(|mut doc| toml_doc_to_string(&mut doc, pretty))
            .map(|s| {
                if idx > 0 {
                    format!("\n{}\n{}", multi_doc_separator.clone(), s)
                } else {
                    s
                }
            })
    })
}

/// Takes Toml documents from the [reader][DocIter] and adds those documents into an
/// [ArrayOfTables](https://toml.io/en/v1.0.0#array-of-tables)
/// with a common key (Toml spec compliant) as the returned iterator is consumed.
///
/// This operation is not **lazy** like the other ones, at the moment it computes all the values
/// and add them to a single [ArrayOfTables](https://toml.io/en/v1.0.0#array-of-tables).
///
/// TODO: Make it lazy. The fix is very easy to say, since the way Toml is designed
/// TODO: we don't need to know about the others tables that belongs to the same array,
/// TODO: we just need to add a `[[{KEY}]]` before the table and prepend the key in all the others.
fn copy_toml_with_root_key(
    iter: impl TomqIter<Document>,
    pretty: bool,
    key: &str,
) -> impl TomqIter<String> {
    let key = key.to_string();
    Some(iter.fold(Ok(Vec::<Document>::new()), |acc, v| match acc {
        Ok(mut tab) => match v {
            Ok(v) => {
                tab.push(v);
                Ok(tab)
            }
            Err(e) => Err(e),
        },
        Err(e) => Err(e),
    }))
    .into_iter()
    .map(move |v| {
        v.map(|v| match v.len() {
            0 => Document::new(),
            // NOTE: Always use common key when asked for, to avoid unexpected behavior, even if
            // we correctly document this, people still might get confused when it happens.
            //1 => v.into_iter().next().unwrap(),
            _ => {
                let mut doc = Document::new();
                let table = doc.as_table_mut();
                let mut arr = ArrayOfTables::new();
                for x in v.into_iter() {
                    arr.push(x.as_table().clone());
                }
                table.insert(key.as_str(), Item::ArrayOfTables(arr));
                doc
            }
        })
    })
    .map(move |v| v.and_then(|mut doc| toml_doc_to_string(&mut doc, pretty)))
}

/// Produces toml output from the document, using a forked version of `toml_edit` pretty printing
/// code (since it is not exposed to other crates).
fn toml_doc_to_string(value: &mut Document, pretty_print: bool) -> Result<String> {
    if pretty_print {
        Pretty.visit_document_mut(value);
        Ok(value.to_string())
    } else {
        Ok(value.to_string())
    }
}

/// Simple helper that interfaces [`Read::read`] calls
/// with [`Iterator::<Item = Result<V>>::next`][`Iterator::next`].
struct IterToRead<I: Iterator, V> {
    iter: I,
    current: Option<V>,
}

impl<I: Iterator, V> IterToRead<I, V> {
    fn new(iter: I) -> Self {
        Self {
            iter,
            current: None,
        }
    }
}

trait IterAsRead: Iterator {
    fn iter_as_read<V>(self) -> IterToRead<Self, V>
    where
        Self: Sized,
    {
        IterToRead::new(self)
    }
}

impl<I> IterAsRead for I where I: Iterator {}

impl<I, V> Read for IterToRead<I, V>
where
    I: Iterator<Item = Result<V>>,
    V: Read,
{
    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
        loop {
            match self.current.as_mut() {
                None => match self.iter.next() {
                    None => return Ok(0),
                    Some(n) => {
                        self.current =
                            Some(n.map_err(|e| std::io::Error::new(ErrorKind::InvalidData, e))?);
                    }
                },
                Some(r) => {
                    let n = r.read(buf)?;
                    if n == 0 {
                        self.current = None;
                    } else {
                        return Ok(n);
                    }
                }
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use crate::opts::RootKey;
    use crate::{parse_toml_document, TomlOpts};
    use std::io::{BufReader, Cursor};

    macro_rules! assert_ok {
        ($cond:expr $(,)?) => {
            match $cond {
                Some(Ok(d)) => {
                    d
                },
                Some(Err(e)) => {
                    assert!(false, "Unexpected error: {}", e)
                },
                None => panic!(),
            }
        };
        ($cond:expr, $($arg:tt)+) => {
            match $cond {
                Some(Ok(d)) => {
                    d
                },
                Some(Err(e)) => {
                    panic!("Expected successful result, got error: {}. Message: {}", e, format!($($arg)+))
                },
                None => panic!($($arg)+),
            }
        };
    }

    macro_rules! assert_not_empty {
        ($cond:expr $(,)?) => {
            assert!(assert_ok!($cond).iter().next().is_some());
        };
        ($cond:expr, $($arg:tt)+) => {
            assert!(assert_ok!($cond, $($arg)+).iter().next().is_some(), $($arg)+);
        };
    }

    macro_rules! assert_empty {
        ($cond:expr $(,)?) => {
            assert!(assert_ok!($cond).iter().next().is_none());
        };
        ($cond:expr, $($arg:tt)+) => {
            assert!(assert_ok!($cond, $($arg)+).iter().next().is_none(), $($arg)+);
        };
    }

    /// Test against the strange corner case which the last empty document is treated
    /// as an entire toml document instead of multi-document separator.
    #[test]
    fn multi_document_last_empty() {
        let doc = r#"
name = "tomq"

---
name = "tomq2"

---
"#;
        let reader = BufReader::with_capacity(1, Cursor::new(doc.as_bytes()));
        let mut doc = parse_toml_document(
            reader,
            &TomlOpts {
                jq_filter: None,
                toml_it: false,
                bat: false,
                retranscode: false,
                skip_invalid: false,
                single_doc: false,
                null_to_empty_doc: false,
                multi_doc_separator: "---".to_string(),
                root_key: RootKey::Flatten,
                fallback_key: None,
                pretty_print: false,
                jq_args: vec![],
            },
        );

        assert_not_empty!(doc.next(), "First document parse");
        assert_not_empty!(doc.next(), "Second document parse");
        assert_empty!(doc.next(), "Third empty document parse");
        assert!(doc.next().is_none(), "Should stop parsing");
        assert!(doc.next().is_none(), "Should stop parsing");
    }
}