roblox-slang 3.0.1

use crate::parser::Translation;
use crate::utils::plurals;
use anyhow::Result;
use std::collections::{BTreeMap, HashSet};

use super::luau::{
    format_generated_luau, sanitize_luau_identifier, sanitized_namespace_prefixes,
    sanitized_path_parts,
};

pub fn generate_type_definitions(
    translations: &[Translation],
    base_locale: &str,
) -> Result<String> {
    let mut code = String::new();

    code.push_str("-- Generated by roblox-slang. DO NOT MODIFY BY HAND!\n\n");

    let base_translations: Vec<_> = translations
        .iter()
        .filter(|t| t.locale == base_locale)
        .collect();

    if base_translations.is_empty() {
        code.push_str("export type Translations = {}\n");
        return Ok(format_generated_luau(&code));
    }

    code.push_str("export type Translations = {\n");
    code.push_str("    new: (locale: string?) -> TranslationsInstance,\n");
    code.push_str("    newForPlayer: (player: Player) -> TranslationsInstance,\n");
    code.push_str("    detectLocale: (player: Player) -> string,\n");
    code.push_str("}\n\n");

    code.push_str("export type TranslationsInstance = {\n");

    code.push_str("    _locale: string,\n");
    code.push_str("    _translator: any,\n");
    code.push_str("    _localeChangedCallbacks: { any },\n\n");

    code.push_str("    setLocale: (self: TranslationsInstance, locale: string) -> (),\n");
    code.push_str("    getLocale: (self: TranslationsInstance) -> string,\n");
    code.push_str("    onLocaleChanged: (self: TranslationsInstance, callback: (newLocale: string, oldLocale: string) -> ()) -> (),\n");
    code.push_str("    getAsset: (self: TranslationsInstance, assetKey: string) -> string,\n\n");

    let mut plural_base_keys: HashSet<String> = HashSet::new();
    let mut regular_translations = Vec::new();

    for translation in &base_translations {
        if plurals::is_plural_key(&translation.key) {
            let base_key = plurals::extract_base_key(&translation.key);
            plural_base_keys.insert(base_key);
        } else {
            regular_translations.push(*translation);
        }
    }

    regular_translations.sort_by(|a, b| a.key.cmp(&b.key));

    for translation in &regular_translations {
        let method_name = sanitize_luau_identifier(&translation.key);
        let params = super::luau::extract_parameters(&translation.value);

        if params.is_empty() {
            code.push_str(&format!(
                "    {}: (self: TranslationsInstance) -> string,\n",
                method_name
            ));
        } else {
            code.push_str(&format!(
                "    {}: (self: TranslationsInstance, params: {{}}) -> string,\n",
                method_name
            ));
        }
    }

    let mut plural_keys_sorted: Vec<_> = plural_base_keys.iter().collect();
    plural_keys_sorted.sort();

    for base_key in &plural_keys_sorted {
        let method_name = sanitize_luau_identifier(base_key);
        code.push_str(&format!(
            "    {}: (self: TranslationsInstance, count: number, params: {{}}?) -> string,\n",
            method_name
        ));
    }

    code.push('\n');

    let namespaces = build_namespace_tree(&regular_translations, &plural_base_keys);
    generate_namespace_types(
        &mut code,
        &namespaces,
        &regular_translations,
        &plural_base_keys,
    );

    code.push_str("}\n");

    Ok(format_generated_luau(&code))
}

fn build_namespace_tree(
    translations: &[&Translation],
    plural_base_keys: &HashSet<String>,
) -> HashSet<String> {
    let mut namespaces = HashSet::new();

    for translation in translations {
        for namespace in sanitized_namespace_prefixes(&translation.key) {
            namespaces.insert(namespace);
        }
    }

    for base_key in plural_base_keys {
        for namespace in sanitized_namespace_prefixes(base_key) {
            namespaces.insert(namespace);
        }
    }

    namespaces
}

fn generate_namespace_types(
    code: &mut String,
    namespaces: &HashSet<String>,
    translations: &[&Translation],
    plural_base_keys: &HashSet<String>,
) {
    let mut root = NamespaceTypeNode::default();

    for namespace in namespaces {
        let parts: Vec<String> = namespace.split('.').map(ToOwned::to_owned).collect();
        add_namespace_path(&mut root, &parts);
    }

    for translation in translations {
        let parts = sanitized_path_parts(&translation.key);
        let params = super::luau::extract_parameters(&translation.value);
        insert_namespace_method(
            &mut root,
            &parts,
            NamespaceMethod::Regular {
                has_params: !params.is_empty(),
            },
        );
    }

    for base_key in plural_base_keys {
        let parts = sanitized_path_parts(base_key);
        insert_namespace_method(&mut root, &parts, NamespaceMethod::Plural);
    }

    render_namespace_node(code, &root, 1);
}

#[derive(Default)]
struct NamespaceTypeNode {
    namespaces: BTreeMap<String, NamespaceTypeNode>,
    methods: BTreeMap<String, NamespaceMethod>,
}

enum NamespaceMethod {
    Regular { has_params: bool },
    Plural,
}

fn add_namespace_path(root: &mut NamespaceTypeNode, parts: &[String]) {
    let mut node = root;
    for part in parts {
        node = node.namespaces.entry(part.clone()).or_default();
    }
}

fn insert_namespace_method(
    root: &mut NamespaceTypeNode,
    parts: &[String],
    method_type: NamespaceMethod,
) {
    if parts.len() < 2 {
        return;
    }

    let mut node = root;
    for part in &parts[..parts.len() - 1] {
        node = node.namespaces.entry(part.clone()).or_default();
    }

    node.methods
        .insert(parts[parts.len() - 1].clone(), method_type);
}

fn render_namespace_node(code: &mut String, node: &NamespaceTypeNode, indent: usize) {
    for (method, method_type) in &node.methods {
        render_namespace_method(code, method, method_type, indent);
    }

    for (namespace, child) in &node.namespaces {
        let indentation = "    ".repeat(indent);
        code.push_str(&format!("{}{}: {{\n", indentation, namespace));
        render_namespace_node(code, child, indent + 1);
        code.push_str(&format!("{}}},\n", indentation));
    }
}

fn render_namespace_method(
    code: &mut String,
    method: &str,
    method_type: &NamespaceMethod,
    indent: usize,
) {
    let indentation = "    ".repeat(indent);
    match method_type {
        NamespaceMethod::Regular { has_params: false } => {
            code.push_str(&format!(
                "{}{}: (self: any) -> string,\n",
                indentation, method
            ));
        }
        NamespaceMethod::Regular { has_params: true } => {
            code.push_str(&format!(
                "{}{}: (self: any, params: {{}}) -> string,\n",
                indentation, method
            ));
        }
        NamespaceMethod::Plural => {
            code.push_str(&format!(
                "{}{}: (self: any, count: number, params: {{}}?) -> string,\n",
                indentation, method
            ));
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_build_namespace_tree() {
        let translations = [
            Translation {
                key: "ui.buttons.buy".to_string(),
                value: "Buy".to_string(),
                locale: "en".to_string(),
                context: None,
            },
            Translation {
                key: "ui.labels.welcome".to_string(),
                value: "Welcome".to_string(),
                locale: "en".to_string(),
                context: None,
            },
        ];

        let refs: Vec<_> = translations.iter().collect();
        let plural_base_keys = HashSet::new();
        let namespaces = build_namespace_tree(&refs, &plural_base_keys);

        assert!(namespaces.contains("ui"));
        assert!(namespaces.contains("ui.buttons"));
        assert!(namespaces.contains("ui.labels"));
    }

    #[test]
    fn test_generate_type_definitions_with_plurals() {
        let translations = vec![
            Translation {
                key: "ui.messages.items(one)".to_string(),
                value: "{count} item".to_string(),
                locale: "en".to_string(),
                context: None,
            },
            Translation {
                key: "ui.messages.items(other)".to_string(),
                value: "{count} items".to_string(),
                locale: "en".to_string(),
                context: None,
            },
        ];

        let code = generate_type_definitions(&translations, "en").unwrap();

        assert!(code.contains("newForPlayer: (player: Player) -> TranslationsInstance"));
        assert!(code.contains("detectLocale: (player: Player) -> string"));

        assert!(code.contains(
            "ui_messages_items: (self: TranslationsInstance, count: number, params: {}?) -> string"
        ));

        assert!(code.contains("items: (self: any, count: number, params: {}?) -> string"));

        assert!(!code.contains("items(one)"));
        assert!(!code.contains("items(other)"));
    }

    #[test]
    fn test_type_definitions_sanitize_namespace_segments() {
        let translations = [
            Translation {
                key: "ui.buttons.buy-now".to_string(),
                value: "Buy Now".to_string(),
                locale: "en".to_string(),
                context: None,
            },
            Translation {
                key: "ui.return.label".to_string(),
                value: "Go Back".to_string(),
                locale: "en".to_string(),
                context: None,
            },
        ];

        let code = generate_type_definitions(&translations, "en").unwrap();

        assert!(
            code.contains("buy_now:"),
            "hyphenated segment must be sanitized in type defs"
        );
        assert!(
            code.contains("_return:"),
            "keyword segment must be sanitized in type defs"
        );

        assert!(
            code.contains("(self: any)"),
            "namespace methods must use self: any"
        );
        assert!(
            !code.contains("buy-now"),
            "raw hyphenated segment must not appear in type defs"
        );
    }

    #[test]
    fn test_type_definitions_dedupe_sanitized_namespace_roots() {
        let translations = [
            Translation {
                key: "ui-menu.a.open".to_string(),
                value: "Open".to_string(),
                locale: "en".to_string(),
                context: None,
            },
            Translation {
                key: "ui_menu.b.close".to_string(),
                value: "Close".to_string(),
                locale: "en".to_string(),
                context: None,
            },
        ];

        let code = generate_type_definitions(&translations, "en").unwrap();

        assert_eq!(
            code.matches("\tui_menu: {\n").count(),
            1,
            "sanitized top-level namespace must only be emitted once"
        );
        assert!(code.contains("\t\ta: {\n"));
        assert!(code.contains("\t\tb: {\n"));
        assert!(code.contains("\t\t\topen: (self: any) -> string,"));
        assert!(code.contains("\t\t\tclose: (self: any) -> string,"));
    }

    #[test]
    fn test_type_definitions_support_deep_namespaces() {
        let translations = [
            Translation {
                key: "ui.buttons.primary.buy".to_string(),
                value: "Buy".to_string(),
                locale: "en".to_string(),
                context: None,
            },
            Translation {
                key: "ui.buttons.primary.items(one)".to_string(),
                value: "{count} item".to_string(),
                locale: "en".to_string(),
                context: None,
            },
            Translation {
                key: "ui.buttons.primary.items(other)".to_string(),
                value: "{count} items".to_string(),
                locale: "en".to_string(),
                context: None,
            },
        ];

        let code = generate_type_definitions(&translations, "en").unwrap();

        assert!(code.contains("\tui: {\n"));
        assert!(code.contains("\t\tbuttons: {\n"));
        assert!(code.contains("\t\t\tprimary: {\n"));
        assert!(code.contains("\t\t\t\tbuy: (self: any) -> string,"));
        assert!(code.contains("\t\t\t\titems: (self: any, count: number, params: {}?) -> string,"));
    }
}