// SPDX-FileCopyrightText: 2026 Milovann Yanatchkov
// SPDX-License-Identifier: MIT

use crate::ast::*;
use chumsky::prelude::*;
use chumsky::{input::Stream, input::ValueInput};
use lexer::Token;
use logos::Logos;

/// Parses a BIMR DSL source string into a `Program` AST.
///
/// Lexes `input` with `logos`, wraps the token iterator in a `chumsky`
/// `Stream`, and runs the `program` combinator. Returns the AST on success or
/// a list of `Simple` parse errors on failure.
pub fn parse(input: &str) -> Result<Program, Vec<Simple<'_, Token>>> {
    let token_iter = Token::lexer(input).spanned().map(|(tok, span)| match tok {
        Ok(tok) => (tok, SimpleSpan::from(span)),
        Err(()) => (Token::Error, span.into()),
    });

    let token_stream = Stream::from_iter(token_iter).map((0..input.len()).into(), |(t, s)| (t, s));

    program().parse(token_stream).into_result()
}

/// Parser combinator for a complete BIMR program.
///
/// A program is zero or more `statement`s collected into a `Program` node.
fn program<'tokens, I>() -> impl Parser<'tokens, I, Program, extra::Err<Simple<'tokens, Token>>>
where
    I: ValueInput<'tokens, Token = Token, Span = SimpleSpan>,
{
    statement()
        .repeated()
        .collect()
        .map(|statements| Program { statements })
}

/// Parser combinator for a single assignment statement.
///
/// Grammar: `identifier = expression`
fn statement<'tokens, I>() -> impl Parser<'tokens, I, Statement, extra::Err<Simple<'tokens, Token>>>
where
    I: ValueInput<'tokens, Token = Token, Span = SimpleSpan>,
{
    identifier()
        .then_ignore(just(Token::Equal))
        .then(expression())
        .map(|(identifier, expression)| Statement {
            assignment: Assignment {
                identifier,
                expression,
            },
        })
}

/// Parser combinator for an expression (RHS of an assignment).
///
/// Tries, in order: function call, tuple, bare atom (identifier / number /
/// string). Function calls are tried first to avoid mis-parsing `Name(…)` as
/// a bare identifier followed by junk.
fn expression<'tokens, I>()
-> impl Parser<'tokens, I, Expression, extra::Err<Simple<'tokens, Token>>>
where
    I: ValueInput<'tokens, Token = Token, Span = SimpleSpan>,
{
    // Atomic expressions — no tuple here to avoid ambiguity with function call parens.
    let atom = || {
        choice((
            identifier().map(Expression::Identifier),
            select! { Token::Float(bits) => Expression::Float(f64::from_bits(bits)) },
            select! { Token::Integer(i) => Expression::Integer(i) },
            select! { Token::String(s) => Expression::String(s) },
        ))
    };

    // Tuple: (expr, expr) or (expr, expr, expr) — at least two elements.
    // Used for inline point literals inside List(...) args.
    // Constructed as a closure so it can be used in multiple places without Clone.
    let tuple = || {
        atom()
            .separated_by(just(Token::Comma))
            .at_least(2)
            .collect::<Vec<_>>()
            .delimited_by(just(Token::LeftParen), just(Token::RightParen))
            .map(Expression::Tuple)
    };

    // A single argument: either `ident=expr` (keyword) or `expr` (positional).
    // atom() is used as the value expression to keep things non-recursive.
    let argument = identifier()
        .then_ignore(just(Token::Equal))
        .then(atom())
        .map(|(key, val)| Argument::Keyword(key, val))
        .or(choice((tuple(), atom())).map(Argument::Positional));

    // Function call: Identifier(arg, ...)
    let function_call = identifier()
        .then(
            argument
                .separated_by(just(Token::Comma))
                .allow_trailing()
                .collect::<Vec<_>>()
                .delimited_by(just(Token::LeftParen), just(Token::RightParen)),
        )
        .map(|(name, arguments)| Expression::FunctionCall(FunctionCall { name, arguments }));

    choice((function_call, tuple(), atom()))
}

/// Parser combinator that matches a single `Token::Identifier` and returns
/// its string value.
fn identifier<'tokens, I>() -> impl Parser<'tokens, I, String, extra::Err<Simple<'tokens, Token>>>
where
    I: ValueInput<'tokens, Token = Token, Span = SimpleSpan>,
{
    select! { Token::Identifier(s) => s }
}

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

    fn parse_one(src: &str) -> Program {
        parse(src).expect("parse failed")
    }

    // ── Basic assignments ─────────────────────────────────────────────────────

    #[test]
    fn test_parse_integer() {
        let prog = parse_one("x = 42");
        assert_eq!(prog.statements.len(), 1);
        let stmt = &prog.statements[0].assignment;
        assert_eq!(stmt.identifier, "x");
        assert!(matches!(stmt.expression, Expression::Integer(42)));
    }

    #[test]
    #[allow(clippy::approx_constant)]
    fn test_parse_float() {
        let prog = parse_one("x = 3.14");
        let expr = &prog.statements[0].assignment.expression;
        assert!(matches!(expr, Expression::Float(v) if (v - 3.14).abs() < 1e-10));
    }

    #[test]
    fn test_parse_negative_int() {
        let prog = parse_one("x = -7");
        let expr = &prog.statements[0].assignment.expression;
        assert!(matches!(expr, Expression::Integer(-7)));
    }

    #[test]
    fn test_parse_string() {
        // The lexer preserves the surrounding quotes in the token value.
        let prog = parse_one(r#"x = "hello""#);
        let expr = &prog.statements[0].assignment.expression;
        assert!(matches!(expr, Expression::String(s) if s == r#""hello""#));
    }

    #[test]
    fn test_parse_identifier() {
        let prog = parse_one("x = y");
        let expr = &prog.statements[0].assignment.expression;
        assert!(matches!(expr, Expression::Identifier(id) if id == "y"));
    }

    // ── Function calls ────────────────────────────────────────────────────────

    #[test]
    fn test_parse_point() {
        let prog = parse_one("p1 = Point(0, 0, 0)");
        let expr = &prog.statements[0].assignment.expression;
        if let Expression::FunctionCall(fc) = expr {
            assert_eq!(fc.name, "Point");
            assert_eq!(fc.arguments.len(), 3);
        } else {
            panic!("expected FunctionCall");
        }
    }

    #[test]
    fn test_parse_line() {
        let prog = parse_one("l1 = Line(p1, p2)");
        let expr = &prog.statements[0].assignment.expression;
        if let Expression::FunctionCall(fc) = expr {
            assert_eq!(fc.name, "Line");
            assert_eq!(fc.arguments.len(), 2);
        } else {
            panic!("expected FunctionCall");
        }
    }

    #[test]
    fn test_parse_wall_with_keywords() {
        let prog = parse_one("w1 = Wall(l1, thickness=200, height=3000)");
        let expr = &prog.statements[0].assignment.expression;
        if let Expression::FunctionCall(fc) = expr {
            assert_eq!(fc.name, "Wall");
            assert_eq!(fc.arguments.len(), 3);
            assert!(matches!(&fc.arguments[0], Argument::Positional(_)));
            assert!(matches!(&fc.arguments[1], Argument::Keyword(k, _) if k == "thickness"));
            assert!(matches!(&fc.arguments[2], Argument::Keyword(k, _) if k == "height"));
        } else {
            panic!("expected FunctionCall");
        }
    }

    #[test]
    fn test_parse_each_operator() {
        let src = r#"e1 = Each(div, "Frame", fw=200, height=300, tangent="auto")"#;
        let prog = parse_one(src);
        let expr = &prog.statements[0].assignment.expression;
        if let Expression::FunctionCall(fc) = expr {
            assert_eq!(fc.name, "Each");
            assert!(fc.arguments.len() >= 2);
        } else {
            panic!("expected FunctionCall");
        }
    }

    // ── Multiple statements ───────────────────────────────────────────────────

    #[test]
    fn test_parse_multiple_statements() {
        let src = "p1 = Point(0, 0, 0)\np2 = Point(100, 0, 0)\nl1 = Line(p1, p2)";
        let prog = parse_one(src);
        assert_eq!(prog.statements.len(), 3);
    }

    #[test]
    fn test_parse_empty_program() {
        let prog = parse_one("");
        assert!(prog.statements.is_empty());
    }

    // ── Inline tuples ─────────────────────────────────────────────────────────

    #[test]
    fn test_parse_tuple_in_list() {
        let src = "lst = List((0, 0, 0), (100, 0, 0))";
        let prog = parse_one(src);
        let expr = &prog.statements[0].assignment.expression;
        if let Expression::FunctionCall(fc) = expr {
            assert_eq!(fc.name, "List");
            assert_eq!(fc.arguments.len(), 2);
            for arg in &fc.arguments {
                assert!(matches!(arg, Argument::Positional(Expression::Tuple(_))));
            }
        } else {
            panic!("expected FunctionCall");
        }
    }

    // ── Comments ──────────────────────────────────────────────────────────────

    #[test]
    fn test_parse_with_comments() {
        let src =
            "# a full-line comment\np1 = Point(0, 0, 0) # trailing comment\nl1 = Line(p1, p1)";
        let prog = parse_one(src);
        assert_eq!(prog.statements.len(), 2);
    }

    #[test]
    fn test_parse_comment_before_statement() {
        let prog = parse_one("# nothing here yet\nx = 42");
        assert_eq!(prog.statements.len(), 1);
        assert_eq!(prog.statements[0].assignment.identifier, "x");
    }

    // ── Error handling ────────────────────────────────────────────────────────

    #[test]
    fn test_parse_error_missing_paren() {
        let result = parse("p1 = Point(0, 0");
        assert!(result.is_err());
    }

    #[test]
    fn test_parse_error_unclosed_string() {
        let result = parse(r#"x = "hello"#);
        assert!(result.is_err());
    }

    #[test]
    fn test_parse_error_bad_syntax() {
        let result = parse("= 42");
        assert!(result.is_err());
    }

    #[test]
    fn test_parse_error_unknown_token() {
        let result = parse("x = @@@");
        assert!(result.is_err());
    }
}