lux/src/formatter.rs

//! Code formatter for Lux
//!
//! Formats Lux source code according to standard style guidelines.

use crate::ast::{
    BehavioralProperty, BinaryOp, Declaration, EffectDecl, Expr, FunctionDecl, HandlerDecl,
    ImplDecl, ImplMethod, LetDecl, Literal, LiteralKind, Pattern, Program, Statement, TraitDecl,
    TypeDecl, TypeDef, TypeExpr, UnaryOp, VariantFields,
};
use crate::lexer::Lexer;
use crate::parser::Parser;

/// Formatter configuration
#[derive(Debug, Clone)]
#[allow(dead_code)]
pub struct FormatConfig {
    /// Number of spaces for indentation
    pub indent_size: usize,
    /// Maximum line width before wrapping
    pub max_width: usize,
    /// Add trailing commas in multi-line constructs
    pub trailing_commas: bool,
}

impl Default for FormatConfig {
    fn default() -> Self {
        Self {
            indent_size: 4,
            max_width: 100,
            trailing_commas: true,
        }
    }
}

/// Format Lux source code
pub fn format(source: &str, config: &FormatConfig) -> Result<String, String> {
    // Parse the source
    let lexer = Lexer::new(source);
    let tokens = lexer.tokenize().map_err(|e| e.message)?;
    let mut parser = Parser::new(tokens);
    let program = parser.parse_program().map_err(|e| e.message)?;

    // Format the AST
    let mut formatter = Formatter::new(config.clone());
    Ok(formatter.format_program(&program))
}

/// Formatter state
struct Formatter {
    config: FormatConfig,
    output: String,
    indent_level: usize,
}

impl Formatter {
    fn new(config: FormatConfig) -> Self {
        Self {
            config,
            output: String::new(),
            indent_level: 0,
        }
    }

    fn indent(&self) -> String {
        " ".repeat(self.indent_level * self.config.indent_size)
    }

    fn write(&mut self, s: &str) {
        self.output.push_str(s);
    }

    fn writeln(&mut self, s: &str) {
        self.output.push_str(s);
        self.output.push('\n');
    }

    fn newline(&mut self) {
        self.output.push('\n');
    }

    fn format_program(&mut self, program: &Program) -> String {
        let mut first = true;
        for decl in &program.declarations {
            if !first {
                self.newline();
            }
            first = false;
            self.format_declaration(decl);
        }
        // Ensure file ends with newline
        if !self.output.ends_with('\n') {
            self.newline();
        }
        self.output.clone()
    }

    fn format_declaration(&mut self, decl: &Declaration) {
        match decl {
            Declaration::Function(f) => self.format_function(f),
            Declaration::Let(l) => self.format_let(l),
            Declaration::Type(t) => self.format_type_decl(t),
            Declaration::Effect(e) => self.format_effect(e),
            Declaration::Handler(h) => self.format_handler(h),
            Declaration::Trait(t) => self.format_trait(t),
            Declaration::Impl(i) => self.format_impl(i),
        }
    }

    fn format_function(&mut self, func: &FunctionDecl) {
        let indent = self.indent();
        self.write(&indent);
        self.write("fn ");
        self.write(&func.name.name);

        // Type parameters
        if !func.type_params.is_empty() {
            self.write("<");
            self.write(
                &func
                    .type_params
                    .iter()
                    .map(|p| p.name.clone())
                    .collect::<Vec<_>>()
                    .join(", "),
            );
            self.write(">");
        }

        // Parameters
        self.write("(");
        let params: Vec<String> = func
            .params
            .iter()
            .map(|p| format!("{}: {}", p.name.name, self.format_type_expr(&p.typ)))
            .collect();
        self.write(&params.join(", "));
        self.write("): ");

        // Return type
        self.write(&self.format_type_expr(&func.return_type));

        // Effects
        if !func.effects.is_empty() {
            self.write(" with {");
            self.write(
                &func
                    .effects
                    .iter()
                    .map(|e| e.name.clone())
                    .collect::<Vec<_>>()
                    .join(", "),
            );
            self.write("}");
        }

        // Properties
        if !func.properties.is_empty() {
            self.write(" is ");
            self.write(
                &func
                    .properties
                    .iter()
                    .map(|p| self.format_property(p))
                    .collect::<Vec<_>>()
                    .join(", "),
            );
        }

        self.write(" =");

        // Body - handle blocks specially to keep `= {` on same line
        if let Expr::Block { statements, result, .. } = &func.body {
            self.write(" {");
            self.newline();
            self.indent_level += 1;
            for stmt in statements {
                let indent = self.indent();
                match stmt {
                    Statement::Let { name, typ, value, .. } => {
                        let type_str = typ.as_ref()
                            .map(|t| format!(": {}", self.format_type_expr(t)))
                            .unwrap_or_default();
                        self.write(&indent);
                        self.writeln(&format!(
                            "let {}{} = {}",
                            name.name,
                            type_str,
                            self.format_expr(value)
                        ));
                    }
                    Statement::Expr(e) => {
                        self.write(&indent);
                        self.writeln(&self.format_expr(e));
                    }
                }
            }
            self.write(&self.indent());
            self.writeln(&self.format_expr(result));
            self.indent_level -= 1;
            self.write(&self.indent());
            self.writeln("}");
        } else {
            let body_str = self.format_expr(&func.body);
            if body_str.contains('\n') {
                self.newline();
                self.indent_level += 1;
                self.write(&self.indent());
                self.write(&body_str);
                self.indent_level -= 1;
                self.newline();
            } else {
                self.write(" ");
                self.writeln(&body_str);
            }
        }
    }

    fn format_let(&mut self, let_decl: &LetDecl) {
        let indent = self.indent();
        self.write(&indent);
        self.write("let ");
        self.write(&let_decl.name.name);

        if let Some(ref typ) = let_decl.typ {
            self.write(": ");
            self.write(&self.format_type_expr(typ));
        }

        self.write(" = ");
        self.write(&self.format_expr(&let_decl.value));
        self.newline();
    }

    fn format_type_decl(&mut self, type_decl: &TypeDecl) {
        let indent = self.indent();
        self.write(&indent);
        self.write("type ");
        self.write(&type_decl.name.name);

        // Type parameters
        if !type_decl.type_params.is_empty() {
            self.write("<");
            self.write(
                &type_decl
                    .type_params
                    .iter()
                    .map(|p| p.name.clone())
                    .collect::<Vec<_>>()
                    .join(", "),
            );
            self.write(">");
        }

        self.write(" =");

        match &type_decl.definition {
            TypeDef::Alias(t) => {
                self.write(" ");
                self.writeln(&self.format_type_expr(t));
            }
            TypeDef::Record(fields) => {
                self.writeln(" {");
                self.indent_level += 1;
                for field in fields {
                    self.write(&self.indent());
                    self.write(&field.name.name);
                    self.write(": ");
                    self.write(&self.format_type_expr(&field.typ));
                    self.writeln(",");
                }
                self.indent_level -= 1;
                self.write(&self.indent());
                self.writeln("}");
            }
            TypeDef::Enum(variants) => {
                self.newline();
                self.indent_level += 1;
                for variant in variants {
                    self.write(&self.indent());
                    self.write("| ");
                    self.write(&variant.name.name);
                    match &variant.fields {
                        VariantFields::Unit => {}
                        VariantFields::Tuple(types) => {
                            self.write("(");
                            self.write(
                                &types
                                    .iter()
                                    .map(|t| self.format_type_expr(t))
                                    .collect::<Vec<_>>()
                                    .join(", "),
                            );
                            self.write(")");
                        }
                        VariantFields::Record(fields) => {
                            self.write(" { ");
                            self.write(
                                &fields
                                    .iter()
                                    .map(|f| format!("{}: {}", f.name.name, self.format_type_expr(&f.typ)))
                                    .collect::<Vec<_>>()
                                    .join(", "),
                            );
                            self.write(" }");
                        }
                    }
                    self.newline();
                }
                self.indent_level -= 1;
            }
        }
    }

    fn format_effect(&mut self, effect: &EffectDecl) {
        let indent = self.indent();
        self.write(&indent);
        self.write("effect ");
        self.write(&effect.name.name);

        if !effect.type_params.is_empty() {
            self.write("<");
            self.write(
                &effect
                    .type_params
                    .iter()
                    .map(|p| p.name.clone())
                    .collect::<Vec<_>>()
                    .join(", "),
            );
            self.write(">");
        }

        self.writeln(" {");
        self.indent_level += 1;

        for op in &effect.operations {
            self.write(&self.indent());
            self.write("fn ");
            self.write(&op.name.name);
            self.write("(");
            let params: Vec<String> = op
                .params
                .iter()
                .map(|p| format!("{}: {}", p.name.name, self.format_type_expr(&p.typ)))
                .collect();
            self.write(&params.join(", "));
            self.write("): ");
            self.writeln(&self.format_type_expr(&op.return_type));
        }

        self.indent_level -= 1;
        self.write(&self.indent());
        self.writeln("}");
    }

    fn format_handler(&mut self, handler: &HandlerDecl) {
        let indent = self.indent();
        self.write(&indent);
        self.write("handler ");
        self.write(&handler.name.name);
        self.write(": ");
        self.write(&handler.effect.name);
        self.writeln(" {");
        self.indent_level += 1;

        for impl_ in &handler.implementations {
            self.write(&self.indent());
            self.write("fn ");
            self.write(&impl_.op_name.name);
            self.write("(");
            self.write(
                &impl_
                    .params
                    .iter()
                    .map(|p| p.name.clone())
                    .collect::<Vec<_>>()
                    .join(", "),
            );
            self.write(") = ");

            let body_str = self.format_expr(&impl_.body);
            if body_str.contains('\n') {
                self.newline();
                self.indent_level += 1;
                self.write(&self.indent());
                self.write(&body_str);
                self.indent_level -= 1;
            } else {
                self.write(&body_str);
            }
            self.newline();
        }

        self.indent_level -= 1;
        self.write(&self.indent());
        self.writeln("}");
    }

    fn format_trait(&mut self, trait_decl: &TraitDecl) {
        let indent = self.indent();
        self.write(&indent);
        self.write("trait ");
        self.write(&trait_decl.name.name);

        if !trait_decl.type_params.is_empty() {
            self.write("<");
            self.write(
                &trait_decl
                    .type_params
                    .iter()
                    .map(|p| p.name.clone())
                    .collect::<Vec<_>>()
                    .join(", "),
            );
            self.write(">");
        }

        self.writeln(" {");
        self.indent_level += 1;

        for method in &trait_decl.methods {
            self.write(&self.indent());
            self.write("fn ");
            self.write(&method.name.name);
            self.write("(");
            let params: Vec<String> = method
                .params
                .iter()
                .map(|p| format!("{}: {}", p.name.name, self.format_type_expr(&p.typ)))
                .collect();
            self.write(&params.join(", "));
            self.write("): ");
            self.writeln(&self.format_type_expr(&method.return_type));
        }

        self.indent_level -= 1;
        self.write(&self.indent());
        self.writeln("}");
    }

    fn format_impl(&mut self, impl_decl: &ImplDecl) {
        let indent = self.indent();
        self.write(&indent);
        self.write("impl ");
        self.write(&impl_decl.trait_name.name);
        self.write(" for ");
        self.write(&self.format_type_expr(&impl_decl.target_type));
        self.writeln(" {");
        self.indent_level += 1;

        for method in &impl_decl.methods {
            self.format_impl_method(method);
        }

        self.indent_level -= 1;
        self.write(&self.indent());
        self.writeln("}");
    }

    fn format_impl_method(&mut self, method: &ImplMethod) {
        let indent = self.indent();
        self.write(&indent);
        self.write("fn ");
        self.write(&method.name.name);
        self.write("(");
        let params: Vec<String> = method
            .params
            .iter()
            .map(|p| format!("{}: {}", p.name.name, self.format_type_expr(&p.typ)))
            .collect();
        self.write(&params.join(", "));
        self.write(")");

        if let Some(ref ret) = method.return_type {
            self.write(": ");
            self.write(&self.format_type_expr(ret));
        }

        self.write(" = ");

        let body_str = self.format_expr(&method.body);
        if self.is_block_expr(&method.body) || body_str.contains('\n') {
            self.newline();
            self.indent_level += 1;
            self.write(&self.indent());
            self.write(&body_str);
            self.indent_level -= 1;
        } else {
            self.write(&body_str);
        }
        self.newline();
    }

    fn format_property(&self, prop: &BehavioralProperty) -> String {
        match prop {
            BehavioralProperty::Pure => "pure".to_string(),
            BehavioralProperty::Total => "total".to_string(),
            BehavioralProperty::Idempotent => "idempotent".to_string(),
            BehavioralProperty::Deterministic => "deterministic".to_string(),
            BehavioralProperty::Commutative => "commutative".to_string(),
        }
    }

    fn format_type_expr(&self, typ: &TypeExpr) -> String {
        match typ {
            TypeExpr::Named(ident) => ident.name.clone(),
            TypeExpr::App(base, args) => {
                format!(
                    "{}<{}>",
                    self.format_type_expr(base),
                    args.iter()
                        .map(|a| self.format_type_expr(a))
                        .collect::<Vec<_>>()
                        .join(", ")
                )
            }
            TypeExpr::Function { params, return_type, effects } => {
                let mut s = format!(
                    "fn({}): {}",
                    params
                        .iter()
                        .map(|p| self.format_type_expr(p))
                        .collect::<Vec<_>>()
                        .join(", "),
                    self.format_type_expr(return_type)
                );
                if !effects.is_empty() {
                    s.push_str(" with {");
                    s.push_str(
                        &effects
                            .iter()
                            .map(|e| e.name.clone())
                            .collect::<Vec<_>>()
                            .join(", "),
                    );
                    s.push('}');
                }
                s
            }
            TypeExpr::Tuple(types) => {
                format!(
                    "({})",
                    types
                        .iter()
                        .map(|t| self.format_type_expr(t))
                        .collect::<Vec<_>>()
                        .join(", ")
                )
            }
            TypeExpr::Record(fields) => {
                format!(
                    "{{ {} }}",
                    fields
                        .iter()
                        .map(|f| format!("{}: {}", f.name.name, self.format_type_expr(&f.typ)))
                        .collect::<Vec<_>>()
                        .join(", ")
                )
            }
            TypeExpr::Unit => "Unit".to_string(),
            TypeExpr::Versioned { base, constraint } => {
                let mut s = self.format_type_expr(base);
                s.push_str(&format!(" {}", constraint));
                s
            }
        }
    }

    fn is_block_expr(&self, expr: &Expr) -> bool {
        matches!(expr, Expr::Block { .. } | Expr::Match { .. })
    }

    fn format_expr(&self, expr: &Expr) -> String {
        match expr {
            Expr::Literal(lit) => self.format_literal(lit),
            Expr::Var(ident) => ident.name.clone(),
            Expr::BinaryOp { left, op, right, .. } => {
                format!(
                    "{} {} {}",
                    self.format_expr(left),
                    self.format_binop(op),
                    self.format_expr(right)
                )
            }
            Expr::UnaryOp { op, operand, .. } => {
                format!("{}{}", self.format_unaryop(op), self.format_expr(operand))
            }
            Expr::Call { func, args, .. } => {
                format!(
                    "{}({})",
                    self.format_expr(func),
                    args.iter()
                        .map(|a| self.format_expr(a))
                        .collect::<Vec<_>>()
                        .join(", ")
                )
            }
            Expr::Field { object, field, .. } => {
                format!("{}.{}", self.format_expr(object), field.name)
            }
            Expr::TupleIndex { object, index, .. } => {
                format!("{}.{}", self.format_expr(object), index)
            }
            Expr::If { condition, then_branch, else_branch, .. } => {
                format!(
                    "if {} then {} else {}",
                    self.format_expr(condition),
                    self.format_expr(then_branch),
                    self.format_expr(else_branch)
                )
            }
            Expr::Match { scrutinee, arms, .. } => {
                let mut s = format!("match {} {{\n", self.format_expr(scrutinee));
                for arm in arms {
                    s.push_str(&format!(
                        "    {} => {},\n",
                        self.format_pattern(&arm.pattern),
                        self.format_expr(&arm.body)
                    ));
                }
                s.push('}');
                s
            }
            Expr::Block { statements, result, .. } => {
                let mut s = String::from("{\n");
                for stmt in statements {
                    match stmt {
                        Statement::Let { name, typ, value, .. } => {
                            let type_str = typ.as_ref()
                                .map(|t| format!(": {}", self.format_type_expr(t)))
                                .unwrap_or_default();
                            s.push_str(&format!(
                                "    let {}{} = {}\n",
                                name.name,
                                type_str,
                                self.format_expr(value)
                            ));
                        }
                        Statement::Expr(e) => {
                            s.push_str(&format!("    {}\n", self.format_expr(e)));
                        }
                    }
                }
                s.push_str(&format!("    {}\n", self.format_expr(result)));
                s.push('}');
                s
            }
            Expr::Lambda { params, body, return_type, .. } => {
                let params_str: Vec<String> = params
                    .iter()
                    .map(|p| format!("{}: {}", p.name.name, self.format_type_expr(&p.typ)))
                    .collect();
                let ret_str = return_type
                    .as_ref()
                    .map(|t| format!(": {}", self.format_type_expr(t)))
                    .unwrap_or_default();
                format!("fn({}){} => {}", params_str.join(", "), ret_str, self.format_expr(body))
            }
            Expr::Let { name, value, body, typ, .. } => {
                let type_str = typ.as_ref()
                    .map(|t| format!(": {}", self.format_type_expr(t)))
                    .unwrap_or_default();
                format!(
                    "let {}{} = {}; {}",
                    name.name,
                    type_str,
                    self.format_expr(value),
                    self.format_expr(body)
                )
            }
            Expr::List { elements, .. } => {
                format!(
                    "[{}]",
                    elements
                        .iter()
                        .map(|e| self.format_expr(e))
                        .collect::<Vec<_>>()
                        .join(", ")
                )
            }
            Expr::Tuple { elements, .. } => {
                format!(
                    "({})",
                    elements
                        .iter()
                        .map(|e| self.format_expr(e))
                        .collect::<Vec<_>>()
                        .join(", ")
                )
            }
            Expr::Record {
                spread, fields, ..
            } => {
                let mut parts = Vec::new();
                if let Some(spread_expr) = spread {
                    parts.push(format!("...{}", self.format_expr(spread_expr)));
                }
                for (name, val) in fields {
                    parts.push(format!("{}: {}", name.name, self.format_expr(val)));
                }
                format!("{{ {} }}", parts.join(", "))
            }
            Expr::EffectOp { effect, operation, args, .. } => {
                format!(
                    "{}.{}({})",
                    effect.name,
                    operation.name,
                    args.iter()
                        .map(|a| self.format_expr(a))
                        .collect::<Vec<_>>()
                        .join(", ")
                )
            }
            Expr::Run { expr, handlers, .. } => {
                if handlers.is_empty() {
                    format!("run {} with {{}}", self.format_expr(expr))
                } else {
                    let mut s = format!("run {} with {{\n", self.format_expr(expr));
                    for (effect, handler) in handlers {
                        s.push_str(&format!("    {} = {},\n", effect.name, self.format_expr(handler)));
                    }
                    s.push('}');
                    s
                }
            }
            Expr::Resume { value, .. } => {
                format!("resume({})", self.format_expr(value))
            }
        }
    }

    fn format_literal(&self, lit: &Literal) -> String {
        match &lit.kind {
            LiteralKind::Int(n) => n.to_string(),
            LiteralKind::Float(f) => format!("{}", f),
            LiteralKind::String(s) => format!("\"{}\"", s.replace('\\', "\\\\").replace('"', "\\\"").replace('{', "\\{").replace('}', "\\}")),
            LiteralKind::Char(c) => format!("'{}'", c),
            LiteralKind::Bool(b) => b.to_string(),
            LiteralKind::Unit => "()".to_string(),
        }
    }

    fn format_binop(&self, op: &BinaryOp) -> &'static str {
        match op {
            BinaryOp::Add => "+",
            BinaryOp::Sub => "-",
            BinaryOp::Mul => "*",
            BinaryOp::Div => "/",
            BinaryOp::Mod => "%",
            BinaryOp::Eq => "==",
            BinaryOp::Ne => "!=",
            BinaryOp::Lt => "<",
            BinaryOp::Le => "<=",
            BinaryOp::Gt => ">",
            BinaryOp::Ge => ">=",
            BinaryOp::And => "&&",
            BinaryOp::Or => "||",
            BinaryOp::Concat => "++",
            BinaryOp::Pipe => "|>",
        }
    }

    fn format_unaryop(&self, op: &UnaryOp) -> &'static str {
        match op {
            UnaryOp::Neg => "-",
            UnaryOp::Not => "!",
        }
    }

    fn format_pattern(&self, pattern: &Pattern) -> String {
        match pattern {
            Pattern::Wildcard(_) => "_".to_string(),
            Pattern::Var(ident) => ident.name.clone(),
            Pattern::Literal(lit) => self.format_literal(lit),
            Pattern::Constructor { name, fields, .. } => {
                if fields.is_empty() {
                    name.name.clone()
                } else {
                    format!(
                        "{}({})",
                        name.name,
                        fields
                            .iter()
                            .map(|f| self.format_pattern(f))
                            .collect::<Vec<_>>()
                            .join(", ")
                    )
                }
            }
            Pattern::Tuple { elements, .. } => {
                format!(
                    "({})",
                    elements
                        .iter()
                        .map(|e| self.format_pattern(e))
                        .collect::<Vec<_>>()
                        .join(", ")
                )
            }
            Pattern::Record { fields, .. } => {
                format!(
                    "{{ {} }}",
                    fields
                        .iter()
                        .map(|(name, pat)| {
                            format!("{}: {}", name.name, self.format_pattern(pat))
                        })
                        .collect::<Vec<_>>()
                        .join(", ")
                )
            }
        }
    }
}

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

    #[test]
    fn test_format_simple_function() {
        let source = "fn   add( x:Int,y:Int ):Int=x+y";
        let result = format(source, &FormatConfig::default()).unwrap();
        assert!(result.contains("fn add(x: Int, y: Int): Int = x + y"));
    }

    #[test]
    fn test_format_let() {
        let source = "let   x=42";
        let result = format(source, &FormatConfig::default()).unwrap();
        assert!(result.contains("let x = 42"));
    }
}