storybook/src/resolve/merge.rs

//! Template composition and merge engine
//!
//! Handles two types of template composition:
//! 1. Template includes (vertical composition): `template Person { include
//!    Human ... }`
//! 2. Character template inheritance (horizontal composition): `character
//!    Martha from Person, Worker { ... }`
//!
//! Also handles legacy @BaseTemplate { ... } syntax for template overrides
//! with:
//! - Set operations (field: value) - replace or add field
//! - Remove operations (remove field) - delete field
//! - Append operations (append field: value) - add new field (error if exists)

use std::collections::HashSet;

use crate::{
    resolve::{
        names::NameTable,
        ResolveError,
        Result,
    },
    syntax::ast::{
        Character,
        Declaration,
        Field,
        OverrideOp,
        Template,
        Value,
    },
};

// ===== Template Composition =====

/// Resolve a template by recursively merging all its includes
///
/// Algorithm:
/// 1. Recursively resolve all included templates (depth-first)
/// 2. Merge included template fields (later includes override earlier ones)
/// 3. Add the template's own fields on top
///
/// Returns the fully merged fields for this template
pub fn resolve_template_includes(
    template: &Template,
    declarations: &[Declaration],
    name_table: &NameTable,
    visited: &mut HashSet<String>,
) -> Result<Vec<Field>> {
    // Detect circular includes
    if !visited.insert(template.name.clone()) {
        return Err(ResolveError::CircularDependency {
            cycle: format!(
                "Circular template include detected: {} -> {}",
                visited.iter().cloned().collect::<Vec<_>>().join(" -> "),
                template.name
            ),
        });
    }

    let mut merged_fields = Vec::new();

    // Resolve all includes first
    for include_name in &template.includes {
        // Look up the included template
        let entry = name_table
            .lookup(std::slice::from_ref(include_name))
            .ok_or_else(|| ResolveError::NameNotFound {
                name: include_name.clone(),
                suggestion: name_table.find_suggestion(include_name),
            })?;

        // Get the template declaration
        let included_template = match &declarations[entry.decl_index] {
            | Declaration::Template(t) => t,
            | _ => {
                return Err(ResolveError::ValidationError {
                    message: format!(
                        "Cannot include '{}': it's not a template",
                        include_name
                    ),
                    help: Some(format!(
                        "The 'include' keyword can only reference templates. '{}' is a different type of declaration. Make sure you're including the correct name and that it refers to a template.",
                        include_name
                    )),
                });
            },
        };

        // Recursively resolve the included template
        let included_fields =
            resolve_template_includes(included_template, declarations, name_table, visited)?;

        // Merge included fields (replacing any existing fields with same name)
        merged_fields = merge_field_lists(merged_fields, included_fields);
    }

    // Add this template's own fields on top
    merged_fields = merge_field_lists(merged_fields, template.fields.clone());

    // Remove this template from visited set (allow it to be used in other branches)
    visited.remove(&template.name);

    Ok(merged_fields)
}

/// Merge character templates into character fields
///
/// Algorithm:
/// 1. Resolve each template (which may itself include other templates)
/// 2. Merge templates left to right (later templates override earlier ones)
/// 3. Add character's own fields on top
/// 4. If any template is strict, validate that all its fields are concrete
///
/// Returns the fully merged fields for this character
pub fn merge_character_templates(
    character: &Character,
    declarations: &[Declaration],
    name_table: &NameTable,
) -> Result<Vec<Field>> {
    let mut merged_fields = Vec::new();
    let mut strict_templates = Vec::new();

    // If character has templates, merge them
    if let Some(template_names) = &character.template {
        for template_name in template_names {
            // Look up the template
            let entry = name_table
                .lookup(std::slice::from_ref(template_name))
                .ok_or_else(|| ResolveError::NameNotFound {
                    name: template_name.clone(),
                    suggestion: name_table.find_suggestion(template_name),
                })?;

            // Get the template declaration
            let template = match &declarations[entry.decl_index] {
                | Declaration::Template(t) => t,
                | _ => {
                    return Err(ResolveError::ValidationError {
                        message: format!(
                            "Character '{}' cannot inherit from '{}': it's not a template",
                            character.name, template_name
                        ),
                        help: Some(format!(
                            "The 'from' keyword can only reference templates. '{}' is a different type of declaration. Make sure you're inheriting from the correct name and that it refers to a template.",
                            template_name
                        )),
                    });
                },
            };

            // Track strict templates for validation
            if template.strict {
                strict_templates.push(template_name.clone());
            }

            // Resolve template (which handles includes recursively)
            let mut visited = HashSet::new();
            let template_fields =
                resolve_template_includes(template, declarations, name_table, &mut visited)?;

            // Merge template fields into accumulated fields
            merged_fields = merge_field_lists(merged_fields, template_fields);
        }
    }

    // Add character's own fields on top
    merged_fields = merge_field_lists(merged_fields, character.fields.clone());

    // Validate strict mode: all strict template fields must have concrete values
    if !strict_templates.is_empty() {
        validate_strict_templates(&character.name, &merged_fields, &strict_templates)?;
    }

    Ok(merged_fields)
}

/// Merge two field lists, with fields from the second list overriding the first
fn merge_field_lists(base: Vec<Field>, override_fields: Vec<Field>) -> Vec<Field> {
    let mut merged = base;

    for field in override_fields {
        // If field exists, replace it; otherwise add it
        if let Some(existing) = merged.iter_mut().find(|f| f.name == field.name) {
            existing.value = field.value.clone();
            existing.span = field.span.clone();
        } else {
            merged.push(field);
        }
    }

    merged
}

/// Validate that strict template requirements are met
///
/// For strict templates, all fields must have concrete values (not ranges)
fn validate_strict_templates(
    character_name: &str,
    fields: &[Field],
    strict_templates: &[String],
) -> Result<()> {
    for field in fields {
        if let Value::Range(_, _) = &field.value {
            return Err(ResolveError::ValidationError {
                message: format!(
                    "Character '{}' inherits from strict template(s) {}, but field '{}' has a range value instead of a concrete value",
                    character_name,
                    strict_templates.join(", "),
                    field.name
                ),
                help: Some(format!(
                    "Strict templates require all fields to have concrete values. Replace the range in '{}' with a specific value. For example, instead of '18..65', use a specific age like '34'.",
                    field.name
                )),
            });
        }
    }

    Ok(())
}

// ===== Legacy Override System =====

/// Apply an override to a base template's fields
///
/// This performs a structural merge:
/// 1. Start with all fields from base
/// 2. Apply each override operation in order
/// 3. Return merged field list
pub fn apply_override(base_fields: Vec<Field>, override_ops: &[OverrideOp]) -> Result<Vec<Field>> {
    let mut merged = base_fields;

    for op in override_ops {
        match op {
            | OverrideOp::Set(field) => {
                // Replace existing field or add new one
                if let Some(existing) = merged.iter_mut().find(|f| f.name == field.name) {
                    existing.value = field.value.clone();
                    existing.span = field.span.clone();
                } else {
                    merged.push(field.clone());
                }
            },

            | OverrideOp::Remove(name) => {
                // Remove field if it exists
                merged.retain(|f| f.name != *name);
            },

            | OverrideOp::Append(field) => {
                // Add field only if it doesn't exist
                if merged.iter().any(|f| f.name == field.name) {
                    return Err(ResolveError::ValidationError {
                        message: format!(
                            "Cannot append field '{}': field already exists",
                            field.name
                        ),
                        help: Some(format!(
                            "The 'append' operation is used to add new fields that don't exist in the base template. The field '{}' already exists. Use 'set' instead to update an existing field, or use a different field name.",
                            field.name
                        )),
                    });
                }
                merged.push(field.clone());
            },
        }
    }

    Ok(merged)
}

/// Recursively resolve overrides in a value
///
/// If the value contains an Override, look up the base template
/// and apply the override operations
pub fn resolve_value_overrides(value: &Value, name_table: &NameTable) -> Result<Value> {
    match value {
        | Value::Override(override_spec) => {
            // Look up the base template
            let _base_entry = name_table.lookup(&override_spec.base).ok_or_else(|| {
                ResolveError::NameNotFound {
                    name: override_spec.base.join("::"),
                    suggestion: name_table
                        .find_suggestion(override_spec.base.last().unwrap_or(&String::new())),
                }
            })?;

            // For now, we'll return an error since we need the actual template fields
            // In a full implementation, we'd extract the fields from the base declaration
            Err(ResolveError::ValidationError {
                message: format!(
                    "Override resolution not yet fully implemented for base '{}'",
                    override_spec.base.join("::")
                ),
                help: Some("Template overrides are not yet supported. This feature is planned for a future release. For now, define characters directly without using template inheritance.".to_string()),
            })
        },

        | Value::List(items) => {
            // Recursively resolve overrides in list items
            let resolved: Result<Vec<_>> = items
                .iter()
                .map(|v| resolve_value_overrides(v, name_table))
                .collect();
            Ok(Value::List(resolved?))
        },

        | Value::Object(fields) => {
            // Recursively resolve overrides in object fields
            let resolved_fields: Result<Vec<_>> = fields
                .iter()
                .map(|f| {
                    let resolved_value = resolve_value_overrides(&f.value, name_table)?;
                    Ok(Field {
                        name: f.name.clone(),
                        value: resolved_value,
                        span: f.span.clone(),
                    })
                })
                .collect();
            Ok(Value::Object(resolved_fields?))
        },

        // Other value types don't contain overrides
        | _ => Ok(value.clone()),
    }
}

/// Check if applying the same override twice gives the same result
/// (idempotence)
pub fn is_idempotent(base: &[Field], ops: &[OverrideOp]) -> bool {
    let result1 = apply_override(base.to_vec(), ops);
    if result1.is_err() {
        return false;
    }

    let intermediate = result1.unwrap();
    let result2 = apply_override(intermediate.clone(), ops);
    if result2.is_err() {
        return false;
    }

    // Should get the same result
    intermediate == result2.unwrap()
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::syntax::ast::Span;

    fn make_field(name: &str, value: i64) -> Field {
        Field {
            name: name.to_string(),
            value: Value::Int(value),
            span: Span::new(0, 10),
        }
    }

    #[test]
    fn test_set_replaces_existing_field() {
        let base = vec![make_field("age", 25), make_field("health", 100)];

        let ops = vec![OverrideOp::Set(make_field("age", 30))];

        let result = apply_override(base, &ops).unwrap();
        assert_eq!(result.len(), 2);

        let age_field = result.iter().find(|f| f.name == "age").unwrap();
        assert_eq!(age_field.value, Value::Int(30));
    }

    #[test]
    fn test_set_adds_new_field() {
        let base = vec![make_field("age", 25)];

        let ops = vec![OverrideOp::Set(make_field("health", 100))];

        let result = apply_override(base, &ops).unwrap();
        assert_eq!(result.len(), 2);
        assert!(result.iter().any(|f| f.name == "health"));
    }

    #[test]
    fn test_remove_deletes_field() {
        let base = vec![
            make_field("age", 25),
            make_field("health", 100),
            make_field("energy", 50),
        ];

        let ops = vec![OverrideOp::Remove("health".to_string())];

        let result = apply_override(base, &ops).unwrap();
        assert_eq!(result.len(), 2);
        assert!(!result.iter().any(|f| f.name == "health"));
        assert!(result.iter().any(|f| f.name == "age"));
        assert!(result.iter().any(|f| f.name == "energy"));
    }

    #[test]
    fn test_remove_nonexistent_field_is_noop() {
        let base = vec![make_field("age", 25)];

        let ops = vec![OverrideOp::Remove("nonexistent".to_string())];

        let result = apply_override(base, &ops).unwrap();
        assert_eq!(result.len(), 1);
        assert_eq!(result[0].name, "age");
    }

    #[test]
    fn test_append_adds_new_field() {
        let base = vec![make_field("age", 25)];

        let ops = vec![OverrideOp::Append(make_field("health", 100))];

        let result = apply_override(base, &ops).unwrap();
        assert_eq!(result.len(), 2);
        assert!(result.iter().any(|f| f.name == "health"));
    }

    #[test]
    fn test_append_existing_field_errors() {
        let base = vec![make_field("age", 25)];

        let ops = vec![OverrideOp::Append(make_field("age", 30))];

        let result = apply_override(base, &ops);
        assert!(result.is_err());
    }

    #[test]
    fn test_multiple_operations() {
        let base = vec![
            make_field("age", 25),
            make_field("health", 100),
            make_field("energy", 50),
        ];

        let ops = vec![
            OverrideOp::Set(make_field("age", 30)),
            OverrideOp::Remove("energy".to_string()),
            OverrideOp::Append(make_field("strength", 75)),
        ];

        let result = apply_override(base, &ops).unwrap();
        assert_eq!(result.len(), 3);

        let age = result.iter().find(|f| f.name == "age").unwrap();
        assert_eq!(age.value, Value::Int(30));

        assert!(!result.iter().any(|f| f.name == "energy"));
        assert!(result.iter().any(|f| f.name == "strength"));
    }

    #[test]
    fn test_set_is_idempotent() {
        let base = vec![make_field("age", 25)];
        let ops = vec![OverrideOp::Set(make_field("age", 30))];

        assert!(is_idempotent(&base, &ops));
    }

    #[test]
    fn test_remove_is_idempotent() {
        let base = vec![make_field("age", 25), make_field("health", 100)];
        let ops = vec![OverrideOp::Remove("health".to_string())];

        assert!(is_idempotent(&base, &ops));
    }

    #[test]
    fn test_append_is_not_idempotent() {
        let base = vec![make_field("age", 25)];
        let ops = vec![OverrideOp::Append(make_field("health", 100))];

        // Append is NOT idempotent because second application would try to
        // append to a list that already has the field
        assert!(!is_idempotent(&base, &ops));
    }

    // ===== Template Composition Tests =====

    use crate::syntax::ast::File;

    fn make_file(declarations: Vec<Declaration>) -> File {
        File { declarations }
    }

    fn make_template(
        name: &str,
        fields: Vec<Field>,
        includes: Vec<&str>,
        strict: bool,
    ) -> Template {
        Template {
            name: name.to_string(),
            fields,
            includes: includes.iter().map(|s| s.to_string()).collect(),
            strict,
            span: Span::new(0, 10),
        }
    }

    fn make_character(name: &str, fields: Vec<Field>, templates: Vec<&str>) -> Character {
        Character {
            name: name.to_string(),
            fields,
            template: if templates.is_empty() {
                None
            } else {
                Some(templates.iter().map(|s| s.to_string()).collect())
            },
            span: Span::new(0, 10),
        }
    }

    #[test]
    fn test_resolve_template_with_no_includes() {
        let template = make_template("Person", vec![make_field("age", 25)], vec![], false);
        let declarations = vec![Declaration::Template(template.clone())];
        let name_table = NameTable::from_file(&make_file(declarations.clone())).unwrap();
        let mut visited = HashSet::new();

        let result =
            resolve_template_includes(&template, &declarations, &name_table, &mut visited).unwrap();

        assert_eq!(result.len(), 1);
        assert_eq!(result[0].name, "age");
        assert_eq!(result[0].value, Value::Int(25));
    }

    #[test]
    fn test_resolve_template_with_single_include() {
        let base = make_template("Human", vec![make_field("age", 0)], vec![], false);
        let derived = make_template("Person", vec![make_field("name", 0)], vec!["Human"], false);

        let declarations = vec![
            Declaration::Template(base),
            Declaration::Template(derived.clone()),
        ];
        let name_table = NameTable::from_file(&make_file(declarations.clone())).unwrap();
        let mut visited = HashSet::new();

        let result =
            resolve_template_includes(&derived, &declarations, &name_table, &mut visited).unwrap();

        assert_eq!(result.len(), 2);
        assert!(result.iter().any(|f| f.name == "age"));
        assert!(result.iter().any(|f| f.name == "name"));
    }

    #[test]
    fn test_resolve_template_with_chained_includes() {
        let base = make_template("Being", vec![make_field("alive", 1)], vec![], false);
        let middle = make_template("Human", vec![make_field("age", 0)], vec!["Being"], false);
        let top = make_template("Person", vec![make_field("name", 0)], vec!["Human"], false);

        let declarations = vec![
            Declaration::Template(base),
            Declaration::Template(middle),
            Declaration::Template(top.clone()),
        ];
        let name_table = NameTable::from_file(&make_file(declarations.clone())).unwrap();
        let mut visited = HashSet::new();

        let result =
            resolve_template_includes(&top, &declarations, &name_table, &mut visited).unwrap();

        assert_eq!(result.len(), 3);
        assert!(result.iter().any(|f| f.name == "alive"));
        assert!(result.iter().any(|f| f.name == "age"));
        assert!(result.iter().any(|f| f.name == "name"));
    }

    #[test]
    fn test_resolve_template_field_override() {
        let base = make_template("Human", vec![make_field("age", 0)], vec![], false);
        let derived = make_template(
            "Person",
            vec![make_field("age", 25)], // Override with concrete value
            vec!["Human"],
            false,
        );

        let declarations = vec![
            Declaration::Template(base),
            Declaration::Template(derived.clone()),
        ];
        let name_table = NameTable::from_file(&make_file(declarations.clone())).unwrap();
        let mut visited = HashSet::new();

        let result =
            resolve_template_includes(&derived, &declarations, &name_table, &mut visited).unwrap();

        assert_eq!(result.len(), 1);
        assert_eq!(result[0].name, "age");
        assert_eq!(result[0].value, Value::Int(25)); // Should be overridden
                                                     // value
    }

    #[test]
    fn test_merge_character_templates_single() {
        let template = make_template("Person", vec![make_field("age", 0)], vec![], false);
        let character = make_character("Martha", vec![make_field("age", 34)], vec!["Person"]);

        let declarations = vec![
            Declaration::Template(template),
            Declaration::Character(character.clone()),
        ];
        let name_table = NameTable::from_file(&make_file(declarations.clone())).unwrap();

        let result = merge_character_templates(&character, &declarations, &name_table).unwrap();

        assert_eq!(result.len(), 1);
        assert_eq!(result[0].name, "age");
        assert_eq!(result[0].value, Value::Int(34)); // Character's value
                                                     // overrides template
    }

    #[test]
    fn test_merge_character_templates_multiple() {
        let physical = make_template("Physical", vec![make_field("height", 0)], vec![], false);
        let mental = make_template("Mental", vec![make_field("iq", 0)], vec![], false);
        let character = make_character(
            "Martha",
            vec![make_field("height", 165), make_field("iq", 120)],
            vec!["Physical", "Mental"],
        );

        let declarations = vec![
            Declaration::Template(physical),
            Declaration::Template(mental),
            Declaration::Character(character.clone()),
        ];
        let name_table = NameTable::from_file(&make_file(declarations.clone())).unwrap();

        let result = merge_character_templates(&character, &declarations, &name_table).unwrap();

        assert_eq!(result.len(), 2);
        assert!(result
            .iter()
            .any(|f| f.name == "height" && f.value == Value::Int(165)));
        assert!(result
            .iter()
            .any(|f| f.name == "iq" && f.value == Value::Int(120)));
    }

    #[test]
    fn test_merge_character_templates_with_includes() {
        let base = make_template("Human", vec![make_field("age", 0)], vec![], false);
        let derived = make_template("Person", vec![make_field("name", 0)], vec!["Human"], false);
        let character = make_character(
            "Martha",
            vec![make_field("age", 34), make_field("name", 1)],
            vec!["Person"],
        );

        let declarations = vec![
            Declaration::Template(base),
            Declaration::Template(derived),
            Declaration::Character(character.clone()),
        ];
        let name_table = NameTable::from_file(&make_file(declarations.clone())).unwrap();

        let result = merge_character_templates(&character, &declarations, &name_table).unwrap();

        assert_eq!(result.len(), 2);
        assert!(result
            .iter()
            .any(|f| f.name == "age" && f.value == Value::Int(34)));
        assert!(result
            .iter()
            .any(|f| f.name == "name" && f.value == Value::Int(1)));
    }

    #[test]
    fn test_strict_template_validation_passes() {
        let template = make_template("Person", vec![make_field("age", 0)], vec![], true);
        let character = make_character("Martha", vec![make_field("age", 34)], vec!["Person"]);

        let declarations = vec![
            Declaration::Template(template),
            Declaration::Character(character.clone()),
        ];
        let name_table = NameTable::from_file(&make_file(declarations.clone())).unwrap();

        let result = merge_character_templates(&character, &declarations, &name_table);
        assert!(result.is_ok());
    }

    #[test]
    fn test_strict_template_validation_fails_with_range() {
        let template = make_template(
            "Person",
            vec![Field {
                name: "age".to_string(),
                value: Value::Range(Box::new(Value::Int(18)), Box::new(Value::Int(65))),
                span: Span::new(0, 10),
            }],
            vec![],
            true,
        );
        let character = make_character("Martha", vec![], vec!["Person"]);

        let declarations = vec![
            Declaration::Template(template),
            Declaration::Character(character.clone()),
        ];
        let name_table = NameTable::from_file(&make_file(declarations.clone())).unwrap();

        let result = merge_character_templates(&character, &declarations, &name_table);
        assert!(result.is_err());
        if let Err(ResolveError::ValidationError { message, .. }) = result {
            assert!(message.contains("strict template"));
            assert!(message.contains("range value"));
        }
    }

    #[test]
    fn test_circular_include_detection() {
        let a = make_template("A", vec![], vec!["B"], false);
        let b = make_template("B", vec![], vec!["A"], false);

        let declarations = vec![Declaration::Template(a.clone()), Declaration::Template(b)];
        let name_table = NameTable::from_file(&make_file(declarations.clone())).unwrap();
        let mut visited = HashSet::new();

        let result = resolve_template_includes(&a, &declarations, &name_table, &mut visited);
        assert!(result.is_err());
        if let Err(ResolveError::CircularDependency { .. }) = result {
            // Expected
        } else {
            panic!("Expected CircularDependency error");
        }
    }

    #[test]
    fn test_merge_field_lists_override() {
        let base = vec![make_field("age", 25), make_field("health", 100)];
        let overrides = vec![make_field("age", 30)];

        let result = merge_field_lists(base, overrides);

        assert_eq!(result.len(), 2);
        let age = result.iter().find(|f| f.name == "age").unwrap();
        assert_eq!(age.value, Value::Int(30));
    }

    #[test]
    fn test_merge_field_lists_add_new() {
        let base = vec![make_field("age", 25)];
        let overrides = vec![make_field("health", 100)];

        let result = merge_field_lists(base, overrides);

        assert_eq!(result.len(), 2);
        assert!(result.iter().any(|f| f.name == "age"));
        assert!(result.iter().any(|f| f.name == "health"));
    }
}