Chapter 7: The Rule Engine: Linting and Deterministic Fixing

// src/rules/mod.rs

use crate::ast::MermaidAst;
use crate::diagnostics::{Diagnostic, DiagnosticLevel, Diagnostics};
use std::fmt::Debug;

// Re-export specific rules for easier access
pub mod missing_graph_decl;
pub mod arrow_normalization;

/// Defines the operation mode for the rule engine.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum FixMode {
    /// Only report diagnostics, no modifications.
    LintOnly,
    /// Apply safe, reversible fixes.
    ApplyFixes,
    /// Apply only guaranteed safe and minimal fixes, fail on any ambiguity.
    StrictFixes,
}

/// The trait that all linting and fixing rules must implement.
///
/// Rules are designed to be deterministic and operate on the AST.
/// They can produce diagnostics and, in fixing modes, modify the AST.
pub trait Rule: Debug + Send + Sync {
    /// Returns the unique name of the rule.
    fn name(&self) -> &'static str;

    /// Checks the AST for issues and adds diagnostics.
    /// This method should NOT modify the AST.
    fn check(&self, ast: &MermaidAst, diagnostics: &mut Diagnostics);

    /// Attempts to apply a fix to the AST.
    /// Returns `true` if any modification was made, `false` otherwise.
    /// This method should only be called in `ApplyFixes` or `StrictFixes` modes.
    /// It should also add diagnostics for the applied fix.
    fn apply_fix(&self, ast: &mut MermaidAst, diagnostics: &mut Diagnostics) -> bool;

    /// Indicates which fix modes this rule supports.
    /// By default, rules support `ApplyFixes`. Override for `StrictFixes`.
    fn supports_fix_mode(&self, mode: FixMode) -> bool {
        match mode {
            FixMode::LintOnly => true, // All rules can lint
            FixMode::ApplyFixes => true,
            FixMode::StrictFixes => false, // By default, rules are not strict-fix safe
        }
    }
}

/// The Rule Engine orchestrates the application of various rules.
#[derive(Default)]
pub struct RuleEngine {
    rules: Vec<Box<dyn Rule>>,
}

impl RuleEngine {
    /// Creates a new RuleEngine with no rules registered.
    pub fn new() -> Self {
        Self { rules: Vec::new() }
    }

    /// Registers a new rule with the engine.
    pub fn register_rule(&mut self, rule: Box<dyn Rule>) {
        self.rules.push(rule);
    }

    /// Runs the rule engine in lint-only mode.
    /// It checks the AST against all registered rules and collects diagnostics.
    pub fn run_lint(&self, ast: &MermaidAst) -> Diagnostics {
        let mut diagnostics = Diagnostics::new();
        for rule in &self.rules {
            rule.check(ast, &mut diagnostics);
        }
        diagnostics
    }

    /// Runs the rule engine to apply fixes to the AST.
    ///
    /// It performs multiple passes until no more fixes can be applied by any rule
    /// or a maximum number of passes is reached to prevent infinite loops.
    ///
    /// Returns the collected diagnostics, including those for applied fixes.
    pub fn run_fix(&self, ast: &mut MermaidAst, fix_mode: FixMode) -> Diagnostics {
        let mut diagnostics = Diagnostics::new();
        if fix_mode == FixMode::LintOnly {
            // Fallback to lint-only if fix mode is incorrectly passed
            return self.run_lint(ast);
        }

        const MAX_FIX_PASSES: usize = 10; // Prevent infinite loops
        let mut changed_in_last_pass = true;
        let mut pass_count = 0;

        while changed_in_last_pass && pass_count < MAX_FIX_PASSES {
            changed_in_last_pass = false;
            pass_count += 1;

            // First, run a check pass to gather all current diagnostics
            // This is important because applying a fix might introduce new issues
            // or resolve existing ones that other rules might detect.
            // We clear and re-populate diagnostics in each pass.
            let mut current_pass_diagnostics = Diagnostics::new();
            for rule in &self.rules {
                rule.check(ast, &mut current_pass_diagnostics);
            }
            // Merge current pass diagnostics into the main diagnostics object
            diagnostics.extend(current_pass_diagnostics);

            // Then, attempt to apply fixes
            for rule in &self.rules {
                if rule.supports_fix_mode(fix_mode) {
                    if rule.apply_fix(ast, &mut diagnostics) {
                        changed_in_last_pass = true;
                    }
                }
            }
        }

        if changed_in_last_pass {
            diagnostics.add_diagnostic(Diagnostic {
                level: DiagnosticLevel::Warning,
                code: "R000_FIX_LIMIT".to_string(),
                message: format!("Reached maximum fix passes ({}). Some issues might remain. Please re-run the tool.", MAX_FIX_PASSES),
                span: None,
                help: Some("This usually indicates a complex interaction between rules or a rule that isn't fully idempotent. Consider manual inspection or re-running the tool.".to_string()),
            });
        }
        diagnostics
    }

    /// Helper to register all default rules.
    pub fn register_default_rules(&mut self) {
        self.register_rule(Box::new(missing_graph_decl::MissingGraphDeclarationRule));
        self.register_rule(Box::new(arrow_normalization::ArrowNormalizationRule));
        // Register other rules here as they are implemented
    }
}

// src/rules/missing_graph_decl.rs

use super::{Rule, FixMode};
use crate::ast::{MermaidAst, Statement, GraphDeclaration, GraphOrientation};
use crate::diagnostics::{Diagnostic, DiagnosticLevel, Diagnostics};
use crate::span::Span; // Assuming Span is defined and used for locations

/// Rule to ensure that a Mermaid diagram starts with a graph or flowchart declaration.
/// If missing, it adds a default 'graph TD' declaration.
#[derive(Debug)]
pub struct MissingGraphDeclarationRule;

impl Rule for MissingGraphDeclarationRule {
    fn name(&self) -> &'static str {
        "missing-graph-declaration"
    }

    fn check(&self, ast: &MermaidAst, diagnostics: &mut Diagnostics) {
        if !ast.statements.iter().any(|stmt| matches!(stmt, Statement::GraphDeclaration(_))) {
            diagnostics.add_diagnostic(Diagnostic {
                level: DiagnosticLevel::Warning,
                code: "R701_MISSING_GRAPH_DECL".to_string(),
                message: "Mermaid diagram is missing a 'graph' or 'flowchart' declaration.".to_string(),
                span: ast.statements.first().map(|s| s.span()), // Point to the start of the file
                help: Some("Consider adding 'graph TD' or 'flowchart TD' at the beginning of your diagram for clarity and strict compliance.".to_string()),
            });
        }
    }

    fn apply_fix(&self, ast: &mut MermaidAst, diagnostics: &mut Diagnostics) -> bool {
        if !ast.statements.iter().any(|stmt| matches!(stmt, Statement::GraphDeclaration(_))) {
            // Create a default 'graph TD' declaration.
            // We'll use a dummy span as this is an inserted node.
            // For a real-world scenario, you might want to infer a more accurate span.
            let default_decl = Statement::GraphDeclaration(GraphDeclaration {
                kind: "graph".to_string(), // Use "graph" as the default
                orientation: Some(GraphOrientation::TD),
                span: Span::new(0, 0), // Dummy span, as it's inserted
            });

            // Insert at the beginning of the statements
            ast.statements.insert(0, default_decl);

            diagnostics.add_diagnostic(Diagnostic {
                level: DiagnosticLevel::Note,
                code: "R701_MISSING_GRAPH_DECL_FIXED".to_string(),
                message: "Added 'graph TD' declaration to the beginning of the diagram.".to_string(),
                span: Some(Span::new(0, 0)), // Point to the very beginning
                help: None,
            });
            true // Modification made
        } else {
            false // No modification needed
        }
    }

    fn supports_fix_mode(&self, mode: FixMode) -> bool {
        // This is a safe fix, so it can be applied in both ApplyFixes and StrictFixes modes.
        matches!(mode, FixMode::ApplyFixes | FixMode::StrictFixes)
    }
}

// src/rules/arrow_normalization.rs

use super::{Rule, FixMode};
use crate::ast::{MermaidAst, Statement, Edge, EdgeArrow, EdgeLine, EdgeStyle};
use crate::diagnostics::{Diagnostic, DiagnosticLevel, Diagnostics};
use crate::span::Span;

/// Rule to normalize various arrow syntaxes to a standard form (e.g., `-->`).
/// This rule primarily targets simple directed arrows.
#[derive(Debug)]
pub struct ArrowNormalizationRule;

impl Rule for ArrowNormalizationRule {
    fn name(&self) -> &'static str {
        "arrow-normalization"
    }

    fn check(&self, ast: &MermaidAst, diagnostics: &mut Diagnostics) {
        // This rule primarily identifies non-standard arrows that _could_ be normalized.
        // It iterates through the AST to find edges and checks their arrow styles.
        for statement in &ast.statements {
            if let Statement::Edge(edge) = statement {
                if let Some(ref edge_style) = edge.style {
                    // Check for non-standard directed arrows that can be normalized to '-->'
                    let is_non_standard_arrow = matches!(
                        (edge_style.line, edge_style.arrow_start, edge_style.arrow_end),
                        (EdgeLine::Solid, EdgeArrow::None, EdgeArrow::Arrow) | // e.g., "---"
                        (EdgeLine::Dotted, EdgeArrow::None, EdgeArrow::Arrow) // e.g., "-.->"
                    );

                    if is_non_standard_arrow {
                        diagnostics.add_diagnostic(Diagnostic {
                            level: DiagnosticLevel::Warning,
                            code: "R702_NON_STANDARD_ARROW".to_string(),
                            message: format!("Non-standard arrow syntax detected: '{}'.", edge.raw_arrow()),
                            span: Some(edge.span()),
                            help: Some("Consider normalizing to '-->' for consistency.".to_string()),
                        });
                    }
                }
            }
        }
    }

    fn apply_fix(&self, ast: &mut MermaidAst, diagnostics: &mut Diagnostics) -> bool {
        let mut changed = false;
        // We need to traverse the AST mutably. This can be complex for deeply nested structures.
        // For simplicity here, we iterate top-level statements. A full implementation
        // would likely use an AST visitor pattern to handle all levels.
        for statement in &mut ast.statements {
            if let Statement::Edge(edge) = statement {
                if let Some(ref mut edge_style) = edge.style {
                    // Apply fix for non-standard directed arrows to '-->'
                    let should_normalize = matches!(
                        (edge_style.line, edge_style.arrow_start, edge_style.arrow_end),
                        (EdgeLine::Solid, EdgeArrow::None, EdgeArrow::Arrow) | // e.g., "---"
                        (EdgeLine::Dotted, EdgeArrow::None, EdgeArrow::Arrow) // e.g., "-.->"
                    );

                    if should_normalize {
                        // Normalize to '-->'
                        edge_style.line = EdgeLine::Solid;
                        edge_style.arrow_start = EdgeArrow::None;
                        edge_style.arrow_end = EdgeArrow::Arrow;

                        diagnostics.add_diagnostic(Diagnostic {
                            level: DiagnosticLevel::Note,
                            code: "R702_ARROW_NORMALIZED".to_string(),
                            message: format!("Normalized arrow syntax to '-->'. Original was: '{}'", edge.raw_arrow()),
                            span: Some(edge.span()),
                            help: None,
                        });
                        changed = true;
                    }
                }
            }
        }
        changed
    }

    fn supports_fix_mode(&self, mode: FixMode) -> bool {
        // This fix is generally safe, but normalizing all arrows might not be
        // desired in all 'strict' contexts if the original syntax was valid.
        // For now, we'll keep it to ApplyFixes. A more advanced rule might
        // distinguish between strictly equivalent forms and stylistic choices.
        matches!(mode, FixMode::ApplyFixes)
    }
}

# Cargo.toml
[dependencies]
clap = { version = "4", features = ["derive"] }
# ... other dependencies

// src/main.rs

use clap::{Parser, Subcommand};
use std::fs;
use std::path::PathBuf;

// Import our modules
mod lexer;
mod parser;
mod ast;
mod diagnostics;
mod validator;
mod rules; // New!
mod span;

use lexer::Lexer;
use parser::Parser as MermaidParser; // Alias to avoid conflict with clap::Parser
use validator::AstValidator;
use diagnostics::DiagnosticLevel;
use rules::{RuleEngine, FixMode}; // New!

/// A strict linter and formatter for Mermaid diagrams.
#[derive(Parser, Debug)]
#[command(author, version, about = "Strict linter and formatter for Mermaid diagrams", long_about = None)]
struct Cli {
    /// Path to the Mermaid file to process
    #[arg(name = "FILE", help = "Path to the Mermaid file to process")]
    file: PathBuf,

    #[command(subcommand)]
    command: Commands,
}

#[derive(Subcommand, Debug)]
enum Commands {
    /// Checks the Mermaid file for errors and warnings (lint mode).
    Lint {
        /// Exit with a non-zero code if any errors or warnings are found.
        #[arg(long)]
        strict_exit: bool,
    },
    /// Attempts to fix common Mermaid syntax issues.
    Fix {
        /// Overwrite the original file with the fixed content.
        #[arg(long)]
        write: bool,
        /// Apply only guaranteed safe and minimal fixes, fail on any ambiguity.
        #[arg(long)]
        strict: bool,
    },
    // Future: Add `format` command here.
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let cli = Cli::parse();
    let source_code = fs::read_to_string(&cli.file)?;

    let mut diagnostics = diagnostics::Diagnostics::new();

    // 1. Lexing
    let lexer = Lexer::new(&source_code);
    let tokens = lexer.collect::<Vec<_>>();
    // Note: Lexer errors are typically collected during tokenization if designed that way.
    // For now, we assume a successful tokenization or handle basic errors in parser.

    // 2. Parsing
    let mut parser = MermaidParser::new(tokens);
    let mut ast = match parser.parse() {
        Ok(ast) => ast,
        Err(parse_diagnostics) => {
            diagnostics.extend(parse_diagnostics);
            diagnostics.print_diagnostics(&source_code);
            if diagnostics.has_errors() {
                eprintln!("Error: Parsing failed. Exiting.");
                std::process::exit(1);
            }
            // If parser returns diagnostics but no AST, we can't proceed.
            // For now, assume it returns an AST even with errors for further processing.
            // A robust parser would return Result<Option<Ast>, Diagnostics>
            // For this example, let's mock an empty AST if parsing fails completely
            return Err("Parsing failed and produced no AST. Cannot continue.".into());
        }
    };
    diagnostics.extend(parser.into_diagnostics()); // Collect parser-specific diagnostics

    // 3. Semantic Validation (from previous chapter)
    let validator = AstValidator::new();
    validator.validate(&ast, &mut diagnostics);

    // 4. Rule Engine (Linting and Fixing) - NEW
    let mut rule_engine = RuleEngine::new();
    rule_engine.register_default_rules(); // Register all our default rules

    match &cli.command {
        Commands::Lint { strict_exit } => {
            let lint_diagnostics = rule_engine.run_lint(&ast);
            diagnostics.extend(lint_diagnostics);

            diagnostics.print_diagnostics(&source_code);

            if *strict_exit && diagnostics.has_errors_or_warnings() {
                eprintln!("Linting failed with errors or warnings. Exiting strictly.");
                std::process::exit(1);
            } else if diagnostics.has_errors() {
                eprintln!("Linting failed with errors. Exiting.");
                std::process::exit(1);
            }
        }
        Commands::Fix { write, strict } => {
            let fix_mode = if *strict { FixMode::StrictFixes } else { FixMode::ApplyFixes };
            let fix_diagnostics = rule_engine.run_fix(&mut ast, fix_mode);
            diagnostics.extend(fix_diagnostics);

            // Re-validate after fixing to catch any new issues introduced by fixes (shouldn't happen with safe fixes)
            let mut post_fix_diagnostics = diagnostics::Diagnostics::new();
            validator.validate(&ast, &mut post_fix_diagnostics);
            diagnostics.extend(post_fix_diagnostics);

            diagnostics.print_diagnostics(&source_code);

            if diagnostics.has_errors() {
                eprintln!("Error: Fixes were applied, but new errors were detected or existing errors could not be resolved. Exiting.");
                std::process::exit(1);
            }

            // If fixes were applied and no errors, output the modified AST
            if diagnostics.has_fix_notes() || diagnostics.has_level(DiagnosticLevel::Note) { // Check for notes which indicate fixes
                println!("\n--- Fixed Mermaid Code ---");
                // TODO: For now, we'll just print a debug representation or a simplified string.
                // A proper formatter (Chapter 8) will convert AST back to clean Mermaid string.
                // For now, let's just show a placeholder or a simplistic AST to string.
                // This will be replaced by the Formatter in the next chapter.
                println!("{:?}", ast); // Placeholder
                println!("--------------------------");

                if *write {
                    // This is where the formatter (Chapter 8) would convert `ast` back to a string
                    // For now, we'll use a crude representation or just indicate overwrite.
                    eprintln!("`--write` option specified. Original file would be overwritten with formatted content.");
                    eprintln!("(Actual file overwrite is pending implementation of the Formatter in Chapter 8.)");
                    // fs::write(&cli.file, ast_to_string_representation)?; // This will be the actual call
                }
            } else {
                println!("\nNo fixes applied or no issues found that could be fixed.");
            }
        }
    }

    Ok(())
}

// src/diagnostics.rs (Relevant parts, assuming previous implementation)

#[derive(Debug, Clone, PartialEq, Eq, Copy)]
pub enum DiagnosticLevel {
    Error,
    Warning,
    Note, // Added this
    Help, // Used as a sub-message, not a primary level
}

// ... other Diagnostic struct and methods ...

impl Diagnostics {
    // ... existing methods ...

    pub fn has_errors_or_warnings(&self) -> bool {
        self.diagnostics.iter().any(|d| d.level == DiagnosticLevel::Error || d.level == DiagnosticLevel::Warning)
    }

    pub fn has_errors(&self) -> bool {
        self.diagnostics.iter().any(|d| d.level == DiagnosticLevel::Error)
    }

    pub fn has_warnings(&self) -> bool {
        self.diagnostics.iter().any(|d| d.level == DiagnosticLevel::Warning)
    }

    pub fn has_fix_notes(&self) -> bool {
        self.diagnostics.iter().any(|d| d.level == DiagnosticLevel::Note && d.code.contains("FIXED"))
    }

    pub fn has_level(&self, level: DiagnosticLevel) -> bool {
        self.diagnostics.iter().any(|d| d.level == level)
    }

    pub fn print_diagnostics(&self, source_code: &str) {
        use colored::Colorize; // Assuming you have 'colored' crate for nice output

        for diag in self.diagnostics.iter().filter(|d| d.level != DiagnosticLevel::Help) { // Don't print Help as primary
            let level_str = match diag.level {
                DiagnosticLevel::Error => "error".red().bold(),
                DiagnosticLevel::Warning => "warning".yellow().bold(),
                DiagnosticLevel::Note => "note".blue().bold(), // Handle Note
                DiagnosticLevel::Help => continue, // Should not be printed as primary
            };

            let code_str = diag.code.bright_black();

            eprintln!("{}: {} [{}]", level_str, diag.message, code_str);

            if let Some(span) = diag.span {
                let lines: Vec<&str> = source_code.lines().collect();
                if let Some(line_num) = span.start_line(source_code) {
                    eprintln!("     {} {}:{}:{}", "-->".blue().bold(), cli_args.file.display(), line_num + 1, span.start_col(source_code) + 1);

                    // Print context line
                    if let Some(line_content) = lines.get(line_num) {
                        eprintln!(" {} | {}", (line_num + 1).to_string().blue(), line_content);
                        // Highlight the exact span
                        eprintln!("   {} {}{}", " ".blue(), " ".repeat(span.start_col(source_code)), "^".repeat(span.len()).green().bold());
                    }
                }
            }

            if let Some(ref help) = diag.help {
                eprintln!("     {} {}", "help:".cyan().bold(), help);
            }
            eprintln!(); // Blank line for separation
        }
    }
}

// src/diagnostics.rs (Corrected print_diagnostics signature)

use colored::Colorize;
use std::path::Path;

// ... DiagnosticLevel and Diagnostic struct ...

impl Diagnostics {
    // ... existing methods ...

    pub fn print_diagnostics(&self, source_code: &str, file_path: &Path) { // Added file_path
        for diag in self.diagnostics.iter().filter(|d| d.level != DiagnosticLevel::Help) {
            let level_str = match diag.level {
                DiagnosticLevel::Error => "error".red().bold(),
                DiagnosticLevel::Warning => "warning".yellow().bold(),
                DiagnosticLevel::Note => "note".blue().bold(),
                DiagnosticLevel::Help => continue,
            };

            let code_str = diag.code.bright_black();

            eprintln!("{}: {} [{}]", level_str, diag.message, code_str);

            if let Some(span) = diag.span {
                let lines: Vec<&str> = source_code.lines().collect();
                if let Some(line_num) = span.start_line(source_code) {
                    eprintln!("     {} {}:{}:{}", "-->".blue().bold(), file_path.display(), line_num + 1, span.start_col(source_code) + 1);

                    if let Some(line_content) = lines.get(line_num) {
                        eprintln!(" {} | {}", (line_num + 1).to_string().blue(), line_content);
                        eprintln!("   {} {}{}", " ".blue(), " ".repeat(span.start_col(source_code)), "^".repeat(span.len()).green().bold());
                    }
                }
            }

            if let Some(ref help) = diag.help {
                eprintln!("     {} {}", "help:".cyan().bold(), help);
            }
            eprintln!();
        }
    }
}

// src/main.rs (Relevant part in main function)
            diagnostics.print_diagnostics(&source_code, &cli.file);