storybook/src/resolve/references.rs

//! Semantic reference tracking and resolution
//!
//! This module provides semantic analysis to find all references to symbols,
//! enabling features like rename refactoring and find-all-references.

use super::names::{
    DeclKind,
    NameTable,
};
use crate::syntax::ast::{
    Behavior,
    Character,
    Declaration,
    Field,
    File,
    Institution,
    LifeArc,
    Location,
    Participant,
    Relationship,
    Schedule,
    Span,
    Species,
    Template,
    Value,
};

/// A reference to a symbol in the code
#[derive(Debug, Clone, PartialEq)]
pub struct Reference {
    /// The name being referenced
    pub name: String,
    /// Kind of symbol being referenced
    pub kind: DeclKind,
    /// Location of the reference (just the identifier)
    pub span: Span,
    /// Index of the file containing this reference
    pub file_index: usize,
    /// Context of the reference
    pub context: ReferenceContext,
}

/// Context describing where and how a symbol is referenced
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ReferenceContext {
    /// Symbol definition/declaration
    Definition,
    /// Used as a type annotation (e.g., `character Alice: Person`)
    TypeAnnotation,
    /// Used in a field value (e.g., `friend: Alice`)
    FieldValue,
    /// Referenced in a behavior tree (e.g., `@WorkAtBakery`)
    BehaviorReference,
    /// Used in a template include
    TemplateInclude,
    /// Used in a relationship participant
    RelationshipParticipant,
    /// Other/unknown context
    Other,
}

/// Find all references to a specific symbol across all files
pub fn find_all_references(
    files: &[File],
    symbol_name: &str,
    symbol_kind: DeclKind,
) -> Vec<Reference> {
    let mut references = Vec::new();

    for (file_index, file) in files.iter().enumerate() {
        // Build name table to validate symbols exist
        let name_table = match NameTable::from_file(file) {
            | Ok(table) => table,
            | Err(_) => continue, // Skip files with errors
        };

        // Find definition if it exists in this file
        if let Some(entry) = name_table.resolve_name(symbol_name) {
            if entry.kind == symbol_kind {
                references.push(Reference {
                    name: symbol_name.to_string(),
                    kind: symbol_kind,
                    span: entry.span.clone(),
                    file_index,
                    context: ReferenceContext::Definition,
                });
            }
        }

        // Walk AST to find all semantic references
        for decl in &file.declarations {
            references.extend(find_references_in_declaration(
                decl,
                symbol_name,
                symbol_kind,
                file_index,
            ));
        }
    }

    references
}

/// Find references within a single declaration
fn find_references_in_declaration(
    decl: &Declaration,
    symbol_name: &str,
    symbol_kind: DeclKind,
    file_index: usize,
) -> Vec<Reference> {
    let mut refs = Vec::new();

    match decl {
        | Declaration::Character(c) => {
            refs.extend(find_references_in_character(
                c,
                symbol_name,
                symbol_kind,
                file_index,
            ));
        },
        | Declaration::Template(t) => {
            refs.extend(find_references_in_template(
                t,
                symbol_name,
                symbol_kind,
                file_index,
            ));
        },
        | Declaration::LifeArc(l) => {
            refs.extend(find_references_in_life_arc(
                l,
                symbol_name,
                symbol_kind,
                file_index,
            ));
        },
        | Declaration::Schedule(s) => {
            refs.extend(find_references_in_schedule(
                s,
                symbol_name,
                symbol_kind,
                file_index,
            ));
        },
        | Declaration::Behavior(b) => {
            refs.extend(find_references_in_behavior(
                b,
                symbol_name,
                symbol_kind,
                file_index,
            ));
        },
        | Declaration::Institution(i) => {
            refs.extend(find_references_in_institution(
                i,
                symbol_name,
                symbol_kind,
                file_index,
            ));
        },
        | Declaration::Relationship(r) => {
            refs.extend(find_references_in_relationship(
                r,
                symbol_name,
                symbol_kind,
                file_index,
            ));
        },
        | Declaration::Location(l) => {
            refs.extend(find_references_in_location(
                l,
                symbol_name,
                symbol_kind,
                file_index,
            ));
        },
        | Declaration::Species(s) => {
            refs.extend(find_references_in_species(
                s,
                symbol_name,
                symbol_kind,
                file_index,
            ));
        },
        | Declaration::Use(_) => {
            // Use statements are handled separately
        },
        | Declaration::Concept(_) |
        Declaration::SubConcept(_) |
        Declaration::ConceptComparison(_) => {
            // TODO: Implement reference finding for type system
        },
    }

    refs
}

fn find_references_in_character(
    c: &Character,
    symbol_name: &str,
    symbol_kind: DeclKind,
    file_index: usize,
) -> Vec<Reference> {
    let mut refs = Vec::new();

    // Check species annotation
    if let Some(ref species) = c.species {
        if species == symbol_name && symbol_kind == DeclKind::Species {
            refs.push(Reference {
                name: symbol_name.to_string(),
                kind: symbol_kind,
                span: c.span.clone(),
                file_index,
                context: ReferenceContext::TypeAnnotation,
            });
        }
    }

    // Check templates (character can have multiple)
    if let Some(ref templates) = c.template {
        for template in templates {
            if template == symbol_name && symbol_kind == DeclKind::Template {
                refs.push(Reference {
                    name: symbol_name.to_string(),
                    kind: symbol_kind,
                    span: c.span.clone(),
                    file_index,
                    context: ReferenceContext::TypeAnnotation,
                });
            }
        }
    }

    // Check fields for identifier references
    refs.extend(find_references_in_fields(
        &c.fields,
        symbol_name,
        symbol_kind,
        file_index,
    ));

    refs
}

fn find_references_in_template(
    t: &Template,
    symbol_name: &str,
    symbol_kind: DeclKind,
    file_index: usize,
) -> Vec<Reference> {
    let mut refs = Vec::new();

    // Check includes
    for include in &t.includes {
        if include == symbol_name && symbol_kind == DeclKind::Template {
            refs.push(Reference {
                name: symbol_name.to_string(),
                kind: symbol_kind,
                span: t.span.clone(),
                file_index,
                context: ReferenceContext::TemplateInclude,
            });
        }
    }

    // Check fields
    refs.extend(find_references_in_fields(
        &t.fields,
        symbol_name,
        symbol_kind,
        file_index,
    ));

    refs
}

fn find_references_in_fields(
    fields: &[Field],
    symbol_name: &str,
    symbol_kind: DeclKind,
    file_index: usize,
) -> Vec<Reference> {
    let mut refs = Vec::new();

    for field in fields {
        // Check if field value is an identifier that references the symbol
        refs.extend(find_references_in_value(
            &field.value,
            symbol_name,
            symbol_kind,
            file_index,
            field.span.clone(),
        ));
    }

    refs
}

fn find_references_in_value(
    value: &Value,
    symbol_name: &str,
    symbol_kind: DeclKind,
    file_index: usize,
    span: Span,
) -> Vec<Reference> {
    let mut refs = Vec::new();

    match value {
        | Value::Identifier(path) => {
            // Check if this identifier references our symbol
            if let Some(name) = path.last() {
                if name == symbol_name {
                    // Identifiers can reference characters, templates, enums, species
                    let matches_kind = matches!(
                        symbol_kind,
                        DeclKind::Character | DeclKind::Template | DeclKind::Species
                    );

                    if matches_kind {
                        refs.push(Reference {
                            name: symbol_name.to_string(),
                            kind: symbol_kind,
                            span,
                            file_index,
                            context: ReferenceContext::FieldValue,
                        });
                    }
                }
            }
        },
        | Value::List(values) => {
            // Recursively check list values
            for v in values {
                refs.extend(find_references_in_value(
                    v,
                    symbol_name,
                    symbol_kind,
                    file_index,
                    span.clone(),
                ));
            }
        },
        | Value::Object(fields) => {
            // Recursively check object fields
            refs.extend(find_references_in_fields(
                fields,
                symbol_name,
                symbol_kind,
                file_index,
            ));
        },
        | Value::Override(override_val) => {
            // Check the base template reference
            if let Some(base_name) = override_val.base.last() {
                if base_name == symbol_name && symbol_kind == DeclKind::Template {
                    refs.push(Reference {
                        name: symbol_name.to_string(),
                        kind: symbol_kind,
                        span: override_val.span.clone(),
                        file_index,
                        context: ReferenceContext::FieldValue,
                    });
                }
            }
        },
        | _ => {
            // Other value types don't contain references
        },
    }

    refs
}

fn find_references_in_life_arc(
    _l: &LifeArc,
    _symbol_name: &str,
    _symbol_kind: DeclKind,
    _file_index: usize,
) -> Vec<Reference> {
    // Life arcs don't typically reference other symbols
    Vec::new()
}

fn find_references_in_schedule(
    _s: &Schedule,
    _symbol_name: &str,
    _symbol_kind: DeclKind,
    _file_index: usize,
) -> Vec<Reference> {
    // Schedules don't typically reference other symbols
    Vec::new()
}

fn find_references_in_behavior(
    _b: &Behavior,
    _symbol_name: &str,
    _symbol_kind: DeclKind,
    _file_index: usize,
) -> Vec<Reference> {
    // TODO: Parse behavior tree nodes to find @BehaviorName references
    // This requires walking the BehaviorNode tree
    Vec::new()
}

fn find_references_in_institution(
    i: &Institution,
    symbol_name: &str,
    symbol_kind: DeclKind,
    file_index: usize,
) -> Vec<Reference> {
    find_references_in_fields(&i.fields, symbol_name, symbol_kind, file_index)
}

fn find_references_in_relationship(
    r: &Relationship,
    symbol_name: &str,
    symbol_kind: DeclKind,
    file_index: usize,
) -> Vec<Reference> {
    let mut refs = Vec::new();

    // Check participant references
    for participant in &r.participants {
        refs.extend(find_references_in_participant(
            participant,
            symbol_name,
            symbol_kind,
            file_index,
        ));
    }

    refs
}

fn find_references_in_participant(
    p: &Participant,
    symbol_name: &str,
    symbol_kind: DeclKind,
    file_index: usize,
) -> Vec<Reference> {
    let mut refs = Vec::new();

    // Check if participant name references the symbol
    if let Some(participant_name) = p.name.last() {
        if participant_name == symbol_name && symbol_kind == DeclKind::Character {
            refs.push(Reference {
                name: symbol_name.to_string(),
                kind: symbol_kind,
                span: p.span.clone(),
                file_index,
                context: ReferenceContext::RelationshipParticipant,
            });
        }
    }

    // Check participant fields
    refs.extend(find_references_in_fields(
        &p.fields,
        symbol_name,
        symbol_kind,
        file_index,
    ));

    refs
}

fn find_references_in_location(
    l: &Location,
    symbol_name: &str,
    symbol_kind: DeclKind,
    file_index: usize,
) -> Vec<Reference> {
    find_references_in_fields(&l.fields, symbol_name, symbol_kind, file_index)
}

fn find_references_in_species(
    s: &Species,
    symbol_name: &str,
    symbol_kind: DeclKind,
    file_index: usize,
) -> Vec<Reference> {
    find_references_in_fields(&s.fields, symbol_name, symbol_kind, file_index)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::syntax::{
        lexer::Lexer,
        FileParser,
    };

    fn parse(source: &str) -> File {
        let lexer = Lexer::new(source);
        FileParser::new().parse(lexer).unwrap()
    }

    #[test]
    fn test_find_character_references_in_field() {
        let source = r#"
character Alice {}
character Bob { friend: Alice }
"#;
        let file = parse(source);
        let files = vec![file];

        let refs = find_all_references(&files, "Alice", DeclKind::Character);

        // Should find: definition + field reference
        assert_eq!(refs.len(), 2);

        let definition = refs
            .iter()
            .find(|r| r.context == ReferenceContext::Definition);
        assert!(definition.is_some());

        let field_ref = refs
            .iter()
            .find(|r| r.context == ReferenceContext::FieldValue);
        assert!(field_ref.is_some());
    }

    #[test]
    fn test_find_template_references() {
        let source = r#"
template Person {}
character Alice from Person {}
"#;
        let file = parse(source);
        let files = vec![file];

        let refs = find_all_references(&files, "Person", DeclKind::Template);

        // Should find: definition + type annotation
        assert_eq!(refs.len(), 2);

        let type_ref = refs
            .iter()
            .find(|r| r.context == ReferenceContext::TypeAnnotation);
        assert!(type_ref.is_some());
    }

    #[test]
    fn test_find_species_references() {
        let source = r#"
species Human {}
character Alice: Human {}
"#;
        let file = parse(source);
        let files = vec![file];

        let refs = find_all_references(&files, "Human", DeclKind::Species);

        // Should find: definition + species annotation
        assert_eq!(refs.len(), 2);
    }

    #[test]
    fn test_find_references_across_multiple_files() {
        let file1 = parse("character Alice {}");
        let file2 = parse("character Bob { friend: Alice }");
        let file3 = parse("character Charlie { mentor: Alice }");

        let files = vec![file1, file2, file3];

        let refs = find_all_references(&files, "Alice", DeclKind::Character);

        // Should find: 1 definition + 2 references
        assert_eq!(refs.len(), 3);

        let def = refs
            .iter()
            .filter(|r| r.context == ReferenceContext::Definition)
            .count();
        assert_eq!(def, 1);

        let field_refs = refs
            .iter()
            .filter(|r| r.context == ReferenceContext::FieldValue)
            .count();
        assert_eq!(field_refs, 2);

        // Check file indices
        assert_eq!(refs.iter().filter(|r| r.file_index == 0).count(), 1);
        assert_eq!(refs.iter().filter(|r| r.file_index == 1).count(), 1);
        assert_eq!(refs.iter().filter(|r| r.file_index == 2).count(), 1);
    }

    #[test]
    fn test_respects_symbol_kind() {
        let source = r#"
character Alice {}
template Person {}
character Bob { friend: Alice }
character Charlie from Person {}
"#;
        let file = parse(source);
        let files = vec![file];

        // Find character Alice
        let char_refs = find_all_references(&files, "Alice", DeclKind::Character);
        // Should find: character definition + field reference
        assert_eq!(char_refs.len(), 2);

        // Find template Person
        let template_refs = find_all_references(&files, "Person", DeclKind::Template);
        // Should find: template definition + from reference
        assert_eq!(template_refs.len(), 2);
    }

    #[test]
    fn test_template_includes() {
        let source = r#"
template Base {}
template Extended { include Base }
"#;
        let file = parse(source);
        let files = vec![file];

        let refs = find_all_references(&files, "Base", DeclKind::Template);

        // Should find: definition + include reference
        assert_eq!(refs.len(), 2);

        let include_ref = refs
            .iter()
            .find(|r| r.context == ReferenceContext::TemplateInclude);
        assert!(include_ref.is_some());
    }

    #[test]
    fn test_relationship_participants() {
        let source = r#"
character Alice {}
character Bob {}
relationship Friends { Alice as friend {} Bob as friend {} }
"#;
        let file = parse(source);
        let files = vec![file];

        let alice_refs = find_all_references(&files, "Alice", DeclKind::Character);
        let bob_refs = find_all_references(&files, "Bob", DeclKind::Character);

        // Each should have: definition + relationship participant
        assert_eq!(alice_refs.len(), 2);
        assert_eq!(bob_refs.len(), 2);
    }

    #[test]
    fn test_no_references_found() {
        let source = "character Alice {}";
        let file = parse(source);
        let files = vec![file];

        // Look for non-existent symbol
        let refs = find_all_references(&files, "Bob", DeclKind::Character);

        // Should find nothing
        assert_eq!(refs.len(), 0);
    }
}