RustAst/src/bin/ast.rs

use clap::Parser;
use myc::ast::compiler::emitter;
use myc::ast::environment::Environment;
use myc::ast::parser::Parser as MycParser;
use myc::utils::tester;
use std::fs;
use std::io::{self, IsTerminal, Read};
use std::path::PathBuf;

#[derive(Parser)]
#[command(author, version, about = "MYC AST Compiler & Benchmarker", long_about = None)]
struct Cli {
    /// The script file to run or benchmark
    #[arg(value_name = "FILE")]
    file: Option<PathBuf>,

    /// Run a script string directly
    #[arg(short, long)]
    eval: Option<String>,

    /// Run benchmarks (Only allowed in Release mode)
    #[arg(short, long)]
    bench: bool,

    /// Update the benchmark baseline (Only allowed in Release mode)
    #[arg(short, long)]
    update_bench: bool,

    /// Dump the compiled AST
    #[arg(short, long)]
    dump: bool,

    /// Run with TracingObserver enabled
    #[arg(short, long)]
    trace: bool,

    /// Library directories to load before execution
    #[arg(short, long)]
    lib: Vec<PathBuf>,

    /// Disable optimization
    #[arg(long)]
    no_opt: bool,

    /// Start MCP server (stdio transport, for LLM integration)
    #[arg(long)]
    mcp: bool,

    /// Emit the parsed source back as Myc code (round-trip test)
    #[arg(long)]
    emit: bool,
}

fn main() {
    let cli = Cli::parse();

    if cli.mcp {
        if let Err(e) = myc::ast::mcp_server::run() {
            eprintln!("MCP server error: {}", e);
            std::process::exit(1);
        }
        return;
    }
    let mut env = Environment::new();
    env.optimization = !cli.no_opt;

    // Load libraries (Now just search paths)
    for lib_path in &cli.lib {
        env.add_search_path(lib_path);
    }

    if cli.bench || cli.update_bench {
        if cfg!(debug_assertions) {
            eprintln!("❌ ERROR: Benchmarks must be run in Release mode!");
            eprintln!("   Use: cargo run --release --bin ast -- --bench");
            std::process::exit(1);
        }

        println!("🚀 Running benchmarks in RELEASE mode...\n");
        let filter = cli
            .file
            .as_ref()
            .and_then(|p| p.file_name())
            .map(|n| n.to_string_lossy().to_string());
        let results = tester::run_benchmarks(cli.update_bench, filter.as_deref());
        for res in results {
            let diff = res
                .diff_pct
                .map_or(String::new(), |d| format!(" ({:+.1}%)", d));
            println!("{}: {} - {:?}{}", res.status, res.name, res.median, diff);
        }
        return;
    }

    // Determine the source code to execute
    let source = if let Some(script_str) = cli.eval {
        Some(script_str)
    } else if let Some(file_path) = cli.file {
        if file_path.to_str() == Some("-") {
            read_stdin()
        } else {
            match fs::read_to_string(&file_path) {
                Ok(content) => Some(content),
                Err(e) => {
                    eprintln!("Error reading file {:?}: {}", file_path, e);
                    std::process::exit(1);
                }
            }
        }
    } else if !io::stdin().is_terminal() {
        read_stdin()
    } else {
        None
    };

    if let Some(content) = source {
        if cli.emit {
            let mut parser = MycParser::new(&content);
            let syntax_ast = parser.parse_expression();
            if parser.diagnostics.has_errors() {
                for diag in &parser.diagnostics.items {
                    eprintln!("ERROR: {}", diag.message);
                }
                std::process::exit(1);
            }
            print!("{}", emitter::emit(&syntax_ast));
        } else if cli.dump {
            match env.dump_ast(&content) {
                Ok(dump) => println!("{}", dump),
                Err(e) => eprintln!("Error dumping AST: {}", e),
            }
        } else if cli.trace {
            execute_trace(&env, &content);
        } else {
            execute(&env, &content);
        }
    } else {
        println!("MYC AST Compiler CLI. Use --help for usage.");
    }
}

fn read_stdin() -> Option<String> {
    let mut buffer = String::new();
    match io::stdin().read_to_string(&mut buffer) {
        Ok(_) => Some(buffer),
        Err(e) => {
            eprintln!("Error reading from stdin: {}", e);
            std::process::exit(1);
        }
    }
}

fn execute(env: &Environment, source: &str) {
    let result = env.compile(source);
    for diag in &result.diagnostics.items {
        let level = format!("{:?}", diag.level).to_uppercase();
        let loc = diag
            .identity
            .as_ref()
            .and_then(|id| id.location)
            .map(|l| format!(" at {}:{}", l.line, l.col))
            .unwrap_or_default();
        eprintln!("{}{} : {}", level, loc, diag.message);
    }

    if result.diagnostics.has_errors() {
        std::process::exit(1);
    }

    match env.run_script_compiled(result.ast.unwrap()) {
        Ok(res) => println!("{}", res),
        Err(e) => {
            eprintln!("Runtime Error: {}", e);
            std::process::exit(1);
        }
    }
}

fn execute_trace(env: &Environment, source: &str) {
    match env.run_debug(source) {
        Ok((Ok(result), logs)) => {
            for line in logs {
                println!("{}", line);
            }
            println!("Result: {}", result);
        }
        Ok((Err(e), logs)) => {
            for line in logs {
                println!("{}", line);
            }
            eprintln!("Runtime Error: {}", e);
            std::process::exit(1);
        }
        Err(e) => {
            eprintln!("Compilation Error: {}", e);
            std::process::exit(1);
        }
    }
}