aam-core 2.8.0

//! Parser stage: builds an Abstract Syntax Tree from tokens.
//!
//! The Parser consumes tokens from the Lexer and produces an AST,
//! preserving line number information for error diagnostics.

use crate::error::{AamlError, ErrorDiagnostics};
use crate::pipeline::lexer::{Token, TokenKind};
use crate::pipeline::tasks::{ExecutionTask, ParseTask};

/// Represents a value in the AST, supporting nested structures.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ValueNode<'a> {
    Literal(std::borrow::Cow<'a, str>),
    Object(std::sync::Arc<[(std::borrow::Cow<'a, str>, Self)]>),
    List(std::sync::Arc<[Self]>),
}

impl std::fmt::Display for ValueNode<'_> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            ValueNode::Literal(s) => write!(f, "{s}"),
            ValueNode::Object(pairs) => {
                let formatted_pairs: Vec<String> =
                    pairs.iter().map(|(k, v)| format!("{k} = {v}")).collect();
                write!(f, "{{ {} }}", formatted_pairs.join(", "))
            }
            ValueNode::List(items) => {
                let formatted_items: Vec<String> = items.iter().map(ToString::to_string).collect();
                write!(f, "[{}]", formatted_items.join(", "))
            }
        }
    }
}

/// A node in the Abstract Syntax Tree.
#[derive(Debug, Clone)]
pub enum AstNode<'a> {
    /// Key-value assignment: `key = value`
    Assignment {
        key: std::borrow::Cow<'a, str>,
        value: ValueNode<'a>,
        line: usize,
    },
    /// Directive: `@directive_name arguments`
    Directive {
        name: std::borrow::Cow<'a, str>,
        args: std::borrow::Cow<'a, str>,
        body: Option<ValueNode<'a>>,
        line: usize,
    },
}

impl AstNode<'_> {
    /// Returns the line number where this node appears
    #[must_use]
    pub const fn line(&self) -> usize {
        match self {
            AstNode::Assignment { line, .. } | AstNode::Directive { line, .. } => *line,
        }
    }
}

/// Trait for parsing tokens into an AST.
pub trait Parser: Send + Sync {
    /// Parses a token stream into an AST.
    ///
    /// # Errors
    /// Returns `AamlError::ParseError` if the token stream is malformed.
    fn parse<'a>(&self, tokens: &[Token<'a>]) -> Result<Vec<AstNode<'a>>, AamlError>;

    /// Parses with recovery and returns both partial AST and all parse errors.
    fn parse_with_recovery<'a>(&self, tokens: &[Token<'a>]) -> ParseOutput<'a>;

    /// Generates parse tasks from an AST
    fn generate_parse_tasks<'a>(&self, ast: &[AstNode<'a>]) -> Vec<ParseTask<'a>>;

    /// Generates execution tasks from an AST
    fn generate_execution_tasks<'a>(&self, ast: &[AstNode<'a>]) -> Vec<ExecutionTask<'a>>;
}

/// Parser output for LSP use-cases where partial AST plus diagnostics are required.
pub struct ParseOutput<'a> {
    pub ast: Vec<AstNode<'a>>,
    pub errors: Vec<AamlError>,
}

/// Default implementation of the Parser stage.
pub struct DefaultParser;

impl DefaultParser {
    #[must_use]
    pub const fn new() -> Self {
        Self
    }

    /// Filters out comment and newline tokens
    fn filter_tokens<'a, 'b>(tokens: &'b [Token<'a>]) -> Vec<&'b Token<'a>> {
        use crate::pipeline::lexer::TokenKind;
        tokens
            .iter()
            .filter(|t| t.kind != TokenKind::Comment)
            .collect()
    }

    /// Parses assignment tokens: `identifier = value`
    fn parse_assignment<'a>(
        tokens: &[&Token<'a>],
        start: usize,
    ) -> Result<(std::borrow::Cow<'a, str>, ValueNode<'a>, usize), AamlError> {
        use crate::pipeline::lexer::TokenKind;

        if start >= tokens.len() {
            return Err(AamlError::ParseError {
                line: tokens.get(start).map_or(1, |t| t.line),
                content: "incomplete assignment".to_string(),
                details: "Expected: key = value".to_string(),
                diagnostics: Some(Box::new(ErrorDiagnostics::new(
                    "Incomplete assignment",
                    "Assignment must have at least 3 tokens: key, =, value".to_string(),
                    "Check format: key = value".to_string(),
                ))),
            });
        }

        if tokens[start].kind != TokenKind::Identifier {
            return Err(AamlError::ParseError {
                line: tokens[start].line,
                content: format!("Expected identifier, got {:?}", tokens[start].kind),
                details: "First token of assignment must be an identifier".to_string(),
                diagnostics: None,
            });
        }

        let mut assign_pos = start + 1;
        while assign_pos < tokens.len()
            && tokens[assign_pos].kind != TokenKind::Assign
            && tokens[assign_pos].kind != TokenKind::Newline
        {
            assign_pos += 1;
        }

        if assign_pos >= tokens.len() || tokens[assign_pos].kind != TokenKind::Assign {
            let got = tokens.get(start + 1).map_or("<eof>", |t| t.text.as_ref());
            return Err(AamlError::ParseError {
                line: tokens.get(start + 1).map_or(tokens[start].line, |t| t.line),
                content: format!("Expected '=', got '{got}'"),
                details: "Assignment operator '=' expected after key".to_string(),
                diagnostics: None,
            });
        }

        let key = if assign_pos == start + 1 {
            tokens[start].text.clone()
        } else {
            let key_text = tokens[start..assign_pos]
                .iter()
                .map(|t| t.text.as_ref())
                .collect::<Vec<_>>()
                .join(" ");
            key_text.into()
        };

        let (value, consumed) = Self::parse_value(tokens, assign_pos + 1)?;
        Ok((key, value, consumed))
    }

    fn parse_braced_literal<'a, 'b>(
        tokens: &'b [&'b Token<'a>],
        start: usize,
    ) -> Result<(ValueNode<'a>, usize), AamlError> {
        use crate::pipeline::lexer::TokenKind;

        let mut depth = 0_i32;
        let mut pos = start;

        while pos < tokens.len() {
            match tokens[pos].kind {
                TokenKind::LeftBrace => depth += 1,
                TokenKind::RightBrace => {
                    depth -= 1;
                    if depth == 0 {
                        let first = &tokens[start];
                        let last = &tokens[pos];
                        let text = &first.source[first.start..last.end];
                        return Ok((ValueNode::Literal(text.into()), pos + 1));
                    }
                }
                TokenKind::Newline if depth == 0 => break,
                _ => {}
            }
            pos += 1;
        }

        Err(AamlError::ParseError {
            line: tokens[start].line,
            content: "unclosed brace".to_string(),
            details: "Expected '}' to close inline object".to_string(),
            diagnostics: None,
        })
    }

    /// Parses a value (which may be literal, inline object, or inline list)
    fn parse_value<'a, 'b>(
        tokens: &'b [&'b Token<'a>],
        start: usize,
    ) -> Result<(ValueNode<'a>, usize), AamlError> {
        use crate::pipeline::lexer::TokenKind;

        if start >= tokens.len() {
            return Err(AamlError::ParseError {
                line: tokens.len(),
                content: "unexpected end of input".to_string(),
                details: "Expected a value after '='".to_string(),
                diagnostics: None,
            });
        }

        match &tokens[start].kind {
            TokenKind::LeftBrace => {
                // Prefer structured inline object; fallback to raw braced literal for nested schema-like values.
                match Self::parse_inline_object(tokens, start) {
                    Ok((obj, consumed)) => Ok((obj, consumed)),
                    Err(_) => Self::parse_braced_literal(tokens, start),
                }
            }
            TokenKind::LeftBracket => {
                // Inline list
                let (list, consumed) = Self::parse_inline_list(tokens, start)?;
                Ok((list, consumed))
            }
            _ => {
                let end_pos = tokens[start..]
                    .iter()
                    .position(|t| {
                        matches!(
                            t.kind,
                            TokenKind::Newline
                                | TokenKind::RightBrace
                                | TokenKind::RightBracket
                                | TokenKind::Comma
                        )
                    })
                    .map_or(tokens.len(), |p| start + p);

                if end_pos == start {
                    return Err(AamlError::ParseError {
                        line: tokens[start].line,
                        content: "invalid value".to_string(),
                        details: "Could not parse literal value".to_string(),
                        diagnostics: None,
                    });
                }

                // Preserve the exact source text between the first and last token
                // so whitespace and `=` characters inside values are not lost.
                let first = &tokens[start];
                let last = &tokens[end_pos - 1];
                let exact = &first.source[first.start..last.end];

                Ok((ValueNode::Literal(exact.into()), end_pos))
            }
        }
    }

    /// Parses an inline object: `{ key = val, key = val, ... }`
    fn parse_inline_object<'a, 'b>(
        tokens: &'b [&'b Token<'a>],
        start: usize,
    ) -> Result<(ValueNode<'a>, usize), AamlError> {
        use crate::pipeline::lexer::TokenKind;

        let mut pairs = Vec::new();
        let mut pos = start + 1;
        let mut expect_field = true;
        let mut has_any_field = false;

        while pos < tokens.len() {
            match tokens[pos].kind {
                TokenKind::RightBrace if expect_field && has_any_field => {
                    return Err(AamlError::ParseError {
                        line: tokens[pos].line,
                        content: "trailing comma in inline object".to_string(),
                        details: "Expected another field after ',' or remove trailing comma"
                            .to_string(),
                        diagnostics: None,
                    });
                }
                TokenKind::RightBrace => return Ok((ValueNode::Object(pairs.into()), pos + 1)),
                TokenKind::Identifier if expect_field => {
                    pos = Self::parse_inline_object_pair(tokens, pos, &mut pairs)?;
                    expect_field = false;
                    has_any_field = true;
                }
                TokenKind::Identifier => {
                    return Err(AamlError::ParseError {
                        line: tokens[pos].line,
                        content: "missing comma in inline object".to_string(),
                        details: "Expected ',' between object fields".to_string(),
                        diagnostics: None,
                    });
                }
                TokenKind::Comma if expect_field => {
                    return Err(AamlError::ParseError {
                        line: tokens[pos].line,
                        content: "unexpected comma in inline object".to_string(),
                        details: "Expected an object field before ','".to_string(),
                        diagnostics: None,
                    });
                }
                TokenKind::Comma => {
                    expect_field = true;
                    pos += 1;
                }
                _ => {
                    return Err(AamlError::ParseError {
                        line: tokens[pos].line,
                        content: "invalid inline object format".to_string(),
                        details: "Expected identifier or closing brace".to_string(),
                        diagnostics: None,
                    });
                }
            }
        }

        Err(AamlError::ParseError {
            line: tokens[start].line,
            content: "unclosed brace".to_string(),
            details: "Expected '}' to close inline object".to_string(),
            diagnostics: None,
        })
    }

    fn parse_inline_object_pair<'a, 'b>(
        tokens: &'b [&'b Token<'a>],
        pos: usize,
        pairs: &mut Vec<(std::borrow::Cow<'a, str>, ValueNode<'a>)>,
    ) -> Result<usize, AamlError> {
        use crate::pipeline::lexer::TokenKind;
        let key: std::borrow::Cow<'a, str> = tokens[pos].text.clone();
        if pos + 1 >= tokens.len() || tokens[pos + 1].kind != TokenKind::Assign {
            return Err(AamlError::ParseError {
                line: tokens[pos].line,
                content: "invalid inline object format".to_string(),
                details: "Expected '=' after key".to_string(),
                diagnostics: None,
            });
        }

        let value_pos = pos + 2;
        let (value, next_pos) = if value_pos < tokens.len()
            && matches!(
                tokens[value_pos].kind,
                TokenKind::RightBrace | TokenKind::Comma
            ) {
            (ValueNode::Literal("".into()), value_pos)
        } else {
            Self::parse_value(tokens, value_pos)?
        };
        pairs.push((key, value));
        Ok(next_pos)
    }

    /// Parses an inline list: `[item, item, ...]`
    fn parse_inline_list<'a, 'b>(
        tokens: &'b [&'b Token<'a>],
        start: usize,
    ) -> Result<(ValueNode<'a>, usize), AamlError> {
        use crate::pipeline::lexer::TokenKind;

        let mut items = Vec::new();
        let mut pos = start + 1;

        while pos < tokens.len() {
            if tokens[pos].kind == TokenKind::RightBracket {
                return Ok((ValueNode::List(items.into()), pos + 1));
            }

            if tokens[pos].kind == TokenKind::Comma {
                pos += 1;
                continue;
            }

            let (value, next_pos) = Self::parse_value(tokens, pos)?;
            items.push(value);
            pos = next_pos;
        }

        Err(AamlError::ParseError {
            line: tokens[start].line,
            content: "unclosed bracket".to_string(),
            details: "Expected ']' to close inline list".to_string(),
            diagnostics: None,
        })
    }

    fn missing_directive_name_error(line: usize) -> AamlError {
        AamlError::ParseError {
            line,
            content: "@".to_string(),
            details: "Directive name expected after '@'".to_string(),
            diagnostics: Some(Box::new(ErrorDiagnostics::new(
                "Missing directive name",
                "Directive requires a name after '@'".to_string(),
                "Use format: @directive_name arguments".to_string(),
            ))),
        }
    }

    fn should_stop_directive_args(kind: &TokenKind, brace_count: i32, bracket_count: i32) -> bool {
        brace_count == 0
            && bracket_count == 0
            && (*kind == TokenKind::At || *kind == TokenKind::Newline)
    }

    fn collect_directive_args(tokens_filtered: &[&Token<'_>], start_pos: usize) -> (String, usize) {
        let mut args = String::new();
        let mut arg_pos = start_pos;
        let mut brace_count = 0;
        let mut bracket_count = 0;

        while arg_pos < tokens_filtered.len() {
            let tk = tokens_filtered[arg_pos];
            match tk.kind {
                TokenKind::LeftBrace => brace_count += 1,
                TokenKind::RightBrace => brace_count -= 1,
                TokenKind::LeftBracket => bracket_count += 1,
                TokenKind::RightBracket => bracket_count -= 1,
                _ => {}
            }

            if Self::should_stop_directive_args(&tk.kind, brace_count, bracket_count) {
                break;
            }

            if !args.is_empty() {
                args.push(' ');
            }
            args.push_str(&tk.text);
            arg_pos += 1;
        }

        (args, arg_pos)
    }

    fn validate_schema_directive_args(line: usize, trimmed_args: &str) -> Result<(), AamlError> {
        let (name_part, body_opt) = trimmed_args.split_once('{').unwrap_or((trimmed_args, ""));
        if name_part.trim().is_empty() {
            return Err(AamlError::ParseError {
                line,
                content: "@schema".to_string(),
                details: "Schema name is missing".to_string(),
                diagnostics: None,
            });
        }

        if !trimmed_args.contains('{') || !trimmed_args.ends_with('}') {
            return Err(AamlError::ParseError {
                line,
                content: "@schema".to_string(),
                details: "Schema block must end with '}'".to_string(),
                diagnostics: None,
            });
        }

        let body = body_opt.trim_end_matches('}').trim();
        for field in body.split([',', '\n']) {
            let field = field.trim();
            if field.is_empty() {
                continue;
            }

            if !field.contains(':') {
                return Err(AamlError::ParseError {
                    line,
                    content: "invalid schema field".to_string(),
                    details: format!("Field '{field}' must be of the form 'name: type'"),
                    diagnostics: None,
                });
            }
        }

        Ok(())
    }

    fn parse_directive<'a>(
        tokens_filtered: &[&Token<'a>],
        pos: &mut usize,
        line: usize,
    ) -> Result<AstNode<'a>, AamlError> {
        if *pos + 1 >= tokens_filtered.len() {
            return Err(Self::missing_directive_name_error(line));
        }

        let dir_name: std::borrow::Cow<'a, str> = tokens_filtered[*pos + 1].text.clone();
        let (args, arg_pos) = Self::collect_directive_args(tokens_filtered, *pos + 2);

        *pos = arg_pos;

        let trimmed_args = args.trim();
        if &*dir_name == "schema" {
            Self::validate_schema_directive_args(line, trimmed_args)?;
        }

        Ok(AstNode::Directive {
            name: dir_name,
            args: trimmed_args.to_string().into(),
            body: None,
            line,
        })
    }

    fn synchronize(tokens_filtered: &[&Token<'_>], mut pos: usize) -> usize {
        while pos < tokens_filtered.len() {
            let kind = &tokens_filtered[pos].kind;
            if *kind == TokenKind::Newline || *kind == TokenKind::RightBrace {
                return pos + 1;
            }
            pos += 1;
        }
        pos
    }

    fn scope_from_key(key: &str) -> std::borrow::Cow<'static, str> {
        if let Some((prefix, _)) = key.rsplit_once('.') {
            return format!("root::{}", prefix.replace('.', "::")).into();
        }
        std::borrow::Cow::Borrowed("root")
    }

    fn emit_assignment_parse_tasks<'a>(
        tasks: &mut Vec<ParseTask<'a>>,
        key: &str,
        value: &ValueNode<'a>,
        line: usize,
    ) {
        tasks.push(ParseTask::ProcessVariable {
            variable_name: key.to_string().into(),
            value: value.to_string().into(),
            scope: Self::scope_from_key(key).into_owned().into(),
            line,
        });

        if let ValueNode::Object(pairs) = value {
            for (field, child_value) in pairs.iter() {
                let child_key = format!("{key}.{field}");
                Self::emit_assignment_parse_tasks(tasks, &child_key, child_value, line);
            }
        }
    }

    fn build_type_task<'a>(args: &str, line: usize) -> ParseTask<'a> {
        let (type_name, type_spec) = args
            .split_once('=')
            .map_or(("", ""), |(lhs, rhs)| (lhs.trim(), rhs.trim()));

        ParseTask::RegisterType {
            type_name: type_name.to_string().into(),
            type_spec: type_spec.to_string().into(),
            line,
        }
    }

    fn build_schema_task<'a>(args: &str, line: usize) -> ParseTask<'a> {
        let schema_name = args.split_whitespace().next().unwrap_or("").to_string();
        let body = args
            .split_once('{')
            .and_then(|(_, b)| b.rsplit_once('}'))
            .map_or("", |(b, _)| b);
        let fields = body
            .split([',', '\n'])
            .filter_map(|field_def| {
                let (field_name, type_name) = field_def.split_once(':')?;
                let field_name = field_name.trim();
                let type_name = type_name.trim();

                if field_name.is_empty() {
                    return None;
                }

                Some(format!("{field_name}:{type_name}"))
            })
            .collect::<Vec<_>>()
            .join(",");

        ParseTask::RegisterSchema {
            schema_name: schema_name.into(),
            fields: fields.into(),
            line,
        }
    }

    fn build_directive_task<'a>(
        name: &std::borrow::Cow<'a, str>,
        args: &std::borrow::Cow<'a, str>,
        line: usize,
    ) -> ParseTask<'a> {
        match &**name {
            "type" => Self::build_type_task(args, line),
            "schema" => Self::build_schema_task(args, line),
            "derive" => ParseTask::ResolveDeriveImport {
                derive_path: args.clone(),
                line,
            },
            _ => ParseTask::ExecuteDirective {
                directive_name: name.clone(),
                arguments: args.clone(),
                line,
            },
        }
    }
}

impl Default for DefaultParser {
    fn default() -> Self {
        Self::new()
    }
}

impl Parser for DefaultParser {
    fn parse<'a>(&self, tokens: &[Token<'a>]) -> Result<Vec<AstNode<'a>>, AamlError> {
        let result = self.parse_with_recovery(tokens);
        if result.errors.is_empty() {
            Ok(result.ast)
        } else {
            Err(result
                .errors
                .into_iter()
                .next()
                .unwrap_or_else(|| AamlError::ParseError {
                    line: 1,
                    content: "parse failed".to_string(),
                    details: "Unknown parser failure".to_string(),
                    diagnostics: None,
                }))
        }
    }

    fn parse_with_recovery<'a>(&self, tokens: &[Token<'a>]) -> ParseOutput<'a> {
        use crate::pipeline::lexer::TokenKind;

        let mut ast: Vec<AstNode<'a>> = Vec::new();
        let mut errors = Vec::new();
        let tokens_filtered = Self::filter_tokens(tokens);
        let mut pos = 0;

        while pos < tokens_filtered.len() {
            let token = tokens_filtered[pos];

            match &token.kind {
                TokenKind::At => {
                    match Self::parse_directive(&tokens_filtered, &mut pos, token.line) {
                        Ok(dir_node) => ast.push(dir_node),
                        Err(err) => {
                            errors.push(err);
                            pos = Self::synchronize(&tokens_filtered, pos + 1);
                        }
                    }
                }
                TokenKind::Identifier => match Self::parse_assignment(&tokens_filtered, pos) {
                    Ok((key, value, new_pos)) => {
                        ast.push(AstNode::Assignment {
                            key,
                            value,
                            line: token.line,
                        });
                        pos = new_pos;
                    }
                    Err(err) => {
                        errors.push(err);
                        pos = Self::synchronize(&tokens_filtered, pos + 1);
                    }
                },
                _ => pos += 1,
            }
        }

        ParseOutput { ast, errors }
    }
    fn generate_parse_tasks<'a>(&self, ast: &[AstNode<'a>]) -> Vec<ParseTask<'a>> {
        let mut tasks = Vec::new();

        for node in ast {
            match node {
                AstNode::Assignment { key, value, line } => {
                    Self::emit_assignment_parse_tasks(&mut tasks, key, value, *line);
                }
                AstNode::Directive {
                    name, args, line, ..
                } => tasks.push(Self::build_directive_task(name, args, *line)),
            }
        }

        tasks
    }

    fn generate_execution_tasks<'a>(&self, ast: &[AstNode<'a>]) -> Vec<ExecutionTask<'a>> {
        let mut tasks = Vec::new();

        for node in ast {
            match node {
                AstNode::Assignment { key, value, line } => {
                    tasks.push(ExecutionTask::SetValue {
                        key: key.clone(),
                        value: value.to_string().into(),
                        line: *line,
                    });
                }
                // Directives translated to execution tasks...
                AstNode::Directive {
                    name, args, line, ..
                } => {
                    if &**name == "import" {
                        tasks.push(ExecutionTask::ImportFile {
                            file_path: args.clone(),
                            merge_strategy: std::borrow::Cow::Borrowed("merge"),
                            line: *line,
                        });
                    } else if &**name == "derive" {
                        tasks.push(ExecutionTask::ExecuteInheritance {
                            derive_path: args.clone(),
                            child_key: std::borrow::Cow::Borrowed(""), // For AAM v2
                            line: *line,
                        });
                    }
                }
            }
        }
        tasks
    }
}

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

    #[test]
    fn test_parse_simple_assignment() {
        let lexer = DefaultLexer::new();
        let tokens = lexer.tokenize("key = value").unwrap();
        let parser = DefaultParser::new();
        let ast = parser.parse(&tokens).unwrap();

        assert_eq!(ast.len(), 1);
        match &ast[0] {
            AstNode::Assignment { key, value, .. } => {
                assert_eq!(&**key, "key");
                if let ValueNode::Literal(s) = value {
                    assert_eq!(&**s, "value");
                } else {
                    panic!("Expected ValueNode::Literal");
                }
            }
            AstNode::Directive { .. } => panic!("Expected assignment"),
        }
    }

    #[test]
    fn test_parse_directive() {
        let lexer = DefaultLexer::new();
        let tokens = lexer.tokenize("@import base.aam").unwrap();
        let parser = DefaultParser::new();
        let ast = parser.parse(&tokens).unwrap();

        assert_eq!(ast.len(), 1);
        match &ast[0] {
            AstNode::Directive { name, args: _, .. } => {
                assert_eq!(&**name, "import");
            }
            AstNode::Assignment { .. } => panic!("Expected directive"),
        }
    }

    #[test]
    fn test_parse_multiple_assignments() {
        let lexer = DefaultLexer::new();
        let tokens = lexer.tokenize("a = b\nc = d").unwrap();
        let parser = DefaultParser::new();
        let ast = parser.parse(&tokens).unwrap();

        assert_eq!(ast.len(), 2);
    }
}