nff/src/main.rs

mod ast;
mod cst;
mod lexer;
mod parser;
mod syntax;

use anyhow::{Context, Result};
use clap::Parser;
use std::fs;
use std::io::{self, Write};
use std::path::Path;
use thiserror::Error;

use crate::cst::CstBuilder;
use crate::lexer::NftablesLexer;
use crate::parser::Parser as NftablesParser;
use crate::syntax::{FormatConfig, IndentStyle, NftablesFormatter};

#[derive(Error, Debug)]
enum FormatterError {
    #[error("File not found: {0}")]
    FileNotFound(String),
    #[error("Invalid file: {0}")]
    InvalidFile(String),
    #[error("Parse error: {0}")]
    ParseError(String),
    #[error("IO error: {0}")]
    Io(#[from] io::Error),
}

#[derive(Parser, Debug)]
#[command(
    name = "nff",
    version = "0.1.0",
    about = "A high-quality nftables formatter and beautifier",
    long_about = "nff (nftables formatter) is a tool for formatting and beautifying nftables configuration files with proper indentation and structure."
)]
struct Args {
    /// nftables config file (e.g: /etc/nftables.conf)
    #[arg(short, long, value_name = "FILE")]
    file: String,

    /// Type of indentation
    #[arg(short, long, default_value = "tabs", value_parser = clap::value_parser!(IndentStyle))]
    indent: IndentStyle,

    /// Output file (writes to stdout if not specified)
    #[arg(short, long, value_name = "FILE")]
    output: Option<String>,

    /// Optimize output by removing excessive empty lines
    #[arg(long)]
    optimize: bool,

    /// Number of spaces per indentation level (only used with --indent=spaces)
    #[arg(long, default_value = "2", value_name = "N")]
    spaces: usize,

    /// Show debug information (tokens, AST, etc.)
    #[arg(long)]
    debug: bool,

    /// Check syntax only, don't format
    #[arg(long)]
    check: bool,
}

fn process_nftables_config(args: Args) -> Result<()> {
    let path = Path::new(&args.file);
    if !path.exists() {
        return Err(FormatterError::FileNotFound(args.file).into());
    }

    if !path.is_file() {
        return Err(FormatterError::InvalidFile("Not a regular file".to_string()).into());
    }

    // Read file contents
    let source = fs::read_to_string(&args.file)
        .with_context(|| format!("Failed to read file: {}", args.file))?;

    // Tokenize
    let mut lexer = NftablesLexer::new(&source);
    let tokens = if args.debug {
        // Use error-recovery tokenization for debug mode
        lexer.tokenize_with_errors()
    } else {
        lexer
            .tokenize()
            .map_err(|e| FormatterError::ParseError(e.to_string()))?
    };

    if args.debug {
        eprintln!("=== TOKENS ===");
        for (i, token) in tokens.iter().enumerate() {
            eprintln!(
                "{:3}: {:?} @ {:?} = '{}'",
                i, token.kind, token.range, token.text
            );
        }
        eprintln!();

        // Build and validate CST
        eprintln!("=== CST ===");
        let cst_tree = CstBuilder::build_tree(&tokens);
        match CstBuilder::validate_tree(&cst_tree) {
            Ok(()) => eprintln!("CST validation passed"),
            Err(e) => eprintln!("CST validation error: {}", e),
        }

        // Also test parse_to_cst
        match CstBuilder::parse_to_cst(&tokens) {
            Ok(_) => eprintln!("CST parsing successful"),
            Err(e) => eprintln!("CST parsing error: {}", e),
        }
        eprintln!();
    }

    // Parse
    let ruleset = if args.debug {
        // Use error-recovery parsing for debug mode
        let (parsed_ruleset, errors) = NftablesParser::parse_with_errors(&source);
        if !errors.is_empty() {
            eprintln!("=== PARSE ERRORS ===");
            for error in &errors {
                eprintln!("Parse error: {}", error);
            }
            eprintln!();
        }
        parsed_ruleset.unwrap_or_else(|| crate::ast::Ruleset::new())
    } else {
        let mut parser = NftablesParser::new(tokens.clone());
        parser
            .parse()
            .map_err(|e| FormatterError::ParseError(e.to_string()))?
    };

    if args.debug {
        eprintln!("=== AST ===");
        eprintln!("{:#?}", ruleset);
        eprintln!();
    }

    if args.check {
        println!("Syntax check passed for: {}", args.file);
        return Ok(());
    }

    // Format
    let config = FormatConfig {
        indent_style: args.indent,
        spaces_per_level: args.spaces,
        optimize: args.optimize,
        max_empty_lines: if args.optimize { 1 } else { 2 },
    };

    let formatter = NftablesFormatter::new(config);
    let formatted_output = formatter.format_ruleset(&ruleset);

    // Write output
    match &args.output {
        Some(output_file) => {
            fs::write(output_file, &formatted_output)
                .with_context(|| format!("Failed to write to output file: {}", output_file))?;
            println!("Formatted output written to: {}", output_file);
        }
        None => {
            io::stdout()
                .write_all(formatted_output.as_bytes())
                .with_context(|| "Failed to write to stdout")?;
        }
    }

    Ok(())
}

fn main() -> Result<()> {
    let args = Args::parse();

    if let Err(e) = process_nftables_config(args) {
        eprintln!("Error: {}", e);

        // Print the error chain
        let mut current = e.source();
        while let Some(cause) = current {
            eprintln!("  Caused by: {}", cause);
            current = cause.source();
        }

        std::process::exit(1);
    }

    Ok(())
}