storybook/src/lsp/semantic_tokens.rs

//! Semantic tokens for enhanced syntax highlighting
//!
//! Provides detailed token type information beyond what a basic grammar can
//! provide, allowing the editor to highlight different kinds of identifiers,
//! keywords, and values with appropriate semantic meaning.

use tower_lsp::lsp_types::{
    SemanticToken,
    SemanticTokenType,
    SemanticTokens,
    SemanticTokensResult,
};

use super::document::Document;
use crate::syntax::{
    ast::{
        Declaration,
        Field,
        SubConceptKind,
        Value,
    },
    lexer::{
        Lexer,
        Token,
    },
};

/// Standard semantic token types supported by LSP
pub const LEGEND_TYPES: &[SemanticTokenType] = &[
    SemanticTokenType::NAMESPACE,   // use paths
    SemanticTokenType::TYPE,        // template names, species names, enum names
    SemanticTokenType::CLASS,       // character declarations
    SemanticTokenType::ENUM,        // enum declarations
    SemanticTokenType::INTERFACE,   // template declarations
    SemanticTokenType::STRUCT,      // institution, location declarations
    SemanticTokenType::PARAMETER,   // action parameters
    SemanticTokenType::VARIABLE,    // character names in schedules
    SemanticTokenType::PROPERTY,    // field names
    SemanticTokenType::ENUM_MEMBER, // enum variant names
    SemanticTokenType::FUNCTION,    // behavior names
    SemanticTokenType::METHOD,      // relationship names
    SemanticTokenType::KEYWORD,     // keywords like "from", "include", "strict"
    SemanticTokenType::STRING,      // string literals
    SemanticTokenType::NUMBER,      // numeric literals
    SemanticTokenType::OPERATOR,    // operators like "..", "->", etc.
];

/// Semantic token modifiers (currently unused but available for future use)
pub const LEGEND_MODIFIERS: &[&str] = &[
    "declaration",
    "definition",
    "readonly",
    "static",
    "deprecated",
    "abstract",
    "async",
    "modification",
    "documentation",
    "defaultLibrary",
];

/// Helper to find identifier positions within a span using the lexer
fn find_identifiers_in_span(
    text: &str,
    span_start: usize,
    span_end: usize,
    target_names: &[String],
) -> Vec<(usize, String)> {
    let span_text = &text[span_start..span_end];
    let lexer = Lexer::new(span_text);
    let tokens: Vec<_> = lexer.collect();

    let mut results = Vec::new();
    for (offset, token, _end) in tokens {
        if let Token::Ident(name) = token {
            if target_names.contains(&name) {
                results.push((span_start + offset, name));
            }
        }
    }
    results
}

/// Recursively highlight behavior tree nodes
fn highlight_behavior_node(
    builder: &mut SemanticTokensBuilder,
    node: &crate::syntax::ast::BehaviorNode,
) {
    use crate::syntax::ast::BehaviorNode;

    match node {
        | BehaviorNode::Selector { children, .. } | BehaviorNode::Sequence { children, .. } => {
            for child in children {
                highlight_behavior_node(builder, child);
            }
        },
        | BehaviorNode::Action(_action_name, params) => {
            // Action names don't have spans, so we'd need to search for them
            // For now, just highlight the parameters
            for param in params {
                highlight_field(builder, param);
            }
        },
        | BehaviorNode::Decorator { child, .. } => {
            highlight_behavior_node(builder, child);
        },
        | BehaviorNode::SubTree(_path) => {
            // SubTree references another behavior by path
            // Would need position tracking to highlight
        },
        | BehaviorNode::Condition(_expr) => {
            // Conditions contain expressions which could be highlighted
            // Would need expression traversal
        },
    }
}

/// Generate semantic tokens for a document
pub fn get_semantic_tokens(doc: &Document) -> Option<SemanticTokensResult> {
    let ast = doc.ast.as_ref()?;
    let mut builder = SemanticTokensBuilder::new(&doc.text);
    let mut positions = doc.positions.clone();

    // Process all top-level declarations
    for decl in &ast.declarations {
        match decl {
            | Declaration::Use(use_decl) => {
                // Highlight use paths as namespaces
                let path_positions = find_identifiers_in_span(
                    &doc.text,
                    use_decl.span.start,
                    use_decl.span.end,
                    &use_decl.path,
                );

                for (offset, segment) in path_positions {
                    let (line, col) = positions.offset_to_position(offset);
                    builder.add_token(
                        line,
                        col,
                        segment.len(),
                        token_type_index(SemanticTokenType::NAMESPACE),
                        0,
                    );
                }
            },
            | Declaration::Character(character) => {
                // Highlight character name as CLASS
                builder.add_token(
                    character.span.start_line,
                    character.span.start_col,
                    character.name.len(),
                    token_type_index(SemanticTokenType::CLASS),
                    0,
                );

                // Highlight species as TYPE
                if let Some(ref species) = character.species {
                    let species_positions = find_identifiers_in_span(
                        &doc.text,
                        character.span.start,
                        character.span.end,
                        std::slice::from_ref(species),
                    );

                    for (offset, species_name) in species_positions {
                        let (line, col) = positions.offset_to_position(offset);
                        builder.add_token(
                            line,
                            col,
                            species_name.len(),
                            token_type_index(SemanticTokenType::TYPE),
                            0,
                        );
                    }
                }

                // Highlight template references
                if let Some(ref templates) = character.template {
                    let template_positions = find_identifiers_in_span(
                        &doc.text,
                        character.span.start,
                        character.span.end,
                        templates,
                    );

                    for (offset, template_name) in template_positions {
                        let (line, col) = positions.offset_to_position(offset);
                        builder.add_token(
                            line,
                            col,
                            template_name.len(),
                            token_type_index(SemanticTokenType::INTERFACE),
                            0,
                        );
                    }
                }

                // Highlight fields
                for field in &character.fields {
                    highlight_field(&mut builder, field);
                }
            },
            | Declaration::Template(template) => {
                // Highlight template name as INTERFACE
                builder.add_token(
                    template.span.start_line,
                    template.span.start_col,
                    template.name.len(),
                    token_type_index(SemanticTokenType::INTERFACE),
                    0,
                );

                // Find and highlight includes using the lexer
                let include_positions = find_identifiers_in_span(
                    &doc.text,
                    template.span.start,
                    template.span.end,
                    &template.includes,
                );

                for (offset, include_name) in include_positions {
                    let (line, col) = positions.offset_to_position(offset);
                    builder.add_token(
                        line,
                        col,
                        include_name.len(),
                        token_type_index(SemanticTokenType::INTERFACE),
                        0,
                    );
                }

                // Highlight fields
                for field in &template.fields {
                    highlight_field(&mut builder, field);
                }
            },
            | Declaration::Species(species) => {
                // Highlight species name as TYPE
                builder.add_token(
                    species.span.start_line,
                    species.span.start_col,
                    species.name.len(),
                    token_type_index(SemanticTokenType::TYPE),
                    0,
                );

                // Highlight fields
                for field in &species.fields {
                    highlight_field(&mut builder, field);
                }
            },
            | Declaration::Institution(institution) => {
                // Highlight institution name as STRUCT
                builder.add_token(
                    institution.span.start_line,
                    institution.span.start_col,
                    institution.name.len(),
                    token_type_index(SemanticTokenType::STRUCT),
                    0,
                );

                // Highlight fields
                for field in &institution.fields {
                    highlight_field(&mut builder, field);
                }
            },
            | Declaration::Location(location) => {
                // Highlight location name as STRUCT
                builder.add_token(
                    location.span.start_line,
                    location.span.start_col,
                    location.name.len(),
                    token_type_index(SemanticTokenType::STRUCT),
                    0,
                );

                // Highlight fields
                for field in &location.fields {
                    highlight_field(&mut builder, field);
                }
            },
            | Declaration::Behavior(behavior) => {
                // Highlight behavior name as FUNCTION
                builder.add_token(
                    behavior.span.start_line,
                    behavior.span.start_col,
                    behavior.name.len(),
                    token_type_index(SemanticTokenType::FUNCTION),
                    0,
                );

                // TODO: Traverse behavior tree to highlight conditions and actions
                // Would need recursive function to walk BehaviorNode tree
                highlight_behavior_node(&mut builder, &behavior.root);
            },
            | Declaration::Relationship(relationship) => {
                // Highlight relationship name as METHOD
                builder.add_token(
                    relationship.span.start_line,
                    relationship.span.start_col,
                    relationship.name.len(),
                    token_type_index(SemanticTokenType::METHOD),
                    0,
                );

                // Highlight participants
                for participant in &relationship.participants {
                    // For qualified paths like "Alice.parent", we want to highlight each segment
                    // The participant has its own span, so we can search within it
                    let participant_names = participant.name.clone();
                    let name_positions = find_identifiers_in_span(
                        &doc.text,
                        participant.span.start,
                        participant.span.end,
                        &participant_names,
                    );

                    for (offset, name) in name_positions {
                        let (line, col) = positions.offset_to_position(offset);
                        builder.add_token(
                            line,
                            col,
                            name.len(),
                            token_type_index(SemanticTokenType::VARIABLE),
                            0,
                        );
                    }
                }

                // Highlight fields
                for field in &relationship.fields {
                    highlight_field(&mut builder, field);
                }
            },
            | Declaration::LifeArc(life_arc) => {
                // Highlight life_arc name as TYPE
                builder.add_token(
                    life_arc.span.start_line,
                    life_arc.span.start_col,
                    life_arc.name.len(),
                    token_type_index(SemanticTokenType::TYPE),
                    0,
                );

                // Highlight states and transitions
                for state in &life_arc.states {
                    // State name as ENUM_MEMBER
                    builder.add_token(
                        state.span.start_line,
                        state.span.start_col,
                        state.name.len(),
                        token_type_index(SemanticTokenType::ENUM_MEMBER),
                        0,
                    );

                    // State fields
                    if let Some(ref fields) = state.on_enter {
                        for field in fields {
                            highlight_field(&mut builder, field);
                        }
                    }
                }
            },
            | Declaration::Schedule(schedule) => {
                // Highlight schedule name as TYPE
                builder.add_token(
                    schedule.span.start_line,
                    schedule.span.start_col,
                    schedule.name.len(),
                    token_type_index(SemanticTokenType::TYPE),
                    0,
                );

                // Highlight block fields
                for block in &schedule.blocks {
                    for field in &block.fields {
                        highlight_field(&mut builder, field);
                    }
                }
            },
            | Declaration::Concept(concept) => {
                // Highlight concept name as TYPE
                builder.add_token(
                    concept.span.start_line,
                    concept.span.start_col,
                    concept.name.len(),
                    token_type_index(SemanticTokenType::TYPE),
                    0,
                );
            },
            | Declaration::SubConcept(sub_concept) => {
                // Highlight parent concept as TYPE
                let parent_positions = find_identifiers_in_span(
                    &doc.text,
                    sub_concept.span.start,
                    sub_concept.span.end,
                    std::slice::from_ref(&sub_concept.parent_concept),
                );
                if let Some((offset, parent_name)) = parent_positions.into_iter().next() {
                    let (line, col) = positions.offset_to_position(offset);
                    builder.add_token(
                        line,
                        col,
                        parent_name.len(),
                        token_type_index(SemanticTokenType::TYPE),
                        0,
                    );
                }

                // Highlight sub_concept name as ENUM
                let name_positions = find_identifiers_in_span(
                    &doc.text,
                    sub_concept.span.start,
                    sub_concept.span.end,
                    std::slice::from_ref(&sub_concept.name),
                );
                if let Some((offset, name)) = name_positions.into_iter().next() {
                    let (line, col) = positions.offset_to_position(offset);
                    builder.add_token(
                        line,
                        col,
                        name.len(),
                        token_type_index(SemanticTokenType::ENUM),
                        0,
                    );
                }

                // Highlight enum variants as ENUM_MEMBER, or record fields
                match &sub_concept.kind {
                    | SubConceptKind::Enum { variants } => {
                        let variant_positions = find_identifiers_in_span(
                            &doc.text,
                            sub_concept.span.start,
                            sub_concept.span.end,
                            variants,
                        );
                        for (offset, variant_name) in variant_positions {
                            let (line, col) = positions.offset_to_position(offset);
                            builder.add_token(
                                line,
                                col,
                                variant_name.len(),
                                token_type_index(SemanticTokenType::ENUM_MEMBER),
                                0,
                            );
                        }
                    },
                    | SubConceptKind::Record { fields } => {
                        for field in fields {
                            highlight_field(&mut builder, field);
                        }
                    },
                }
            },
            | Declaration::ConceptComparison(comparison) => {
                // Highlight comparison name as TYPE
                builder.add_token(
                    comparison.span.start_line,
                    comparison.span.start_col,
                    comparison.name.len(),
                    token_type_index(SemanticTokenType::TYPE),
                    0,
                );

                // Highlight variant names as ENUM_MEMBER
                for variant in &comparison.variants {
                    let variant_positions = find_identifiers_in_span(
                        &doc.text,
                        variant.span.start,
                        variant.span.end,
                        std::slice::from_ref(&variant.name),
                    );
                    if let Some((offset, name)) = variant_positions.into_iter().next() {
                        let (line, col) = positions.offset_to_position(offset);
                        builder.add_token(
                            line,
                            col,
                            name.len(),
                            token_type_index(SemanticTokenType::ENUM_MEMBER),
                            0,
                        );
                    }
                }
            },
        }
    }

    Some(SemanticTokensResult::Tokens(SemanticTokens {
        result_id: None,
        data: builder.build(),
    }))
}

/// Helper to highlight a field
fn highlight_field(builder: &mut SemanticTokensBuilder, field: &Field) {
    // Highlight field name as PROPERTY
    builder.add_token(
        field.span.start_line,
        field.span.start_col,
        field.name.len(),
        token_type_index(SemanticTokenType::PROPERTY),
        0,
    );

    // Highlight field value
    highlight_value(builder, &field.value);
}

/// Helper to highlight a value
fn highlight_value(builder: &mut SemanticTokensBuilder, value: &Value) {
    match value {
        | Value::Text(_s) => {
            // Text literals are already highlighted by the grammar
            // but we could add semantic highlighting here if needed
        },
        | Value::Number(_) | Value::Decimal(_) => {
            // Number literals are already highlighted by the grammar
        },
        | Value::Boolean(_) => {
            // Boolean literals are already highlighted by the grammar
        },
        | Value::Identifier(_path) => {
            // Identifiers could be highlighted as TYPE references
            // but we'd need precise position tracking
        },
        | Value::List(items) => {
            for item in items {
                highlight_value(builder, item);
            }
        },
        | Value::Object(fields) => {
            for field in fields {
                highlight_field(builder, field);
            }
        },
        | Value::Range(start, end) => {
            highlight_value(builder, start);
            highlight_value(builder, end);
        },
        | Value::ProseBlock(_) => {
            // Prose blocks are already highlighted by the grammar
        },
        | Value::Override(_) => {
            // Override values need their own handling
        },
        | Value::Time(_) | Value::Duration(_) => {
            // Time/duration literals are already highlighted by the grammar
        },
        | Value::Any => {
            // Any keyword is already highlighted by the grammar
        },
    }
}

/// Get the index of a semantic token type in the legend
fn token_type_index(token_type: SemanticTokenType) -> u32 {
    LEGEND_TYPES
        .iter()
        .position(|t| t == &token_type)
        .unwrap_or(0) as u32
}

/// Builder for semantic tokens with proper delta encoding
struct SemanticTokensBuilder {
    tokens: Vec<(usize, usize, usize, u32, u32)>, // (line, col, length, type, modifiers)
}

impl SemanticTokensBuilder {
    fn new(_text: &str) -> Self {
        Self { tokens: Vec::new() }
    }

    fn add_token(
        &mut self,
        line: usize,
        col: usize,
        length: usize,
        token_type: u32,
        modifiers: u32,
    ) {
        self.tokens.push((line, col, length, token_type, modifiers));
    }

    fn build(mut self) -> Vec<SemanticToken> {
        // Sort tokens by position (line, then column)
        self.tokens
            .sort_by_key(|(line, col, _, _, _)| (*line, *col));

        // Convert to delta-encoded format required by LSP
        let mut result = Vec::new();
        let mut prev_line = 0;
        let mut prev_col = 0;

        for (line, col, length, token_type, modifiers) in self.tokens {
            let delta_line = line - prev_line;
            let delta_start = if delta_line == 0 { col - prev_col } else { col };

            result.push(SemanticToken {
                delta_line: delta_line as u32,
                delta_start: delta_start as u32,
                length: length as u32,
                token_type,
                token_modifiers_bitset: modifiers,
            });

            prev_line = line;
            prev_col = col;
        }

        result
    }
}