dioxus_iconify/

generator.rs

use anyhow::{Context, Result};
use indoc::{formatdoc, indoc};
use std::collections::{BTreeMap, HashMap, HashSet};
use std::fs;
use std::path::{Path, PathBuf};

use crate::api::IconifyIcon;
use crate::naming::IconIdentifier;

const MOD_RS_TEMPLATE: &str = indoc! {r#"// Auto-generated by dioxus-iconify - DO NOT EDIT
    use dioxus::prelude::*;

    #[derive(Clone, Copy, PartialEq)]
    pub struct IconData {
        pub name: &'static str,
        pub body: &'static str,
        pub view_box: &'static str,
        pub width: &'static str,
        pub height: &'static str,
    }

    #[component]
    pub fn Icon(
        data: IconData,
        /// Optional size to set both width and height
        #[props(default, into)]
        size: String,
        /// Additional attributes to extend the svg element
        #[props(extends = GlobalAttributes)]
        attributes: Vec<Attribute>,
    ) -> Element {
        let (width, height) = if size.is_empty() {
            (data.width, data.height)
        } else {
            (size.as_str(), size.as_str())
        };

        rsx! {
            svg {
                view_box: "{data.view_box}",
                width: "{width}",
                height: "{height}",
                dangerous_inner_html: "{data.body}",
                ..attributes,
            }
        }
    }
    "#};

/// Represents a generated icon constant
#[derive(Debug, Clone)]
struct IconConst {
    name: String,
    full_icon_name: String,
    body: String,
    view_box: String,
    width: String,
    height: String,
}

impl IconConst {
    fn from_api_icon(identifier: &IconIdentifier, icon: &IconifyIcon) -> Self {
        Self {
            name: identifier.to_const_name(),
            full_icon_name: identifier.full_name.clone(),
            body: icon.body.clone(),
            view_box: icon
                .view_box
                .clone()
                .unwrap_or_else(|| "0 0 24 24".to_string()),
            width: icon.width.unwrap_or(24).to_string(),
            height: icon.height.unwrap_or(24).to_string(),
        }
    }

    fn to_rust_code(&self) -> String {
        // we use non upper case to be able to switch/wrap to struct or enum i the future
        formatdoc! { "

            #[allow(non_upper_case_globals)]
            pub const {}: IconData = IconData {{
                name: \"{}\",
                body: r#\"{}\"#,
                view_box: \"{}\",
                width: \"{}\",
                height: \"{}\",
            }};
            ",
            self.name,
            self.full_icon_name,
            self.body,
            self.view_box,
            self.width,
            self.height
        }
    }
}

/// Icon code generator
pub struct Generator {
    icons_dir: PathBuf,
}

impl Generator {
    /// Create a new generator with the specified icons directory
    pub fn new(icons_dir: PathBuf) -> Self {
        Self { icons_dir }
    }

    /// List all generated icons grouped by collection
    pub fn list_icons(&self) -> Result<BTreeMap<String, Vec<String>>> {
        let mut icons_by_collection: BTreeMap<String, Vec<String>> = BTreeMap::new();

        // Check if icons directory exists
        if !self.icons_dir.exists() {
            return Ok(icons_by_collection);
        }

        // Read all .rs files in the icons directory (except mod.rs)
        let entries = fs::read_dir(&self.icons_dir).context("Failed to read icons directory")?;

        for entry in entries {
            let entry = entry.context("Failed to read directory entry")?;
            let path = entry.path();

            // Skip if not a file or if it's mod.rs
            if !path.is_file() || path.file_name() == Some("mod.rs".as_ref()) {
                continue;
            }

            // Skip if not a .rs file
            if path.extension() != Some("rs".as_ref()) {
                continue;
            }

            // Parse the collection file
            let icons = self.parse_collection_file(&path)?;

            // Get collection name from file name
            if let Some(collection_name) = path.file_stem().and_then(|s| s.to_str()) {
                let icon_names: Vec<String> = icons
                    .values()
                    .map(|icon| icon.full_icon_name.clone())
                    .collect();

                if !icon_names.is_empty() {
                    icons_by_collection.insert(collection_name.to_string(), icon_names);
                }
            }
        }

        Ok(icons_by_collection)
    }

    /// Get all icon identifiers from generated files
    pub fn get_all_icon_identifiers(&self) -> Result<Vec<String>> {
        let icons_by_collection = self.list_icons()?;
        let mut all_icons = Vec::new();

        for icon_names in icons_by_collection.values() {
            all_icons.extend(icon_names.clone());
        }

        Ok(all_icons)
    }

    /// Initialize the icons directory with mod.rs if it doesn't exist
    pub fn init(&self) -> Result<()> {
        // Create icons directory if it doesn't exist
        if !self.icons_dir.exists() {
            fs::create_dir_all(&self.icons_dir).context("Failed to create icons directory")?;
        }

        // Create mod.rs if it doesn't exist
        let mod_rs_path = self.icons_dir.join("mod.rs");
        if !mod_rs_path.exists() {
            fs::write(&mod_rs_path, MOD_RS_TEMPLATE).context("Failed to create mod.rs")?;
        }

        Ok(())
    }

    /// Add icons to the generated code
    pub fn add_icons(&self, icons: &[(IconIdentifier, IconifyIcon)]) -> Result<()> {
        // Ensure icons directory and mod.rs exist
        self.init()?;

        // Group icons by collection
        let mut icons_by_collection: HashMap<String, Vec<(IconIdentifier, IconifyIcon)>> =
            HashMap::new();

        for (identifier, icon) in icons {
            icons_by_collection
                .entry(identifier.collection.clone())
                .or_default()
                .push((identifier.clone(), icon.clone()));
        }

        // Generate/update each collection file
        for (collection, collection_icons) in &icons_by_collection {
            self.update_collection_file(collection, collection_icons)?;
        }

        // Update mod.rs with module declarations
        self.update_mod_rs(&icons_by_collection.keys().cloned().collect::<Vec<_>>())?;

        Ok(())
    }

    /// Regenerate mod.rs with the latest template
    /// This is useful for updating the Icon component definition after CLI updates
    pub fn regenerate_mod_rs(&self) -> Result<()> {
        let mod_rs_path = self.icons_dir.join("mod.rs");

        // If mod.rs doesn't exist, just use init
        if !mod_rs_path.exists() {
            return self.init();
        }

        // Read existing mod.rs to extract module declarations
        let content = fs::read_to_string(&mod_rs_path).context("Failed to read mod.rs")?;

        // Extract existing module declarations
        let mut existing_modules: HashSet<String> = HashSet::new();
        for line in content.lines() {
            if line.trim().starts_with("pub mod ")
                && let Some(module_name) = line
                    .trim()
                    .strip_prefix("pub mod ")
                    .and_then(|s| s.strip_suffix(';'))
            {
                existing_modules.insert(module_name.trim().to_string());
            }
        }

        // Regenerate mod.rs with latest template
        let mut new_content = MOD_RS_TEMPLATE.to_string();
        new_content.push('\n');

        // Add module declarations in alphabetical order
        let mut sorted_modules: Vec<_> = existing_modules.iter().collect();
        sorted_modules.sort();

        for module in sorted_modules {
            new_content.push_str(&format!("pub mod {};\n", module));
        }

        fs::write(&mod_rs_path, new_content).context("Failed to update mod.rs")?;

        Ok(())
    }

    /// Update a collection file (e.g., mdi.rs) with new icons
    fn update_collection_file(
        &self,
        collection: &str,
        new_icons: &[(IconIdentifier, IconifyIcon)],
    ) -> Result<()> {
        let module_name = collection.replace('-', "_");
        let file_path = self.icons_dir.join(format!("{}.rs", module_name));

        // Read existing icons if file exists
        let mut existing_icons: BTreeMap<String, IconConst> = BTreeMap::new();
        if file_path.exists() {
            existing_icons = self.parse_collection_file(&file_path)?;
        }

        // Add/update new icons
        for (identifier, icon) in new_icons {
            let icon_const = IconConst::from_api_icon(identifier, icon);
            existing_icons.insert(icon_const.name.clone(), icon_const);
        }

        // Generate file content
        let content = self.generate_collection_file(collection, &existing_icons)?;

        // Write to file
        fs::write(&file_path, content)
            .context(format!("Failed to write collection file {:?}", file_path))?;

        println!(
            "✓ Updated {}.rs with {} icon(s)",
            module_name,
            new_icons.len()
        );

        Ok(())
    }

    /// Parse existing icons from a collection file
    fn parse_collection_file(&self, path: &Path) -> Result<BTreeMap<String, IconConst>> {
        let content =
            fs::read_to_string(path).context(format!("Failed to read file {:?}", path))?;

        let mut icons = BTreeMap::new();

        // Simple regex-free parsing: look for "pub const NAME: IconData = IconData {"
        // and extract the data between braces
        let lines: Vec<&str> = content.lines().collect();
        let mut i = 0;

        while i < lines.len() {
            let line = lines[i].trim();

            // Look for "pub const NAME: IconData"
            if line.starts_with("pub const ")
                && line.contains(": IconData")
                && let Some(name_end) = line.find(':')
            {
                let name = line[10..name_end].trim().to_string();

                // Parse the IconData struct (next several lines)
                if let Some(icon_const) = self.parse_icon_data(&lines, &mut i, &name) {
                    icons.insert(name, icon_const);
                }
            }

            i += 1;
        }

        Ok(icons)
    }

    /// Parse IconData struct from lines
    fn parse_icon_data(
        &self,
        lines: &[&str],
        index: &mut usize,
        const_name: &str,
    ) -> Option<IconConst> {
        let mut full_icon_name = String::new();
        let mut body = String::new();
        let mut view_box = String::new();
        let mut width = String::new();
        let mut height = String::new();

        // Look ahead to find all fields
        let mut j = *index;
        while j < lines.len() {
            let line = lines[j].trim();

            if line.contains("name:") {
                full_icon_name = extract_string_value(line);
            } else if line.contains("body:") {
                // Body might span multiple lines in raw string
                body = extract_raw_string_value(lines, &mut j);
            } else if line.contains("view_box:") {
                view_box = extract_string_value(line);
            } else if line.contains("width:") {
                width = extract_string_value(line);
            } else if line.contains("height:") {
                height = extract_string_value(line);
            }

            // End of struct
            if line.contains("};") {
                break;
            }

            j += 1;
        }

        *index = j;

        if !full_icon_name.is_empty() && !body.is_empty() {
            Some(IconConst {
                name: const_name.to_string(),
                full_icon_name,
                body,
                view_box,
                width,
                height,
            })
        } else {
            None
        }
    }

    /// Generate content for a collection file
    fn generate_collection_file(
        &self,
        collection: &str,
        icons: &BTreeMap<String, IconConst>,
    ) -> Result<String> {
        let mut content = formatdoc! { "
            // Auto-generated by dioxus-iconify - DO NOT EDIT
            // Collection: {}
            use super::IconData;
            ",
            collection
        };

        // Add each icon const in alphabetical order (BTreeMap maintains order)
        for icon_const in icons.values() {
            content.push_str(&icon_const.to_rust_code());
        }

        Ok(content)
    }

    /// Update mod.rs with module declarations
    fn update_mod_rs(&self, collections: &[String]) -> Result<()> {
        let mod_rs_path = self.icons_dir.join("mod.rs");

        // Read existing mod.rs
        let content = fs::read_to_string(&mod_rs_path).context("Failed to read mod.rs")?;

        // Extract existing module declarations
        let mut existing_modules: HashSet<String> = HashSet::new();
        for line in content.lines() {
            if line.trim().starts_with("pub mod ")
                && let Some(module_name) = line
                    .trim()
                    .strip_prefix("pub mod ")
                    .and_then(|s| s.strip_suffix(';'))
            {
                existing_modules.insert(module_name.trim().to_string());
            }
        }

        // Add new modules
        let mut needs_update = false;
        for collection in collections {
            let module_name = collection.replace('-', "_");
            if !existing_modules.contains(&module_name) {
                existing_modules.insert(module_name);
                needs_update = true;
            }
        }

        // Regenerate mod.rs if we have new modules
        if needs_update {
            let mut new_content = MOD_RS_TEMPLATE.to_string();
            new_content.push('\n');

            // Add module declarations in alphabetical order
            let mut sorted_modules: Vec<_> = existing_modules.iter().collect();
            sorted_modules.sort();

            for module in sorted_modules {
                new_content.push_str(&format!("pub mod {};\n", module));
            }

            fs::write(&mod_rs_path, new_content).context("Failed to update mod.rs")?;
        }

        Ok(())
    }
}

/// Extract a string value from a line like `name: "value",`
fn extract_string_value(line: &str) -> String {
    if let Some(start) = line.find('"')
        && let Some(end) = line.rfind('"')
        && end > start
    {
        return line[start + 1..end].to_string();
    }
    String::new()
}

/// Extract a raw string value that might span multiple lines
fn extract_raw_string_value(lines: &[&str], index: &mut usize) -> String {
    let line = lines[*index];

    // Look for r#"..."#
    if let Some(start) = line.find("r#\"") {
        let start_pos = start + 3;

        // Check if it ends on the same line
        if let Some(end) = line[start_pos..].find("\"#") {
            return line[start_pos..start_pos + end].to_string();
        }

        // Multi-line: collect until we find "#
        let mut result = line[start_pos..].to_string();
        *index += 1;

        while *index < lines.len() {
            let next_line = lines[*index];
            if let Some(end) = next_line.find("\"#") {
                result.push_str(&next_line[..end]);
                break;
            }
            result.push_str(next_line);
            result.push('\n');
            *index += 1;
        }

        result
    } else {
        String::new()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::api::IconifyIcon;
    use tempfile::TempDir;

    #[test]
    fn test_list_icons_empty_directory() -> Result<()> {
        let temp_dir = TempDir::new()?;
        let generator = Generator::new(temp_dir.path().join("icons"));

        let icons = generator.list_icons()?;
        assert!(
            icons.is_empty(),
            "Should return empty map for non-existent directory"
        );

        Ok(())
    }

    #[test]
    fn test_list_icons_with_generated_icons() -> Result<()> {
        let temp_dir = TempDir::new()?;
        let icons_dir = temp_dir.path().join("icons");
        let generator = Generator::new(icons_dir.clone());

        // Add some test icons
        let test_icon1 = IconifyIcon {
            body: r#"<path d="M10 20v-6h4v6h5v-8h3L12 3 2 12h3v8z"/>"#.to_string(),
            width: Some(24),
            height: Some(24),
            view_box: Some("0 0 24 24".to_string()),
        };

        let test_icon2 = IconifyIcon {
            body: r#"<circle cx="12" cy="12" r="10"/>"#.to_string(),
            width: Some(24),
            height: Some(24),
            view_box: Some("0 0 24 24".to_string()),
        };

        let identifier1 = IconIdentifier::parse("mdi:home")?;
        let identifier2 = IconIdentifier::parse("mdi:settings")?;
        let identifier3 = IconIdentifier::parse("heroicons:arrow-left")?;

        generator.add_icons(&[
            (identifier1, test_icon1.clone()),
            (identifier2, test_icon2.clone()),
            (identifier3, test_icon1.clone()),
        ])?;

        // List the icons
        let icons = generator.list_icons()?;

        // Should have 2 collections
        assert_eq!(icons.len(), 2, "Should have 2 collections");
        assert!(icons.contains_key("mdi"), "Should have mdi collection");
        assert!(
            icons.contains_key("heroicons"),
            "Should have heroicons collection"
        );

        // Check mdi collection has 2 icons
        let mdi_icons = icons.get("mdi").unwrap();
        assert_eq!(mdi_icons.len(), 2, "mdi should have 2 icons");
        assert!(mdi_icons.contains(&"mdi:home".to_string()));
        assert!(mdi_icons.contains(&"mdi:settings".to_string()));

        // Check heroicons collection has 1 icon
        let heroicons_icons = icons.get("heroicons").unwrap();
        assert_eq!(heroicons_icons.len(), 1, "heroicons should have 1 icon");
        assert!(heroicons_icons.contains(&"heroicons:arrow-left".to_string()));

        Ok(())
    }

    #[test]
    fn test_get_all_icon_identifiers() -> Result<()> {
        let temp_dir = TempDir::new()?;
        let icons_dir = temp_dir.path().join("icons");
        let generator = Generator::new(icons_dir.clone());

        // Test with no icons
        let empty_icons = generator.get_all_icon_identifiers()?;
        assert!(
            empty_icons.is_empty(),
            "Should return empty vec for no icons"
        );

        // Add some test icons
        let test_icon = IconifyIcon {
            body: r#"<path d="M10 20v-6h4v6h5v-8h3L12 3 2 12h3v8z"/>"#.to_string(),
            width: Some(24),
            height: Some(24),
            view_box: Some("0 0 24 24".to_string()),
        };

        let identifier1 = IconIdentifier::parse("mdi:home")?;
        let identifier2 = IconIdentifier::parse("mdi:settings")?;
        let identifier3 = IconIdentifier::parse("heroicons:arrow-left")?;

        generator.add_icons(&[
            (identifier1, test_icon.clone()),
            (identifier2, test_icon.clone()),
            (identifier3, test_icon.clone()),
        ])?;

        // Get all identifiers
        let all_icons = generator.get_all_icon_identifiers()?;

        // Should have 3 icons
        assert_eq!(all_icons.len(), 3, "Should have 3 icons");
        assert!(
            all_icons.contains(&"mdi:home".to_string()),
            "Should contain mdi:home"
        );
        assert!(
            all_icons.contains(&"mdi:settings".to_string()),
            "Should contain mdi:settings"
        );
        assert!(
            all_icons.contains(&"heroicons:arrow-left".to_string()),
            "Should contain heroicons:arrow-left"
        );

        Ok(())
    }

    #[test]
    fn test_regenerate_mod_rs_updates_template() -> Result<()> {
        let temp_dir = TempDir::new()?;
        let icons_dir = temp_dir.path().join("icons");
        let generator = Generator::new(icons_dir.clone());

        // Add some test icons
        let test_icon = IconifyIcon {
            body: r#"<path d="M10 20v-6h4v6h5v-8h3L12 3 2 12h3v8z"/>"#.to_string(),
            width: Some(24),
            height: Some(24),
            view_box: Some("0 0 24 24".to_string()),
        };

        let identifier1 = IconIdentifier::parse("mdi:home")?;
        let identifier2 = IconIdentifier::parse("heroicons:arrow-left")?;

        generator.add_icons(&[
            (identifier1, test_icon.clone()),
            (identifier2, test_icon.clone()),
        ])?;

        let mod_rs_path = icons_dir.join("mod.rs");
        assert!(mod_rs_path.exists(), "mod.rs should exist");

        // Simulate an old version by writing outdated content
        let old_content = r#"// Auto-generated by dioxus-iconify - DO NOT EDIT
use dioxus::prelude::*;

// OLD VERSION WITHOUT SIZE PARAMETER
#[component]
pub fn Icon(data: IconData) -> Element {
    rsx! { svg {} }
}

pub mod heroicons;
pub mod mdi;
"#;
        fs::write(&mod_rs_path, old_content)?;

        // Verify the old content is there
        let content_before = fs::read_to_string(&mod_rs_path)?;
        assert!(
            content_before.contains("OLD VERSION"),
            "Should have old version marker"
        );
        assert!(
            !content_before.contains("size: Option<String>"),
            "Should not have size parameter yet"
        );

        // Regenerate mod.rs
        generator.regenerate_mod_rs()?;

        // Verify the new content has the latest template
        let content_after = fs::read_to_string(&mod_rs_path)?;
        assert!(
            !content_after.contains("OLD VERSION"),
            "Should not have old version marker"
        );
        assert!(
            content_after.contains("size: String"),
            "Should have size parameter from latest template"
        );
        assert!(
            content_after.contains("pub mod heroicons;"),
            "Should preserve heroicons module"
        );
        assert!(
            content_after.contains("pub mod mdi;"),
            "Should preserve mdi module"
        );

        Ok(())
    }
}