// Rust implementation of xacro (XML macro language)

pub mod error;
pub mod eval;
pub mod lexer;
pub mod macros;
pub mod parser;
pub mod symbols;
pub mod urdf_validator;
pub mod utils;

#[cfg(feature = "python")]
pub mod python;

use std::collections::HashMap;
use std::path::{Path, PathBuf};
use xmltree::{Element, XMLNode};

use crate::error::{Result, XacroError};
use crate::symbols::SymbolTable;
use crate::urdf_validator::validate_urdf;

pub struct XacroProcessor {
    filestack: Vec<PathBuf>,
    macrostack: Vec<String>,
    all_includes: Vec<PathBuf>,
    #[allow(dead_code)]
    verbosity: u8,
    symbols: SymbolTable,
    macros: macros::MacroTable,
    format_output: bool,
    remove_first_joint: bool,
    validate_urdf: bool,
    validation_verbose: bool,
}

impl XacroProcessor {
    pub fn new(verbosity: u8) -> Self {
        Self {
            filestack: Vec::new(),
            macrostack: Vec::new(),
            all_includes: Vec::new(),
            verbosity,
            symbols: SymbolTable::new(),
            macros: macros::MacroTable::new(),
            format_output: true, // Format output by default
            remove_first_joint: false,
            validate_urdf: true, // Validation enabled by default
            validation_verbose: true,
        }
    }

    pub fn set_format_output(&mut self, format: bool) {
        self.format_output = format;
    }

    pub fn set_remove_first_joint(&mut self, remove: bool) {
        self.remove_first_joint = remove;
    }

    pub fn set_validate_urdf(&mut self, validate: bool) {
        self.validate_urdf = validate;
    }

    pub fn set_validation_verbose(&mut self, verbose: bool) {
        self.validation_verbose = verbose;
    }

    pub fn init_stacks(&mut self, file: Option<PathBuf>) {
        self.filestack.clear();
        if let Some(f) = file {
            self.filestack.push(f);
        }
        self.macrostack.clear();
    }

    pub fn process_file(
        &mut self,
        input_file: &Path,
        mappings: Option<HashMap<String, String>>,
    ) -> Result<Element> {
        self.init_stacks(Some(input_file.to_path_buf()));

        // Set substitution args
        if let Some(mappings) = mappings {
            self.symbols.set_substitution_args(mappings);
        }

        // Parse the document
        let content = std::fs::read_to_string(input_file).map_err(XacroError::Io)?;

        let mut doc = Element::parse(content.as_bytes())
            .map_err(|e| XacroError::Parse(format!("XML parsing error: {e}")))?;

        // Process the document
        self.process_doc(&mut doc)?;

        // Remove first joint if requested (after all processing is complete)
        if self.remove_first_joint {
            self.remove_first_joint_from_doc(&mut doc);
        }

        // Add xacro auto-generated banner
        let banner_text = format!(
            r#"
 =================================================================================== 
 |    This document was autogenerated by xacro from {:30} | 
 |    EDITING THIS FILE BY HAND IS NOT RECOMMENDED  {:30} | 
 =================================================================================== "#,
            input_file.display().to_string(),
            ""
        );

        // Create a new root with the banner comment
        let mut root = Element::new("root");
        root.children.push(XMLNode::Comment(banner_text));
        root.children.push(XMLNode::Element(doc));

        // Extract the actual document element
        if let Some(XMLNode::Element(doc_elem)) = root.children.pop() {
            // Validate URDF if requested
            if self.validate_urdf {
                let urdf_string = self.element_to_string(&doc_elem);
                match validate_urdf(&urdf_string, self.validation_verbose) {
                    Ok(validation_result) => {
                        if !validation_result.is_valid {
                            let error_msg = format!(
                                "URDF validation failed with {} error(s): {}",
                                validation_result.errors.len(),
                                validation_result.errors.join("; ")
                            );
                            return Err(XacroError::Parse(error_msg));
                        }
                    }
                    Err(e) => {
                        return Err(XacroError::Parse(format!("URDF validation error: {e}")));
                    }
                }
            }

            Ok(doc_elem)
        } else {
            unreachable!()
        }
    }

    pub fn process_string(
        &mut self,
        xml_string: &str,
        mappings: Option<HashMap<String, String>>,
    ) -> Result<Element> {
        self.init_stacks(None);

        // Set substitution args
        if let Some(mappings) = mappings {
            self.symbols.set_substitution_args(mappings);
        }

        // Parse the string
        let mut doc = Element::parse(xml_string.as_bytes())
            .map_err(|e| XacroError::Parse(format!("XML parsing error: {e}")))?;

        // Process the document
        self.process_doc(&mut doc)?;

        // Remove first joint if requested (after all processing is complete)
        if self.remove_first_joint {
            self.remove_first_joint_from_doc(&mut doc);
        }

        // Validate URDF if requested
        if self.validate_urdf {
            let urdf_string = self.element_to_string(&doc);
            match validate_urdf(&urdf_string, self.validation_verbose) {
                Ok(validation_result) => {
                    if !validation_result.is_valid {
                        let error_msg = format!(
                            "URDF validation failed with {} error(s): {}",
                            validation_result.errors.len(),
                            validation_result.errors.join("; ")
                        );
                        return Err(XacroError::Parse(error_msg));
                    }
                }
                Err(e) => {
                    return Err(XacroError::Parse(format!("URDF validation error: {e}")));
                }
            }
        }

        Ok(doc)
    }

    fn process_doc(&mut self, doc: &mut Element) -> Result<()> {
        // Apply xacro:targetNamespace as global xmlns (if defined)
        if let Some(target_ns) = doc.attributes.remove("xacro:targetNamespace") {
            doc.attributes.insert("xmlns".to_string(), target_ns);
        }

        // Process all elements
        self.eval_all(doc)?;

        Ok(())
    }

    fn eval_all(&mut self, element: &mut Element) -> Result<()> {
        // Two-phase processing:
        // Phase 1: Process includes and collect macro definitions
        self.process_includes_and_macros(element)?;

        // Phase 2: Expand macro calls and process everything else
        self.expand_macro_calls(element)?;

        Ok(())
    }

    fn process_includes_and_macros(&mut self, element: &mut Element) -> Result<()> {
        let mut i = 0;
        while i < element.children.len() {
            if let XMLNode::Element(child) = &mut element.children[i] {
                // Use more flexible matching for XML elements
                let element_local_name = if child.name.contains('}') {
                    child.name.split('}').next_back().unwrap_or(&child.name)
                } else if child.name.contains(':') {
                    child.name.split(':').next_back().unwrap_or(&child.name)
                } else {
                    &child.name
                };

                match element_local_name {
                    // Handle include elements
                    "include" => {
                        let included_children = self.process_include(child)?;
                        // Replace the include element with the included children
                        element.children.splice(i..=i, included_children);
                        continue;
                    }
                    // Handle property definitions
                    "property" => {
                        self.grab_property(child)?;
                        element.children.remove(i);
                        continue;
                    }
                    // Handle macro definitions
                    "macro" => {
                        self.grab_macro(child)?;
                        element.children.remove(i);
                        continue;
                    }
                    _ => {
                        // Recursively process children for includes and macros
                        self.process_includes_and_macros(child)?;
                    }
                }
            }
            i += 1;
        }
        Ok(())
    }

    fn expand_macro_calls(&mut self, element: &mut Element) -> Result<()> {
        let mut i = 0;
        while i < element.children.len() {
            match &mut element.children[i] {
                XMLNode::Element(child) => {
                    // Use more flexible matching for XML elements
                    let element_local_name = if child.name.contains('}') {
                        child.name.split('}').next_back().unwrap_or(&child.name)
                    } else if child.name.contains(':') {
                        child.name.split(':').next_back().unwrap_or(&child.name)
                    } else {
                        &child.name
                    };

                    // Check if it's a macro call (not a built-in xacro element)
                    let is_builtin_xacro = matches!(
                        element_local_name,
                        "include" | "property" | "macro" | "if" | "unless"
                    );

                    let is_macro_call = self.macros.contains(element_local_name);

                    if is_macro_call && !is_builtin_xacro {
                        match self.handle_macro_call(child) {
                            Ok(true) => {
                                // Macro was expanded
                                if child.name == "expanded_macro" {
                                    let expanded_children = std::mem::take(&mut child.children);
                                    element.children.splice(i..=i, expanded_children);
                                } else {
                                    // Single element replacement - process it recursively
                                    self.expand_macro_calls(child)?;
                                }
                                continue;
                            }
                            Ok(false) => {
                                // Macro not found - this is an error for xacro elements
                                return Err(XacroError::Parse(format!(
                                    "Unknown macro: {}",
                                    child.name
                                )));
                            }
                            Err(e) => return Err(e),
                        }
                    } else {
                        // Recursively process child
                        self.expand_macro_calls(child)?;
                    }
                }
                XMLNode::Text(text) => {
                    *text = self.eval_text(text)?;
                }
                _ => {}
            }
            i += 1;
        }

        // Evaluate attributes after processing children
        let mut attributes_to_remove = Vec::new();
        for (key, value) in element.attributes.iter_mut() {
            let evaluated = self.eval_text(value)?;
            // Handle special cases for different element types
            if element.name == "origin"
                && key == "rpy"
                && (evaluated == "None" || evaluated.is_empty())
            {
                // For origin elements, remove rpy attribute if it evaluates to None or empty
                // The URDF parser will default to [0, 0, 0] for missing rpy
                attributes_to_remove.push(key.clone());
            } else if element.name == "material"
                && key == "name"
                && (evaluated == "None" || evaluated.is_empty())
            {
                // For material elements, name is required - provide a default if missing
                *value = "default_material".to_string();
            } else if evaluated == "None" || evaluated.is_empty() {
                // For other cases, remove the attribute if it evaluated to None or empty
                attributes_to_remove.push(key.clone());
            } else {
                *value = evaluated;
            }
        }

        // Remove attributes that evaluated to None or empty
        for key in attributes_to_remove {
            element.attributes.remove(&key);
        }

        Ok(())
    }

    fn eval_text(&self, text: &str) -> Result<String> {
        let current_file = self.filestack.last().map(|p| p.as_path());
        eval::eval_text(text, &self.symbols, current_file)
    }

    fn process_include(&mut self, element: &Element) -> Result<Vec<XMLNode>> {
        let filename = element
            .attributes
            .get("filename")
            .ok_or_else(|| XacroError::Parse("include missing 'filename' attribute".into()))?;

        // Evaluate the filename to handle $(find package_name) expressions
        let evaluated_filename = self.eval_text(filename)?;

        // Resolve the file path
        let current_file = self.filestack.last().map(|p| p.as_path());
        let include_path = if std::path::Path::new(&evaluated_filename).is_absolute() {
            std::path::PathBuf::from(evaluated_filename)
        } else if let Some(current) = current_file {
            let parent = current
                .parent()
                .unwrap_or_else(|| std::path::Path::new("."));
            parent.join(evaluated_filename)
        } else {
            std::path::PathBuf::from(evaluated_filename)
        };

        // Check if file exists
        if !include_path.exists() {
            return Err(XacroError::Io(std::io::Error::new(
                std::io::ErrorKind::NotFound,
                format!("Include file not found: {}", include_path.display()),
            )));
        }

        // Parse the included file
        let mut included_doc = parser::parse_file(&include_path)?;

        // Add to filestack for proper context
        self.filestack.push(include_path.clone());

        // Process the included document to collect macros and properties only
        self.process_includes_and_macros(&mut included_doc)?;

        // Remove from filestack
        self.filestack.pop();

        // Track included file
        self.all_includes.push(include_path);

        // Return the children of the processed document
        Ok(included_doc.children)
    }

    fn grab_property(&mut self, element: &Element) -> Result<()> {
        let name = element
            .attributes
            .get("name")
            .ok_or_else(|| XacroError::Parse("property missing 'name' attribute".into()))?;

        let value = element.attributes.get("value");
        let default = element.attributes.get("default");

        if let Some(value) = value {
            self.symbols.set(name.clone(), value.clone());
        } else if let Some(default) = default {
            if !self.symbols.contains(name) {
                self.symbols.set(name.clone(), default.clone());
            }
        }

        Ok(())
    }

    fn grab_macro(&mut self, element: &Element) -> Result<()> {
        let name = element
            .attributes
            .get("name")
            .ok_or_else(|| XacroError::Parse("macro missing 'name' attribute".into()))?;

        let params = element
            .attributes
            .get("params")
            .map(|s| s.as_str())
            .unwrap_or("");

        // Convert to string and re-parse to avoid any evaluation during cloning
        let element_xml = parser::element_to_string(element);
        let mut macro_element = parser::parse_string(&element_xml)?;

        // Remove the name and params attributes from the macro body
        macro_element.attributes.remove("name");
        macro_element.attributes.remove("params");

        let macro_def = macros::Macro::new(name.clone(), params, macro_element);
        self.macros.insert(name.clone(), macro_def);

        Ok(())
    }

    #[allow(dead_code)]
    fn process_conditional(&mut self, element: &mut Element) -> Result<bool> {
        let value = element
            .attributes
            .get("value")
            .ok_or_else(|| XacroError::Parse("conditional missing 'value' attribute".into()))?;

        let evaluated = self.eval_text(value)?;
        let keep = utils::get_boolean_value(&evaluated)?;

        let keep = if element.name == "xacro:unless" {
            !keep
        } else {
            keep
        };

        if keep {
            self.eval_all(element)?;
        }

        Ok(keep)
    }

    fn handle_macro_call(&mut self, element: &mut Element) -> Result<bool> {
        // Extract macro name using same logic as in eval_all
        let macro_name = if element.name.contains('}') {
            element.name.split('}').next_back().unwrap_or(&element.name)
        } else if element.name.contains(':') {
            element.name.split(':').next_back().unwrap_or(&element.name)
        } else {
            &element.name
        };

        if let Some(macro_def) = self.macros.get(macro_name).cloned() {
            // Expand the macro
            let expanded_elements = self.expand_macro(&macro_def, element)?;

            // Replace the current element with expanded content
            // If there's only one element, replace directly; otherwise use children
            if expanded_elements.len() == 1 {
                *element = expanded_elements.into_iter().next().unwrap();
            } else {
                // Multiple elements: replace with a container element
                element.name = "expanded_macro".to_string();
                element.attributes.clear();
                element.children = expanded_elements
                    .into_iter()
                    .map(xmltree::XMLNode::Element)
                    .collect();
            }

            Ok(true)
        } else {
            Err(XacroError::Parse(format!(
                "Unknown macro: {}",
                element.name
            )))
        }
    }

    fn expand_macro(
        &mut self,
        macro_def: &macros::Macro,
        call_element: &Element,
    ) -> Result<Vec<Element>> {
        // Build parameter mapping from call attributes
        let mut param_values = HashMap::new();

        // Add default values first
        for (param, default) in &macro_def.defaults {
            param_values.insert(param.clone(), default.clone());
        }

        // Override with call attributes
        for (attr_name, attr_value) in &call_element.attributes {
            if macro_def.has_param(attr_name) {
                let evaluated_value = self.eval_text(attr_value)?;
                param_values.insert(attr_name.clone(), evaluated_value);
            }
        }

        // Clone macro body and substitute parameters
        let mut expanded_body = macro_def.body.clone();

        // Add macro parameters to symbol table temporarily for evaluation
        let mut temp_symbols_added = Vec::new();
        for (param, value) in &param_values {
            if !self.symbols.contains(param) {
                self.symbols.set(param.clone(), value.clone());
                temp_symbols_added.push(param.clone());
            }
        }

        self.substitute_params_in_element(&mut expanded_body, &param_values)?;

        // Handle child elements (like <origin> inside macro calls)
        for child in &call_element.children {
            if let xmltree::XMLNode::Element(child_elem) = child {
                expanded_body
                    .children
                    .push(xmltree::XMLNode::Element(child_elem.clone()));
            }
        }

        // Return the children of the macro body (the actual content)
        let mut result = Vec::new();

        for child in expanded_body.children {
            if let xmltree::XMLNode::Element(mut element) = child {
                // CRITICAL: Recursively process the expanded content
                // This is what was missing - expanded macros need to be evaluated too!
                // Keep symbols available during recursive evaluation
                self.eval_all(&mut element)?;
                result.push(element);
            }
        }

        // Don't remove macro parameters immediately - they might be needed by nested macros
        // Instead, let them persist in the symbol table
        // This is a more permissive approach that matches xacro behavior

        // Note: Joint removal is now handled after all processing

        Ok(result)
    }

    fn remove_first_joint_from_doc(&mut self, doc: &mut Element) {
        // Simply remove the first joint element found in the document
        self.remove_first_joint_simple(doc);
    }

    #[allow(dead_code)]
    fn remove_first_joint_from_expansion(&mut self, elements: &mut Vec<Element>) {
        // Remove the first joint element from a macro expansion
        for i in 0..elements.len() {
            if elements[i].name == "joint" {
                elements.remove(i);
                return;
            }
        }
    }

    #[allow(clippy::only_used_in_recursion)]
    fn remove_first_joint_simple(&mut self, element: &mut Element) -> bool {
        // Find and remove the first joint element in document order
        for i in 0..element.children.len() {
            if let xmltree::XMLNode::Element(child_element) = &element.children[i] {
                if child_element.name == "joint" {
                    element.children.remove(i);
                    return true; // Found and removed
                }
            }
        }

        // If no joint found at this level, search children recursively
        for child in &mut element.children {
            if let xmltree::XMLNode::Element(child_element) = child {
                if self.remove_first_joint_simple(child_element) {
                    return true; // Found and removed in child
                }
            }
        }

        false // No joint found
    }

    fn substitute_params_in_element(
        &mut self,
        element: &mut Element,
        params: &HashMap<String, String>,
    ) -> Result<()> {
        // Substitute in attributes
        for (_, value) in element.attributes.iter_mut() {
            *value = self.substitute_params_in_text(value, params)?;
        }

        // Substitute in text content
        for child in &mut element.children {
            match child {
                xmltree::XMLNode::Element(child_elem) => {
                    self.substitute_params_in_element(child_elem, params)?;
                }
                xmltree::XMLNode::Text(text) => {
                    *text = self.substitute_params_in_text(text, params)?;
                }
                _ => {}
            }
        }

        Ok(())
    }

    fn substitute_params_in_text(
        &mut self,
        text: &str,
        params: &HashMap<String, String>,
    ) -> Result<String> {
        let mut result = text.to_string();

        // Replace ${param} with parameter values
        for (param, value) in params {
            let pattern = format!("${{{param}}}");
            result = result.replace(&pattern, value);
        }

        // After parameter substitution, always evaluate expressions
        // This handles global properties and math expressions
        result = self.eval_text(&result)?;

        Ok(result)
    }

    pub fn element_to_string(&self, element: &Element) -> String {
        if self.format_output {
            let formatted = self.element_to_formatted_string(element, 0);
            format!("<?xml version=\"1.0\" encoding=\"UTF-8\"?>\n{formatted}")
        } else {
            crate::parser::element_to_string(element)
        }
    }

    #[allow(clippy::only_used_in_recursion)]
    fn element_to_formatted_string(&self, element: &Element, indent_level: usize) -> String {
        let mut result = String::new();
        let indent = "  ".repeat(indent_level);

        // Start tag
        result.push_str(&format!("{}<{}", indent, element.name));

        // Attributes
        for (key, value) in &element.attributes {
            result.push_str(&format!(" {key}=\"{value}\""));
        }

        if element.children.is_empty() {
            // Self-closing tag
            result.push_str(" />");
        } else {
            result.push('>');

            // Check if we have only text content
            let only_text = element
                .children
                .iter()
                .all(|child| matches!(child, XMLNode::Text(_)));

            if only_text && element.children.len() == 1 {
                // Single text node - no newlines
                if let Some(XMLNode::Text(text)) = element.children.first() {
                    result.push_str(text);
                }
            } else {
                // Mixed or multiple content - use newlines and indentation
                result.push('\n');

                for child in &element.children {
                    match child {
                        XMLNode::Element(child_elem) => {
                            result.push_str(
                                &self.element_to_formatted_string(child_elem, indent_level + 1),
                            );
                            result.push('\n');
                        }
                        XMLNode::Text(text) => {
                            let trimmed = text.trim();
                            if !trimmed.is_empty() {
                                result.push_str(&format!(
                                    "{}{}\n",
                                    "  ".repeat(indent_level + 1),
                                    trimmed
                                ));
                            }
                        }
                        XMLNode::Comment(comment) => {
                            result.push_str(&format!(
                                "{}<!--{}-->\n",
                                "  ".repeat(indent_level + 1),
                                comment
                            ));
                        }
                        _ => {} // Handle other node types if needed
                    }
                }

                result.push_str(&indent);
            }

            // End tag
            result.push_str(&format!("</{}>", element.name));
        }

        result
    }
}
