peanuts/src/writer.rs

use std::{
    collections::{HashSet, VecDeque},
    pin::pin,
    str::FromStr,
    task::Poll,
};

use async_recursion::async_recursion;
use circular::Buffer;
use futures::{Future, FutureExt, Sink, SinkExt};
use tokio::io::{AsyncWrite, AsyncWriteExt};

use crate::{
    declaration::{Declaration, VersionInfo},
    element::{escape_str, Content, Element, IntoContent, IntoElement, Name, NamespaceDeclaration},
    endable::Endable,
    error::Error,
    xml::{self, composers::Composer, parsers_complete::Parser, ETag, XMLDecl},
    Result, XMLNS_NS, XML_NS,
};

// pub struct Writer<W, C = Composer> {
#[derive(Debug)]
pub struct Writer<W> {
    inner: Endable<W>,
    depth: Vec<Name>,
    namespace_declarations: Vec<HashSet<NamespaceDeclaration>>,
}

impl<W> Writer<W> {
    pub fn new(writer: W) -> Self {
        let mut default_declarations = HashSet::new();
        default_declarations.insert(NamespaceDeclaration {
            prefix: Some("xml".to_string()),
            namespace: XML_NS.to_string(),
        });
        default_declarations.insert(NamespaceDeclaration {
            prefix: Some("xmlns".to_string()),
            namespace: XMLNS_NS.to_string(),
        });
        Self {
            inner: Endable::new(writer),
            depth: Vec::new(),
            namespace_declarations: vec![default_declarations],
        }
    }

    pub fn into_inner(self) -> W {
        self.inner.into_inner()
    }
}

impl<W: AsyncWrite + Unpin + Send> Writer<W> {
    pub async fn write_declaration(&mut self, version: VersionInfo) -> Result<()> {
        let writer = self.inner.try_as_mut()?;
        let declaration = Declaration::version(version);
        let version_info;
        match declaration.version_info {
            VersionInfo::One => version_info = xml::VersionInfo::SingleQuoted(xml::VersionNum::One),
            VersionInfo::OneDotOne => {
                version_info = xml::VersionInfo::SingleQuoted(xml::VersionNum::OneDotOne)
            }
        }
        let declaration = xml::XMLDecl {
            version_info,
            encoding_decl: None,
            sd_decl: None,
        };
        declaration.write(writer).await?;
        Ok(())
    }

    pub async fn write_full(&mut self, into_element: &impl IntoElement) -> Result<()> {
        let element = into_element.into_element();
        Ok(self.write_element(&element).await?)
    }

    pub async fn write_start(&mut self, into_element: &impl IntoElement) -> Result<()> {
        let element = into_element.into_element();
        Ok(self.write_element_start(&element).await?)
    }

    pub async fn write_all_content(&mut self, into_element: &impl IntoElement) -> Result<()> {
        for content in &into_element.get_content() {
            self.write_content(content).await?;
        }
        Ok(())
    }

    pub async fn write(&mut self, into_content: &impl IntoContent) -> Result<()> {
        let content = into_content.into_content();
        Ok(self.write_content(&content).await?)
    }

    #[async_recursion]
    pub async fn write_element(&mut self, element: &Element) -> Result<()> {
        if element.content.is_empty() {
            self.write_empty(element).await?;
        } else {
            self.write_element_start(element).await?;
            for content in &element.content {
                self.write_content(content).await?;
            }
            self.write_end().await?;
        }
        Ok(())
    }

    pub async fn write_empty(&mut self, element: &Element) -> Result<()> {
        let writer = self.inner.try_as_mut()?;
        let mut namespace_declarations_stack: Vec<_> = self
            .namespace_declarations
            .iter()
            .flatten()
            .chain(&element.namespace_declaration_overrides)
            .collect();

        let mut namespace_declarations = element.namespace_declaration_overrides.clone();

        let default_namespace_declaration;
        let prefix;
        if let Some(namespace) = &element.name.namespace {
            if let Some(name_namespace_declaration) = namespace_declarations_stack
                .iter()
                .rfind(|namespace_declaration| namespace_declaration.namespace == *namespace)
            {
                prefix = name_namespace_declaration.prefix.as_ref();
            } else {
                default_namespace_declaration = NamespaceDeclaration {
                    prefix: None,
                    namespace: namespace.clone(),
                };
                if namespace_declarations.insert(default_namespace_declaration.clone()) {
                    namespace_declarations_stack.push(&default_namespace_declaration);
                    prefix = None
                } else {
                    return Err(Error::DuplicateNameSpaceDeclaration(NamespaceDeclaration {
                        prefix: None,
                        namespace: namespace.clone(),
                    }));
                }
            }
        } else {
            prefix = None
        }

        let name;
        if let Some(prefix) = &prefix {
            name = xml::QName::PrefixedName(xml::PrefixedName {
                prefix: xml::Prefix::parse_full(prefix)?,
                local_part: xml::LocalPart::parse_full(&element.name.local_name)?,
            })
        } else {
            name = xml::QName::UnprefixedName(xml::UnprefixedName::parse_full(
                &element.name.local_name,
            )?)
        }

        let mut attributes = Vec::new();

        for namespace_declaration in namespace_declarations.iter() {
            let ns_name = namespace_declaration
                .prefix
                .as_ref()
                .map(|prefix| -> Result<_> {
                    Ok(xml::NSAttName::PrefixedAttName(xml::PrefixedAttName(
                        xml::NCName::parse_full(&prefix)?,
                    )))
                })
                .unwrap_or(Ok(xml::NSAttName::DefaultAttName))?;
            let value = xml::AttValue::from(namespace_declaration.namespace.as_str());
            let xml_attribute = xml::Attribute::NamespaceDeclaration { ns_name, value };
            attributes.push(xml_attribute);
        }

        for (name, value) in &element.attributes {
            let prefix;
            if let Some(namespace) = &name.namespace {
                let name_namespace_declaration = namespace_declarations_stack
                    .iter()
                    .rfind(|namespace_declaration| namespace_declaration.namespace == *namespace)
                    .ok_or(Error::UndeclaredNamespace(namespace.clone()))?;
                prefix = name_namespace_declaration.prefix.as_ref();
            } else {
                prefix = None
            }

            let att_name;
            if let Some(prefix) = &prefix {
                att_name = xml::QName::PrefixedName(xml::PrefixedName {
                    prefix: xml::Prefix::parse_full(prefix)?,
                    local_part: xml::LocalPart::parse_full(&name.local_name)?,
                })
            } else {
                att_name =
                    xml::QName::UnprefixedName(xml::UnprefixedName::parse_full(&name.local_name)?)
            }

            let value = xml::AttValue::from(value.as_str());

            let xml_attribute = xml::Attribute::Attribute {
                name: att_name,
                value,
            };
            attributes.push(xml_attribute);
        }

        let tag = xml::EmptyElemTag { name, attributes };

        tag.write(writer).await?;

        if self.depth.is_empty() {
            self.inner.end();
        }

        Ok(())
    }

    pub async fn write_element_start(&mut self, element: &Element) -> Result<()> {
        let writer = self.inner.try_as_mut()?;
        let mut namespace_declarations_stack: Vec<_> = self
            .namespace_declarations
            .iter()
            .flatten()
            .chain(&element.namespace_declaration_overrides)
            .collect();

        let mut namespace_declarations = element.namespace_declaration_overrides.clone();

        let default_namespace_declaration;
        let prefix;
        if let Some(namespace) = &element.name.namespace {
            if let Some(name_namespace_declaration) = namespace_declarations_stack
                .iter()
                .rfind(|namespace_declaration| namespace_declaration.namespace == *namespace)
            {
                prefix = name_namespace_declaration.prefix.as_ref();
            } else {
                default_namespace_declaration = NamespaceDeclaration {
                    prefix: None,
                    namespace: namespace.clone(),
                };
                if namespace_declarations.insert(default_namespace_declaration.clone()) {
                    namespace_declarations_stack.push(&default_namespace_declaration);
                    prefix = None
                } else {
                    return Err(Error::DuplicateNameSpaceDeclaration(NamespaceDeclaration {
                        prefix: None,
                        namespace: namespace.clone(),
                    }));
                }
            }
        } else {
            prefix = None
        }

        let name;
        if let Some(prefix) = &prefix {
            name = xml::QName::PrefixedName(xml::PrefixedName {
                prefix: xml::Prefix::parse_full(prefix)?,
                local_part: xml::LocalPart::parse_full(&element.name.local_name)?,
            })
        } else {
            name = xml::QName::UnprefixedName(xml::UnprefixedName::parse_full(
                &element.name.local_name,
            )?)
        }

        let mut attributes = Vec::new();

        for namespace_declaration in namespace_declarations.iter() {
            let ns_name = namespace_declaration
                .prefix
                .as_ref()
                .map(|prefix| -> Result<_> {
                    Ok(xml::NSAttName::PrefixedAttName(xml::PrefixedAttName(
                        xml::NCName::parse_full(&prefix)?,
                    )))
                })
                .unwrap_or(Ok(xml::NSAttName::DefaultAttName))?;
            let value = xml::AttValue::from(namespace_declaration.namespace.as_str());
            let xml_attribute = xml::Attribute::NamespaceDeclaration { ns_name, value };
            attributes.push(xml_attribute);
        }

        for (name, value) in &element.attributes {
            let prefix;
            if let Some(namespace) = &name.namespace {
                let name_namespace_declaration = namespace_declarations_stack
                    .iter()
                    .rfind(|namespace_declaration| namespace_declaration.namespace == *namespace)
                    .ok_or(Error::UndeclaredNamespace(namespace.clone()))?;
                prefix = name_namespace_declaration.prefix.as_ref();
            } else {
                prefix = None
            }

            let att_name;
            if let Some(prefix) = &prefix {
                att_name = xml::QName::PrefixedName(xml::PrefixedName {
                    prefix: xml::Prefix::parse_full(prefix)?,
                    local_part: xml::LocalPart::parse_full(&name.local_name)?,
                })
            } else {
                att_name =
                    xml::QName::UnprefixedName(xml::UnprefixedName::parse_full(&name.local_name)?)
            }

            let value = xml::AttValue::from(value.as_str());

            let xml_attribute = xml::Attribute::Attribute {
                name: att_name,
                value,
            };
            attributes.push(xml_attribute);
        }

        let s_tag = xml::STag { name, attributes };

        s_tag.write(writer).await?;

        self.depth.push(element.name.clone());
        self.namespace_declarations
            .push(namespace_declarations.clone());
        Ok(())
    }

    pub async fn write_content(&mut self, content: &Content) -> Result<()> {
        match content {
            Content::Element(element) => self.write_element(element).await?,
            Content::Text(text) => {
                self.inner
                    .try_as_mut()?
                    .write_all(escape_str(text).as_bytes())
                    .await?
            }
            // TODO: comments and PI
            Content::PI => {}
            Content::Comment(_) => {}
        }
        Ok(())
    }

    pub async fn write_end(&mut self) -> Result<()> {
        let writer = self.inner.try_as_mut()?;
        if let Some(name) = &self.depth.pop() {
            let e_tag;
            let namespace_declarations_stack: Vec<_> =
                self.namespace_declarations.iter().flatten().collect();

            let prefix;
            if let Some(namespace) = &name.namespace {
                let name_namespace_declaration = namespace_declarations_stack
                    .iter()
                    .rfind(|namespace_declaration| namespace_declaration.namespace == *namespace)
                    .ok_or(Error::UndeclaredNamespace(namespace.clone()))?;
                prefix = name_namespace_declaration.prefix.as_ref();
            } else {
                prefix = None
            }

            if let Some(prefix) = &prefix {
                e_tag = xml::ETag {
                    name: xml::QName::PrefixedName(xml::PrefixedName {
                        prefix: xml::Prefix::parse_full(prefix)?,
                        local_part: xml::LocalPart::parse_full(&name.local_name)?,
                    }),
                };
            } else {
                e_tag = xml::ETag {
                    name: xml::QName::UnprefixedName(xml::UnprefixedName::parse_full(
                        &name.local_name,
                    )?),
                };
            }
            e_tag.write(writer).await?;
            self.namespace_declarations.pop();

            if self.depth.is_empty() {
                self.inner.end();
            }
            Ok(())
        } else {
            return Err(Error::NotInElement("".to_string()));
        }
    }
}

#[cfg(test)]
mod test {
    use crate::{
        reader::{test::*, Reader},
        writer::Writer,
    };

    #[tokio::test]
    async fn test_element_write() {
        let mock = MockAsyncReader::new(TEST_DOC);
        let mut reader = Reader::new(mock);
        let element = reader.read_element().await.unwrap();
        let stdout = tokio::io::stdout();
        let mut writer = Writer::new(stdout);
        writer.write_element(&element).await.unwrap();
    }
}