xdoc-macros 0.1.1

use std::collections::BTreeMap;

use proc_macro2::{Ident, Literal, TokenStream, TokenTree};
use quote::{format_ident, quote};

use crate::ast::{
    AttributeTemplate, AttributeValue, ComponentTemplate, ElementTemplate, RustPath, Template,
    TemplateNode, XmlName,
};

#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) struct CodegenError {
    message: String,
}

impl CodegenError {
    fn new(message: impl Into<String>) -> Self {
        Self {
            message: message.into(),
        }
    }

    pub(crate) fn message(&self) -> &str {
        &self.message
    }
}

type CodegenResult<T> = Result<T, CodegenError>;

pub(crate) fn generate_template(
    template: &Template,
    xdoc_path: TokenStream,
) -> CodegenResult<TokenStream> {
    let context = CodegenContext::new(xdoc_path);
    let scope = NamespaceScope::default();
    let nodes = template
        .nodes
        .iter()
        .map(|node| generate_node(node, &scope, &context))
        .collect::<CodegenResult<Vec<_>>>()?;
    let xdoc = context.xdoc();

    Ok(quote! {
        (|| -> #xdoc::core::XmlResult<#xdoc::macros::XmlTemplate> {
            let mut __xdoc_fragment = #xdoc::builder::fragment();
            #(
                __xdoc_fragment = __xdoc_fragment.child(#nodes)?;
            )*
            ::core::result::Result::Ok(#xdoc::macros::XmlTemplate::from_fragment(__xdoc_fragment))
        })()
    })
}

fn generate_node(
    node: &TemplateNode,
    scope: &NamespaceScope,
    context: &CodegenContext,
) -> CodegenResult<TokenStream> {
    match node {
        TemplateNode::Element(element) => generate_element(element, scope, context),
        TemplateNode::Text(text) => {
            let text = Literal::string(text);
            let xdoc = context.xdoc();
            Ok(quote! { #xdoc::builder::text(#text) })
        }
        TemplateNode::Expr(expr) => {
            let tokens = &expr.tokens;
            Ok(quote! { #tokens })
        }
        TemplateNode::Comment(comment) => {
            let comment = Literal::string(comment);
            let xdoc = context.xdoc();
            Ok(quote! { #xdoc::builder::XmlNode::Comment(::std::string::String::from(#comment)) })
        }
        TemplateNode::Component(component) => generate_component(component, scope, context),
    }
}

fn generate_element(
    element: &ElementTemplate,
    parent_scope: &NamespaceScope,
    context: &CodegenContext,
) -> CodegenResult<TokenStream> {
    let namespace_declarations = namespace_declarations(&element.attributes)?;
    let scope = parent_scope.with_declarations(&namespace_declarations);
    let element_start = generate_element_start(&element.name, &scope, context)?;
    let namespace_steps = namespace_declarations
        .iter()
        .map(generate_namespace_step)
        .collect::<CodegenResult<Vec<_>>>()?;
    let attributes = element
        .attributes
        .iter()
        .filter(|attribute| namespace_declaration(attribute).is_none())
        .map(|attribute| generate_attribute_step(attribute, &scope))
        .collect::<CodegenResult<Vec<_>>>()?;
    let children = element
        .children
        .iter()
        .map(|child| generate_child_step(child, &scope, context))
        .collect::<CodegenResult<Vec<_>>>()?;
    let xdoc = context.xdoc();

    Ok(quote! {
        (|| -> #xdoc::core::XmlResult<#xdoc::builder::ElementBuilder> {
            let mut __xdoc_element = #element_start?;
            #(
                #namespace_steps
            )*
            #(
                #attributes
            )*
            #(
                #children
            )*
            ::core::result::Result::Ok(__xdoc_element)
        })()
    })
}

fn generate_attribute_step(
    attribute: &AttributeTemplate,
    scope: &NamespaceScope,
) -> CodegenResult<TokenStream> {
    let value = generate_attribute_value(&attribute.value);

    match &attribute.name.prefix {
        Some(prefix) => {
            let uri = scope.resolve_prefix(prefix).ok_or_else(|| {
                CodegenError::new(format!(
                    "namespace prefix `{prefix}` is not declared for attribute `{}`",
                    attribute.name.lexical()
                ))
            })?;
            let prefix = Literal::string(prefix);
            let local = Literal::string(&attribute.name.local);
            let uri = Literal::string(uri);
            Ok(quote! {
                __xdoc_element = __xdoc_element.qualified_attr(#prefix, #local, #uri, #value)?;
            })
        }
        None => {
            let name = Literal::string(&attribute.name.local);
            Ok(quote! {
                __xdoc_element = __xdoc_element.attr(#name, #value)?;
            })
        }
    }
}

fn generate_component(
    component: &ComponentTemplate,
    scope: &NamespaceScope,
    context: &CodegenContext,
) -> CodegenResult<TokenStream> {
    let component_path = &component.path.tokens;
    let props_path = component_props_path(&component.path)?;
    let props = component
        .props
        .iter()
        .map(generate_component_prop)
        .collect::<CodegenResult<Vec<_>>>()?;

    if component.self_closing {
        return Ok(quote! {
            #component_path(#props_path {
                #(
                    #props
                )*
            })
        });
    }

    let children = component
        .children
        .iter()
        .map(|child| generate_child_fragment_step(child, scope, context))
        .collect::<CodegenResult<Vec<_>>>()?;
    let xdoc = context.xdoc();

    Ok(quote! {
        (|| -> #xdoc::core::XmlResult<_> {
            let mut __xdoc_component_children = #xdoc::builder::fragment();
            #(
                #children
            )*
            #component_path(
                #props_path {
                    #(
                        #props
                    )*
                },
                #xdoc::component::children(__xdoc_component_children)?,
            )
        })()
    })
}

fn generate_component_prop(attribute: &AttributeTemplate) -> CodegenResult<TokenStream> {
    if attribute.name.prefix.is_some() {
        return Err(CodegenError::new(format!(
            "component prop `{}` must be a simple Rust field identifier",
            attribute.name.lexical()
        )));
    }

    let field = format_ident!("{}", attribute.name.local);
    let value = generate_attribute_value(&attribute.value);
    Ok(quote! {
        #field: #value,
    })
}

fn generate_attribute_value(value: &AttributeValue) -> TokenStream {
    match value {
        AttributeValue::Literal(value) => {
            let value = Literal::string(value);
            quote! { #value }
        }
        AttributeValue::Expr(expr) => {
            let tokens = &expr.tokens;
            quote! { #tokens }
        }
    }
}

fn generate_child_step(
    child: &TemplateNode,
    scope: &NamespaceScope,
    context: &CodegenContext,
) -> CodegenResult<TokenStream> {
    let child = generate_node(child, scope, context)?;
    Ok(quote! {
        __xdoc_element = __xdoc_element.child(#child)?;
    })
}

fn generate_child_fragment_step(
    child: &TemplateNode,
    scope: &NamespaceScope,
    context: &CodegenContext,
) -> CodegenResult<TokenStream> {
    let child = generate_node(child, scope, context)?;
    Ok(quote! {
        __xdoc_component_children = __xdoc_component_children.child(#child)?;
    })
}

fn component_props_path(path: &RustPath) -> CodegenResult<TokenStream> {
    let mut tokens = path.tokens.clone().into_iter().collect::<Vec<_>>();
    let Some(last_ident_index) = tokens
        .iter()
        .rposition(|token| matches!(token, TokenTree::Ident(_)))
    else {
        return Err(CodegenError::new(format!(
            "component path `{}` must end with an identifier",
            path.source()
        )));
    };

    let TokenTree::Ident(component_ident) = &tokens[last_ident_index] else {
        unreachable!("rposition matched an identifier");
    };
    let props_ident = Ident::new(&format!("{}Props", component_ident), component_ident.span());
    tokens[last_ident_index] = TokenTree::Ident(props_ident);

    Ok(tokens.into_iter().collect())
}

fn generate_element_start(
    name: &XmlName,
    scope: &NamespaceScope,
    context: &CodegenContext,
) -> CodegenResult<TokenStream> {
    let xdoc = context.xdoc();
    match &name.prefix {
        Some(prefix) => {
            let uri = scope.resolve_prefix(prefix).ok_or_else(|| {
                CodegenError::new(format!(
                    "namespace prefix `{prefix}` is not declared for element `{}`",
                    name.lexical()
                ))
            })?;
            let prefix = Literal::string(prefix);
            let local = Literal::string(&name.local);
            let uri = Literal::string(uri);
            Ok(quote! { #xdoc::builder::ElementBuilder::qualified(#prefix, #local, #uri) })
        }
        None => {
            let local = Literal::string(&name.local);
            if let Some(uri) = scope.default_namespace.as_deref() {
                let uri = Literal::string(uri);
                Ok(quote! { #xdoc::builder::ElementBuilder::namespaced(#local, #uri) })
            } else {
                Ok(quote! { #xdoc::builder::element(#local) })
            }
        }
    }
}

#[derive(Debug, Clone)]
struct CodegenContext {
    xdoc_path: TokenStream,
}

impl CodegenContext {
    fn new(xdoc_path: TokenStream) -> Self {
        Self { xdoc_path }
    }

    fn xdoc(&self) -> &TokenStream {
        &self.xdoc_path
    }
}

fn generate_namespace_step(declaration: &NamespaceDeclaration) -> CodegenResult<TokenStream> {
    let uri = Literal::string(&declaration.uri);
    match &declaration.prefix {
        Some(prefix) => {
            let prefix = Literal::string(prefix);
            Ok(quote! {
                __xdoc_element = __xdoc_element.namespace(#prefix, #uri)?;
            })
        }
        None => Ok(quote! {
            __xdoc_element = __xdoc_element.default_namespace(#uri)?;
        }),
    }
}

fn namespace_declarations(
    attributes: &[AttributeTemplate],
) -> CodegenResult<Vec<NamespaceDeclaration>> {
    attributes
        .iter()
        .filter_map(namespace_declaration)
        .collect::<CodegenResult<Vec<_>>>()
}

fn namespace_declaration(
    attribute: &AttributeTemplate,
) -> Option<CodegenResult<NamespaceDeclaration>> {
    let prefix = match (&attribute.name.prefix, attribute.name.local.as_str()) {
        (None, "xmlns") => None,
        (Some(prefix), local) if prefix == "xmlns" => Some(local.to_owned()),
        _ => return None,
    };

    match &attribute.value {
        AttributeValue::Literal(uri) => Some(Ok(NamespaceDeclaration {
            prefix,
            uri: uri.clone(),
        })),
        AttributeValue::Expr(_) => Some(Err(CodegenError::new(format!(
            "namespace declaration `{}` must use a string literal",
            attribute.name.lexical()
        )))),
    }
}

#[derive(Debug, Clone, Default)]
struct NamespaceScope {
    default_namespace: Option<String>,
    prefixed: BTreeMap<String, String>,
}

impl NamespaceScope {
    fn with_declarations(&self, declarations: &[NamespaceDeclaration]) -> Self {
        let mut scope = self.clone();
        for declaration in declarations {
            match &declaration.prefix {
                Some(prefix) => {
                    scope
                        .prefixed
                        .insert(prefix.clone(), declaration.uri.clone());
                }
                None => {
                    scope.default_namespace = Some(declaration.uri.clone());
                }
            }
        }
        scope
    }

    fn resolve_prefix(&self, prefix: &str) -> Option<&str> {
        self.prefixed.get(prefix).map(String::as_str)
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
struct NamespaceDeclaration {
    prefix: Option<String>,
    uri: String,
}