blots_core/

formatter.rs

use crate::ast::{BinaryOp, DoStatement, Expr, RecordEntry, RecordKey, SpannedExpr};
use crate::ast_to_source::{expr_to_source};
use crate::values::LambdaArg;

const DEFAULT_MAX_COLUMNS: usize = 80;
const INDENT_SIZE: usize = 2;

/// Format a Blots expression with intelligent line breaking
pub fn format_expr(expr: &SpannedExpr, max_columns: Option<usize>) -> String {
    let max_cols = max_columns.unwrap_or(DEFAULT_MAX_COLUMNS);
    format_expr_impl(expr, max_cols, 0)
}

/// Internal formatting implementation with indentation tracking
fn format_expr_impl(expr: &SpannedExpr, max_cols: usize, indent: usize) -> String {
    // Special handling for lambdas to ensure correct argument formatting
    if let Expr::Lambda { args, body } = &expr.node {
        return format_lambda(args, body, max_cols, indent);
    }

    // Special handling for do blocks - they're inherently multi-line
    // but we want to check if the opening fits on the current line
    if let Expr::DoBlock { .. } = &expr.node {
        return format_multiline(expr, max_cols, indent);
    }

    // First, try single-line formatting using our custom formatter
    let single_line = format_single_line(expr);

    // For expressions with do blocks inside, we can't use single_line.len()
    // because it already contains newlines. Instead, check the first line only.
    let first_line = single_line.lines().next().unwrap_or(&single_line);
    let current_line_length = indent + first_line.len();

    // If the first line fits and there are no newlines, use single-line format
    if !single_line.contains('\n') && current_line_length <= max_cols {
        return single_line;
    }

    // Otherwise, apply smart multi-line formatting based on expression type
    format_multiline(expr, max_cols, indent)
}

/// Format an expression on a single line (respecting our formatting rules)
fn format_single_line(expr: &SpannedExpr) -> String {
    match &expr.node {
        Expr::Output { expr: inner_expr } => {
            format!("output {}", format_single_line(inner_expr))
        }
        Expr::Lambda { args, body } => {
            let args_str: Vec<String> = args.iter().map(lambda_arg_to_str).collect();
            let args_part = if args.len() == 1 && matches!(args[0], LambdaArg::Required(_)) {
                args_str[0].clone()
            } else {
                format!("({})", args_str.join(", "))
            };
            format!("{} => {}", args_part, format_single_line(body))
        }
        Expr::Call { func, args } => {
            let func_str = match &func.node {
                Expr::Lambda { .. } => format!("({})", format_single_line(func)),
                _ => format_single_line(func),
            };
            let args_str: Vec<String> = args.iter().map(format_single_line).collect();
            format!("{}({})", func_str, args_str.join(", "))
        }
        Expr::List(items) => {
            let items_str: Vec<String> = items.iter().map(format_single_line).collect();
            format!("[{}]", items_str.join(", "))
        }
        Expr::Record(entries) => {
            let entries_str: Vec<String> = entries.iter().map(format_record_entry_single_line).collect();
            format!("{{{}}}", entries_str.join(", "))
        }
        // For everything else, use the existing expr_to_source
        _ => expr_to_source(expr),
    }
}

/// Format a record entry on a single line
fn format_record_entry_single_line(entry: &RecordEntry) -> String {
    match &entry.key {
        RecordKey::Static(key) => format!("{}: {}", key, format_single_line(&entry.value)),
        RecordKey::Dynamic(key_expr) => {
            format!("[{}]: {}", format_single_line(key_expr), format_single_line(&entry.value))
        }
        RecordKey::Shorthand(name) => name.clone(),
        RecordKey::Spread(expr) => format_single_line(expr),
    }
}

/// Format an expression across multiple lines
fn format_multiline(expr: &SpannedExpr, max_cols: usize, indent: usize) -> String {
    match &expr.node {
        Expr::Output { expr: inner_expr } => {
            // Format as "output " + formatted inner expression
            let formatted_inner = format_expr_impl(inner_expr, max_cols, indent);
            format!("output {}", formatted_inner)
        }
        Expr::Assignment { ident, value } => {
            format_assignment_multiline(ident, value, max_cols, indent)
        }
        Expr::List(items) => format_list_multiline(items, max_cols, indent),
        Expr::Record(entries) => format_record_multiline(entries, max_cols, indent),
        Expr::Conditional { condition, then_expr, else_expr } => {
            format_conditional_multiline(condition, then_expr, else_expr, max_cols, indent)
        }
        Expr::Call { func, args } => format_call_multiline(func, args, max_cols, indent),
        Expr::BinaryOp { op, left, right } => {
            format_binary_op_multiline(op, left, right, max_cols, indent)
        }
        Expr::DoBlock { statements, return_expr } => {
            format_do_block_multiline(statements, return_expr, max_cols, indent)
        }
        // For other expression types, fall back to single-line
        _ => expr_to_source(expr),
    }
}

/// Format an assignment with line breaks
fn format_assignment_multiline(ident: &str, value: &SpannedExpr, max_cols: usize, indent: usize) -> String {
    // The assignment itself doesn't add indentation, but the value might need it
    // Format as: ident = <formatted_value>
    // The value should be formatted at the same indentation level, not pushed over

    let prefix = format!("{} = ", ident);

    // Format the value at the current indentation level
    // (not at prefix_len which would cause excessive indentation)
    let formatted_value = format_expr_impl(value, max_cols, indent);

    format!("{}{}", prefix, formatted_value)
}

/// Format a list with line breaks
fn format_list_multiline(items: &[SpannedExpr], max_cols: usize, indent: usize) -> String {
    if items.is_empty() {
        return "[]".to_string();
    }

    let inner_indent = indent + INDENT_SIZE;
    let indent_str = make_indent(inner_indent);

    let mut result = "[".to_string();

    for item in items.iter() {
        result.push('\n');
        result.push_str(&indent_str);
        result.push_str(&format_expr_impl(item, max_cols, inner_indent));

        // Add comma after each item (including last for multi-line)
        result.push(',');
    }

    result.push('\n');
    result.push_str(&make_indent(indent));
    result.push(']');

    result
}

/// Format a record with line breaks
fn format_record_multiline(entries: &[RecordEntry], max_cols: usize, indent: usize) -> String {
    if entries.is_empty() {
        return "{}".to_string();
    }

    let inner_indent = indent + INDENT_SIZE;
    let indent_str = make_indent(inner_indent);

    let mut result = "{".to_string();

    for entry in entries {
        result.push('\n');
        result.push_str(&indent_str);
        result.push_str(&format_record_entry(entry, max_cols, inner_indent));

        // Add comma after each entry (including last for multi-line)
        result.push(',');
    }

    result.push('\n');
    result.push_str(&make_indent(indent));
    result.push('}');

    result
}

/// Format a single record entry
fn format_record_entry(entry: &RecordEntry, max_cols: usize, indent: usize) -> String {
    match &entry.key {
        RecordKey::Static(key) => {
            format!("{}: {}", key, format_expr_impl(&entry.value, max_cols, indent))
        }
        RecordKey::Dynamic(key_expr) => {
            format!(
                "[{}]: {}",
                format_expr_impl(key_expr, max_cols, indent),
                format_expr_impl(&entry.value, max_cols, indent)
            )
        }
        RecordKey::Shorthand(name) => name.clone(),
        RecordKey::Spread(expr) => format_expr_impl(expr, max_cols, indent),
    }
}

/// Format a lambda (handles both single-line and multi-line)
fn format_lambda(args: &[LambdaArg], body: &SpannedExpr, max_cols: usize, indent: usize) -> String {
    let args_str: Vec<String> = args.iter().map(lambda_arg_to_str).collect();

    // For single required arguments, omit parentheses
    let args_part = if args.len() == 1 && matches!(args[0], LambdaArg::Required(_)) {
        format!("{} =>", args_str[0])
    } else {
        format!("({}) =>", args_str.join(", "))
    };

    // Special handling for do blocks - keep "=> do {" together
    if let Expr::DoBlock { .. } = &body.node {
        let body_formatted = format_expr_impl(body, max_cols, indent);
        return format!("{} {}", args_part, body_formatted);
    }

    // Try single-line first for other body types
    let single_line_body = format_expr_impl(body, max_cols, indent);
    let single_line = format!("{} {}", args_part, single_line_body);

    // Check only if it's actually single-line and fits
    if !single_line.contains('\n') && indent + single_line.len() <= max_cols {
        return single_line;
    }

    // Otherwise, put body on next line with increased indentation
    let body_indent = indent + INDENT_SIZE;
    format!(
        "{}\n{}{}",
        args_part,
        make_indent(body_indent),
        format_expr_impl(body, max_cols, body_indent)
    )
}

/// Format a conditional with line breaks
fn format_conditional_multiline(
    condition: &SpannedExpr,
    then_expr: &SpannedExpr,
    else_expr: &SpannedExpr,
    max_cols: usize,
    indent: usize,
) -> String {
    let cond_str = format_expr_impl(condition, max_cols, indent);

    // Try to fit "if <condition> then" on one line
    let if_then_prefix = format!("if {} then", cond_str);

    if indent + if_then_prefix.len() <= max_cols {
        // Put then/else clauses on new lines
        let inner_indent = indent + INDENT_SIZE;
        format!(
            "{}\n{}{}\nelse\n{}{}",
            if_then_prefix,
            make_indent(inner_indent),
            format_expr_impl(then_expr, max_cols, inner_indent),
            make_indent(inner_indent),
            format_expr_impl(else_expr, max_cols, inner_indent)
        )
    } else {
        // Everything on separate lines
        let inner_indent = indent + INDENT_SIZE;
        format!(
            "if\n{}{}\nthen\n{}{}\nelse\n{}{}",
            make_indent(inner_indent),
            format_expr_impl(condition, max_cols, inner_indent),
            make_indent(inner_indent),
            format_expr_impl(then_expr, max_cols, inner_indent),
            make_indent(inner_indent),
            format_expr_impl(else_expr, max_cols, inner_indent)
        )
    }
}

/// Format a function call with line breaks
fn format_call_multiline(func: &SpannedExpr, args: &[SpannedExpr], max_cols: usize, indent: usize) -> String {
    let func_str = match &func.node {
        Expr::Lambda { .. } => format!("({})", format_expr_impl(func, max_cols, indent)),
        _ => format_expr_impl(func, max_cols, indent),
    };

    if args.is_empty() {
        return format!("{}()", func_str);
    }

    // Try formatting args on separate lines
    let inner_indent = indent + INDENT_SIZE;
    let indent_str = make_indent(inner_indent);

    let mut result = format!("{}(", func_str);

    for (i, arg) in args.iter().enumerate() {
        result.push('\n');
        result.push_str(&indent_str);
        result.push_str(&format_expr_impl(arg, max_cols, inner_indent));

        if i < args.len() - 1 {
            result.push(',');
        } else {
            // Trailing comma on last arg for multi-line
            result.push(',');
        }
    }

    result.push('\n');
    result.push_str(&make_indent(indent));
    result.push(')');

    result
}

/// Format a binary operation with line breaks
fn format_binary_op_multiline(
    op: &BinaryOp,
    left: &SpannedExpr,
    right: &SpannedExpr,
    max_cols: usize,
    indent: usize,
) -> String {
    let op_str = binary_op_str(op);
    let left_str = format_expr_impl(left, max_cols, indent);

    // Special handling for via/into/where with lambda on the right
    // Try to keep "via lambda" together on the same line
    if matches!(op, BinaryOp::Via | BinaryOp::Into | BinaryOp::Where)
        && let Expr::Lambda { .. } = &right.node {
            // Format the right side (lambda with possible do block)
            let right_str = format_expr_impl(right, max_cols, indent);

            // Check if the first line of the whole expression fits
            // (for lambdas with do blocks, this would be "left via i => do {")
            let first_line_of_right = right_str.lines().next().unwrap_or(&right_str);
            let first_line_combined = format!("{} {} {}", left_str, op_str, first_line_of_right);

            if indent + first_line_combined.len() <= max_cols {
                // The opening line fits! Return the full formatted expression
                // If right_str is multi-line, this will preserve that structure
                if right_str.contains('\n') {
                    // Multi-line lambda (like with do block)
                    let remaining_lines = right_str.lines().skip(1).collect::<Vec<_>>().join("\n");
                    return format!("{} {} {}\n{}", left_str, op_str, first_line_of_right, remaining_lines);
                } else {
                    // Single-line lambda
                    return format!("{} {} {}", left_str, op_str, right_str);
                }
            }

            // If it doesn't fit, break after the operator
            let continued_indent = indent;
            let right_formatted = format_expr_impl(right, max_cols, continued_indent);
            return format!("{} {}\n{}{}", left_str, op_str, make_indent(continued_indent), right_formatted);
        }

    // Default: break before the operator with indentation
    let right_indent = indent + INDENT_SIZE;
    format!(
        "{}\n{}{} {}",
        left_str,
        make_indent(right_indent),
        op_str,
        format_expr_impl(right, max_cols, right_indent)
    )
}

/// Format a do block (always multi-line)
fn format_do_block_multiline(statements: &[DoStatement], return_expr: &SpannedExpr, max_cols: usize, indent: usize) -> String {
    let inner_indent = indent + INDENT_SIZE;
    let indent_str = make_indent(inner_indent);

    let mut result = "do {".to_string();

    for stmt in statements {
        match stmt {
            DoStatement::Expression(e) => {
                result.push('\n');
                result.push_str(&indent_str);
                result.push_str(&format_expr_impl(e, max_cols, inner_indent));
            }
            DoStatement::Comment(c) => {
                result.push('\n');
                result.push_str(&indent_str);
                result.push_str(c);
            }
        }
    }

    result.push('\n');
    result.push_str(&indent_str);
    result.push_str("return ");
    result.push_str(&format_expr_impl(return_expr, max_cols, inner_indent));
    result.push('\n');
    result.push_str(&make_indent(indent));
    result.push('}');

    result
}

/// Convert lambda argument to string
fn lambda_arg_to_str(arg: &LambdaArg) -> String {
    match arg {
        LambdaArg::Required(name) => name.clone(),
        LambdaArg::Optional(name) => format!("{}?", name),
        LambdaArg::Rest(name) => format!("...{}", name),
    }
}

/// Convert binary operator to string
fn binary_op_str(op: &BinaryOp) -> &'static str {
    match op {
        BinaryOp::Add => "+",
        BinaryOp::Subtract => "-",
        BinaryOp::Multiply => "*",
        BinaryOp::Divide => "/",
        BinaryOp::Modulo => "%",
        BinaryOp::Power => "^",
        BinaryOp::Equal => "==",
        BinaryOp::NotEqual => "!=",
        BinaryOp::Less => "<",
        BinaryOp::LessEq => "<=",
        BinaryOp::Greater => ">",
        BinaryOp::GreaterEq => ">=",
        BinaryOp::DotEqual => ".==",
        BinaryOp::DotNotEqual => ".!=",
        BinaryOp::DotLess => ".<",
        BinaryOp::DotLessEq => ".<=",
        BinaryOp::DotGreater => ".>",
        BinaryOp::DotGreaterEq => ".>=",
        BinaryOp::And => "&&",
        BinaryOp::NaturalAnd => "and",
        BinaryOp::Or => "||",
        BinaryOp::NaturalOr => "or",
        BinaryOp::Via => "via",
        BinaryOp::Into => "into",
        BinaryOp::Where => "where",
        BinaryOp::Coalesce => "??",
    }
}

/// Generate indentation string
fn make_indent(indent: usize) -> String {
    " ".repeat(indent)
}

/// Join formatted statements with appropriate spacing based on their original positions
/// Preserves up to 2 empty lines between statements
pub fn join_statements_with_spacing(
    statements: &[(String, usize, usize)], // (formatted_statement, start_line, end_line)
) -> String {
    if statements.is_empty() {
        return String::new();
    }

    let mut result = String::new();

    for (i, (stmt, _start_line, end_line)) in statements.iter().enumerate() {
        result.push_str(stmt);

        // Add newlines between statements
        if i < statements.len() - 1 {
            let next_start_line = statements[i + 1].1;

            // Calculate how many lines apart they are
            // If they're on consecutive lines (end_line = 1, next_start = 2), gap = 0
            // If there's one empty line (end_line = 1, next_start = 3), gap = 1
            // If there's two empty lines (end_line = 1, next_start = 4), gap = 2
            let line_gap = next_start_line.saturating_sub(*end_line).saturating_sub(1);

            // Preserve up to 2 empty lines (which means up to 3 newlines total)
            // 0 empty lines = 1 newline
            // 1 empty line = 2 newlines
            // 2 empty lines = 3 newlines
            // 3+ empty lines = 3 newlines (capped at 2)
            let newlines = std::cmp::min(line_gap + 1, 3);

            for _ in 0..newlines {
                result.push('\n');
            }
        }
    }

    result
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::parser::get_pairs;
    use crate::expressions::pairs_to_expr;

    fn parse_test_expr(source: &str) -> SpannedExpr {
        use crate::parser::Rule;

        let pairs = get_pairs(source).unwrap();

        // Extract the actual expression from the statement wrapper
        for pair in pairs {
            if pair.as_rule() == Rule::statement {
                if let Some(inner_pair) = pair.into_inner().next() {
                    return pairs_to_expr(inner_pair.into_inner()).unwrap();
                }
            }
        }

        panic!("No statement found in parsed input");
    }

    #[test]
    fn test_format_short_list() {
        let expr = parse_test_expr("[1, 2, 3]");
        let formatted = format_expr(&expr, Some(80));
        assert_eq!(formatted, "[1, 2, 3]");
    }

    #[test]
    fn test_format_long_list() {
        let expr = parse_test_expr("[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]");
        let formatted = format_expr(&expr, Some(40));
        assert!(formatted.contains("\n"));
        assert!(formatted.contains("[\n"));
    }

    #[test]
    fn test_format_short_record() {
        let expr = parse_test_expr("{x: 1, y: 2}");
        let formatted = format_expr(&expr, Some(80));
        assert_eq!(formatted, "{x: 1, y: 2}");
    }

    #[test]
    fn test_format_long_record() {
        let expr = parse_test_expr("{name: \"Alice\", age: 30, email: \"alice@example.com\", address: \"123 Main St\"}");
        let formatted = format_expr(&expr, Some(40));
        assert!(formatted.contains("\n"));
        assert!(formatted.contains("{\n"));
    }

    #[test]
    fn test_format_conditional() {
        let expr = parse_test_expr("if very_long_condition_variable > 100 then \"yes\" else \"no\"");
        let formatted = format_expr(&expr, Some(30));
        assert!(formatted.contains("\n"));
    }

    #[test]
    fn test_format_binary_op() {
        let expr = parse_test_expr("very_long_variable_name + another_very_long_variable_name");
        let formatted = format_expr(&expr, Some(30));
        assert!(formatted.contains("\n"));
    }

    #[test]
    fn test_format_lambda() {
        let expr = parse_test_expr("(x, y) => x + y");
        let formatted = format_expr(&expr, Some(80));
        assert_eq!(formatted, "(x, y) => x + y");
    }

    #[test]
    fn test_format_nested_list() {
        let expr = parse_test_expr("[[1, 2, 3], [4, 5, 6], [7, 8, 9]]");
        let formatted = format_expr(&expr, Some(20));
        assert!(formatted.contains("\n"));
    }

    #[test]
    fn test_format_function_call() {
        let expr = parse_test_expr("map([1, 2, 3], x => x * 2)");
        let formatted = format_expr(&expr, Some(80));
        assert_eq!(formatted, "map([1, 2, 3], x => x * 2)");
    }

    #[test]
    fn test_format_do_block() {
        let expr = parse_test_expr("do { x = 1\n  return x }");
        let formatted = format_expr(&expr, Some(80));
        assert!(formatted.contains("do {"));
        assert!(formatted.contains("return"));
    }

    #[test]
    fn test_multiple_statements() {
        use crate::parser::Rule;

        let source = "x = [1, 2, 3, 4, 5]\ny = {name: \"Alice\", age: 30}\nz = x + y";
        let pairs = get_pairs(source).unwrap();

        let mut formatted_statements = Vec::new();

        for pair in pairs {
            if pair.as_rule() == Rule::statement {
                if let Some(inner_pair) = pair.into_inner().next() {
                    if let Ok(expr) = pairs_to_expr(inner_pair.into_inner()) {
                        let formatted = format_expr(&expr, Some(80));
                        formatted_statements.push(formatted);
                    }
                }
            }
        }

        // Should have formatted all 3 statements
        assert_eq!(formatted_statements.len(), 3);
        assert_eq!(formatted_statements[0], "x = [1, 2, 3, 4, 5]");
        assert_eq!(formatted_statements[1], "y = {name: \"Alice\", age: 30}");
        assert_eq!(formatted_statements[2], "z = x + y");
    }

    #[test]
    fn test_comments_are_preserved() {
        use crate::parser::Rule;

        let source = "// Comment 1\nx = [1, 2, 3]\n// Comment 2\ny = x + 1";
        let pairs = get_pairs(source).unwrap();

        let mut formatted_statements = Vec::new();

        for pair in pairs {
            if pair.as_rule() == Rule::statement {
                if let Some(inner_pair) = pair.into_inner().next() {
                    match inner_pair.as_rule() {
                        Rule::comment => {
                            // Preserve comment as-is
                            formatted_statements.push(inner_pair.as_str().to_string());
                        }
                        _ => {
                            // Format as expression
                            if let Ok(expr) = pairs_to_expr(inner_pair.into_inner()) {
                                let formatted = format_expr(&expr, Some(80));
                                formatted_statements.push(formatted);
                            }
                        }
                    }
                }
            }
        }

        // All 4 items should be preserved
        assert_eq!(formatted_statements.len(), 4);
        assert_eq!(formatted_statements[0], "// Comment 1");
        assert_eq!(formatted_statements[1], "x = [1, 2, 3]");
        assert_eq!(formatted_statements[2], "// Comment 2");
        assert_eq!(formatted_statements[3], "y = x + 1");
    }

    #[test]
    fn test_comments_with_formatted_code() {
        use crate::parser::Rule;

        let source = "// Configuration\nconfig = {name: \"test\", debug: true}\n// Process data\nresult = [1, 2, 3]";
        let pairs = get_pairs(source).unwrap();

        let mut formatted_statements = Vec::new();

        for pair in pairs {
            if pair.as_rule() == Rule::statement {
                if let Some(inner_pair) = pair.into_inner().next() {
                    match inner_pair.as_rule() {
                        Rule::comment => {
                            formatted_statements.push(inner_pair.as_str().to_string());
                        }
                        _ => {
                            if let Ok(expr) = pairs_to_expr(inner_pair.into_inner()) {
                                let formatted = format_expr(&expr, Some(80));
                                formatted_statements.push(formatted);
                            }
                        }
                    }
                }
            }
        }

        let result = formatted_statements.join("\n");

        // Should have 4 statements: 2 comments + 2 expressions
        assert_eq!(formatted_statements.len(), 4);
        assert_eq!(formatted_statements[0], "// Configuration");
        assert!(formatted_statements[1].starts_with("config = "));
        assert_eq!(formatted_statements[2], "// Process data");
        assert_eq!(formatted_statements[3], "result = [1, 2, 3]");

        // Verify the full result preserves comments
        assert!(result.contains("// Configuration"));
        assert!(result.contains("// Process data"));
    }

    #[test]
    fn test_output_declaration_with_assignment() {
        use crate::parser::Rule;

        let source = "output result = 42";
        let pairs = get_pairs(source).unwrap();

        let mut formatted_statements = Vec::new();

        for pair in pairs {
            if pair.as_rule() == Rule::statement {
                if let Some(inner_pair) = pair.into_inner().next() {
                    match inner_pair.as_rule() {
                        Rule::output_declaration => {
                            // Should preserve the output keyword
                            let mut inner = inner_pair.into_inner();
                            let assignment_or_ident = inner.next().unwrap();

                            let formatted = if assignment_or_ident.as_rule() == Rule::assignment {
                                // Extract identifier and value from assignment
                                let mut assignment_inner = assignment_or_ident.into_inner();
                                let ident = assignment_inner.next().unwrap().as_str();
                                let value_expr = pairs_to_expr(assignment_inner.next().unwrap().into_inner()).unwrap();
                                let value_formatted = format_expr(&value_expr, Some(80));
                                format!("{} = {}", ident, value_formatted)
                            } else {
                                assignment_or_ident.as_str().to_string()
                            };

                            formatted_statements.push(format!("output {}", formatted));
                        }
                        _ => {
                            if let Ok(expr) = pairs_to_expr(inner_pair.into_inner()) {
                                let formatted = format_expr(&expr, Some(80));
                                formatted_statements.push(formatted);
                            }
                        }
                    }
                }
            }
        }

        assert_eq!(formatted_statements.len(), 1);
        assert_eq!(formatted_statements[0], "output result = 42");
    }

    #[test]
    fn test_output_statement_separate() {
        use crate::parser::Rule;

        let source = "result = 42\noutput result";
        let pairs = get_pairs(source).unwrap();

        let mut formatted_statements = Vec::new();

        for pair in pairs {
            if pair.as_rule() == Rule::statement {
                if let Some(inner_pair) = pair.into_inner().next() {
                    match inner_pair.as_rule() {
                        Rule::output_declaration => {
                            // Should preserve the output keyword
                            let mut inner = inner_pair.into_inner();
                            let assignment_or_ident = inner.next().unwrap();

                            let formatted = if assignment_or_ident.as_rule() == Rule::assignment {
                                // Extract identifier and value from assignment
                                let mut assignment_inner = assignment_or_ident.into_inner();
                                let ident = assignment_inner.next().unwrap().as_str();
                                let value_expr = pairs_to_expr(assignment_inner.next().unwrap().into_inner()).unwrap();
                                let value_formatted = format_expr(&value_expr, Some(80));
                                format!("{} = {}", ident, value_formatted)
                            } else {
                                assignment_or_ident.as_str().to_string()
                            };

                            formatted_statements.push(format!("output {}", formatted));
                        }
                        _ => {
                            if let Ok(expr) = pairs_to_expr(inner_pair.into_inner()) {
                                let formatted = format_expr(&expr, Some(80));
                                formatted_statements.push(formatted);
                            }
                        }
                    }
                }
            }
        }

        assert_eq!(formatted_statements.len(), 2);
        assert_eq!(formatted_statements[0], "result = 42");
        assert_eq!(formatted_statements[1], "output result");
    }

    #[test]
    fn test_output_with_complex_expression() {
        use crate::parser::Rule;

        let source = "output total = [1, 2, 3] into sum";
        let pairs = get_pairs(source).unwrap();

        let mut formatted_statements = Vec::new();

        for pair in pairs {
            if pair.as_rule() == Rule::statement {
                if let Some(inner_pair) = pair.into_inner().next() {
                    match inner_pair.as_rule() {
                        Rule::output_declaration => {
                            // Should preserve the output keyword
                            let mut inner = inner_pair.into_inner();
                            let assignment_or_ident = inner.next().unwrap();

                            let formatted = if assignment_or_ident.as_rule() == Rule::assignment {
                                // Extract identifier and value from assignment
                                let mut assignment_inner = assignment_or_ident.into_inner();
                                let ident = assignment_inner.next().unwrap().as_str();
                                let value_expr = pairs_to_expr(assignment_inner.next().unwrap().into_inner()).unwrap();
                                let value_formatted = format_expr(&value_expr, Some(80));
                                format!("{} = {}", ident, value_formatted)
                            } else {
                                assignment_or_ident.as_str().to_string()
                            };

                            formatted_statements.push(format!("output {}", formatted));
                        }
                        _ => {
                            if let Ok(expr) = pairs_to_expr(inner_pair.into_inner()) {
                                let formatted = format_expr(&expr, Some(80));
                                formatted_statements.push(formatted);
                            }
                        }
                    }
                }
            }
        }

        assert_eq!(formatted_statements.len(), 1);
        assert_eq!(formatted_statements[0], "output total = [1, 2, 3] into sum");
    }

    #[test]
    fn test_end_of_line_comments() {
        use crate::parser::Rule;

        let source = "x = 5  // this is an end-of-line comment\ny = 10";
        let pairs = get_pairs(source).unwrap();

        let mut formatted_statements = Vec::new();

        for pair in pairs {
            if pair.as_rule() == Rule::statement {
                let mut inner_pairs = pair.into_inner();

                if let Some(first_pair) = inner_pairs.next() {
                    let formatted = match first_pair.as_rule() {
                        Rule::comment => first_pair.as_str().to_string(),
                        _ => {
                            if let Ok(expr) = pairs_to_expr(first_pair.into_inner()) {
                                format_expr(&expr, Some(80))
                            } else {
                                continue;
                            }
                        }
                    };

                    // Check for end-of-line comment (second element in statement)
                    if let Some(eol_comment) = inner_pairs.next() {
                        if eol_comment.as_rule() == Rule::comment {
                            formatted_statements.push(format!("{}  {}", formatted, eol_comment.as_str()));
                        } else {
                            formatted_statements.push(formatted);
                        }
                    } else {
                        formatted_statements.push(formatted);
                    }
                }
            }
        }

        assert_eq!(formatted_statements.len(), 2);
        assert_eq!(formatted_statements[0], "x = 5  // this is an end-of-line comment");
        assert_eq!(formatted_statements[1], "y = 10");
    }

    #[test]
    fn test_multiple_end_of_line_comments() {
        use crate::parser::Rule;

        let source = "a = 1  // comment 1\nb = 2  // comment 2\nc = 3";
        let pairs = get_pairs(source).unwrap();

        let mut formatted_statements = Vec::new();

        for pair in pairs {
            if pair.as_rule() == Rule::statement {
                let mut inner_pairs = pair.into_inner();

                if let Some(first_pair) = inner_pairs.next() {
                    let formatted = match first_pair.as_rule() {
                        Rule::comment => first_pair.as_str().to_string(),
                        _ => {
                            if let Ok(expr) = pairs_to_expr(first_pair.into_inner()) {
                                format_expr(&expr, Some(80))
                            } else {
                                continue;
                            }
                        }
                    };

                    // Check for end-of-line comment
                    if let Some(eol_comment) = inner_pairs.next() {
                        if eol_comment.as_rule() == Rule::comment {
                            formatted_statements.push(format!("{}  {}", formatted, eol_comment.as_str()));
                        } else {
                            formatted_statements.push(formatted);
                        }
                    } else {
                        formatted_statements.push(formatted);
                    }
                }
            }
        }

        assert_eq!(formatted_statements.len(), 3);
        assert_eq!(formatted_statements[0], "a = 1  // comment 1");
        assert_eq!(formatted_statements[1], "b = 2  // comment 2");
        assert_eq!(formatted_statements[2], "c = 3");
    }

    #[test]
    fn test_eol_comments_not_joined_with_next_line() {
        use crate::parser::Rule;

        // Ensure statements with EOL comments remain on their own line
        let source = "x = 1  // first value\ny = 2  // second value\nz = x + y";
        let pairs = get_pairs(source).unwrap();

        let mut formatted_statements = Vec::new();

        for pair in pairs {
            if pair.as_rule() == Rule::statement {
                let mut inner_pairs = pair.into_inner();

                if let Some(first_pair) = inner_pairs.next() {
                    let formatted = match first_pair.as_rule() {
                        Rule::comment => first_pair.as_str().to_string(),
                        _ => {
                            if let Ok(expr) = pairs_to_expr(first_pair.into_inner()) {
                                format_expr(&expr, Some(80))
                            } else {
                                continue;
                            }
                        }
                    };

                    // Check for end-of-line comment
                    if let Some(eol_comment) = inner_pairs.next() {
                        if eol_comment.as_rule() == Rule::comment {
                            formatted_statements.push(format!("{}  {}", formatted, eol_comment.as_str()));
                        } else {
                            formatted_statements.push(formatted);
                        }
                    } else {
                        formatted_statements.push(formatted);
                    }
                }
            }
        }

        // Join with newlines - each statement should be on its own line
        let result = formatted_statements.join("\n");

        assert_eq!(formatted_statements.len(), 3);
        // Verify no statement got joined into one line
        assert_eq!(result.lines().count(), 3);
        assert_eq!(formatted_statements[0], "x = 1  // first value");
        assert_eq!(formatted_statements[1], "y = 2  // second value");
        assert_eq!(formatted_statements[2], "z = x + y");

        // Ensure the result doesn't contain any line with multiple statements
        for line in result.lines() {
            // Count equals signs - should only be 1 per line
            assert_eq!(line.matches('=').count(), 1, "Line should not contain multiple statements: {}", line);
        }
    }

    #[test]
    fn test_eol_comments_with_line_breaking() {
        use crate::parser::Rule;

        // Test that expressions can still be broken across lines when they have EOL comments
        let source = "longRecord = {name: \"Alice\", age: 30, email: \"alice@example.com\", address: \"123 Main St\"}  // user data";
        let pairs = get_pairs(source).unwrap();

        let mut formatted_statements = Vec::new();

        for pair in pairs {
            if pair.as_rule() == Rule::statement {
                let mut inner_pairs = pair.into_inner();

                if let Some(first_pair) = inner_pairs.next() {
                    let formatted = match first_pair.as_rule() {
                        Rule::comment => first_pair.as_str().to_string(),
                        _ => {
                            if let Ok(expr) = pairs_to_expr(first_pair.into_inner()) {
                                // Use a shorter line limit to force breaking
                                format_expr(&expr, Some(40))
                            } else {
                                continue;
                            }
                        }
                    };

                    // Check for end-of-line comment
                    if let Some(eol_comment) = inner_pairs.next() {
                        if eol_comment.as_rule() == Rule::comment {
                            formatted_statements.push(format!("{}  {}", formatted, eol_comment.as_str()));
                        } else {
                            formatted_statements.push(formatted);
                        }
                    } else {
                        formatted_statements.push(formatted);
                    }
                }
            }
        }

        assert_eq!(formatted_statements.len(), 1);
        let result = &formatted_statements[0];

        // The comment should be at the end
        assert!(result.ends_with("// user data"));
        // The expression should be formatted (likely multi-line)
        assert!(result.contains("longRecord = "));
    }

    #[test]
    fn test_actual_line_breaking_behavior() {
        use crate::parser::Rule;

        // Test what actually happens with long lines
        let source = "x = {name: \"Alice\", age: 30, email: \"alice@example.com\", address: \"123 Main St\", city: \"Springfield\"}";
        let pairs = get_pairs(source).unwrap();

        for pair in pairs {
            if pair.as_rule() == Rule::statement {
                if let Some(inner_pair) = pair.into_inner().next() {
                    if let Ok(expr) = pairs_to_expr(inner_pair.into_inner()) {
                        let formatted = format_expr(&expr, Some(40));
                        println!("Formatted output:\n{}", formatted);
                        println!("Line count: {}", formatted.lines().count());

                        // Check if it's actually breaking lines
                        if formatted.lines().count() > 1 {
                            println!("✓ Lines were broken");
                        } else {
                            println!("✗ No line breaking occurred - output is {} chars", formatted.len());
                        }
                    }
                }
            }
        }
    }

    #[test]
    fn test_multiline_input_gets_collapsed() {
        use crate::parser::Rule;

        // Test if multiline input gets collapsed to a single line
        let source = "x = [\n  1,\n  2,\n  3\n]";
        let pairs = get_pairs(source).unwrap();

        for pair in pairs {
            if pair.as_rule() == Rule::statement {
                if let Some(inner_pair) = pair.into_inner().next() {
                    if let Ok(expr) = pairs_to_expr(inner_pair.into_inner()) {
                        let formatted = format_expr(&expr, Some(80));
                        println!("Input:\n{}", source);
                        println!("Output:\n{}", formatted);

                        // Currently it collapses to single line
                        assert_eq!(formatted, "x = [1, 2, 3]");
                    }
                }
            }
        }
    }

    #[test]
    fn test_empty_lines_between_statements() {
        use crate::parser::Rule;

        // Test that empty lines between statements are preserved (up to 2)
        let source = "x = 1\n\ny = 2\n\n\n\nz = 3";
        let pairs = get_pairs(source).unwrap();

        let mut statements_with_positions = Vec::new();

        for pair in pairs {
            if pair.as_rule() == Rule::statement {
                let start_line = pair.as_span().start_pos().line_col().0;
                let end_line = pair.as_span().end_pos().line_col().0;

                if let Some(inner_pair) = pair.into_inner().next() {
                    if let Ok(expr) = pairs_to_expr(inner_pair.into_inner()) {
                        let formatted = format_expr(&expr, Some(80));
                        statements_with_positions.push((formatted, start_line, end_line));
                    }
                }
            }
        }

        println!("Statement positions: {:?}", statements_with_positions);

        // Use the new helper function to join with spacing
        let result = join_statements_with_spacing(&statements_with_positions);
        println!("Output with preserved spacing:\n{}", result);

        // We want to preserve 1 empty line between x and y, and cap at 2 between y and z
        // So it should be:
        // x = 1
        // <blank>
        // y = 2
        // <blank>
        // <blank>
        // z = 3

        assert_eq!(result.lines().count(), 6); // 3 statements + 1 blank + 2 blanks
        assert_eq!(result, "x = 1\n\ny = 2\n\n\nz = 3");
    }

    #[test]
    fn test_assignment_with_long_list_indentation() {
        use crate::parser::Rule;

        // Test that assignments with lists don't have excessive indentation
        let source = "ingredients = [{name: \"sugar\", amount: 1}, {name: \"flour\", amount: 2}]";
        let pairs = get_pairs(source).unwrap();

        for pair in pairs {
            if pair.as_rule() == Rule::statement {
                if let Some(inner_pair) = pair.into_inner().next() {
                    if let Ok(expr) = pairs_to_expr(inner_pair.into_inner()) {
                        let formatted = format_expr(&expr, Some(40));
                        println!("Formatted:\n{}", formatted);

                        // The list items should be indented by 2 spaces, not pushed over by the "ingredients = " prefix
                        assert!(formatted.contains("[\n  {"));
                        assert!(!formatted.contains("            ")); // Should not have excessive indentation
                    }
                }
            }
        }
    }
}