no-block-pls 0.1.0

use clap::{Parser, ValueEnum};
use no_block_pls::{inject_guard_module, instrument_async_only};
use std::fs;
use std::path::{Path, PathBuf};
use walkdir::WalkDir;

fn find_rust_files(dir: &Path) -> Vec<PathBuf> {
    WalkDir::new(dir)
        .into_iter()
        .filter_map(|entry| entry.ok())
        .filter(|entry| {
            !entry
                .path()
                .components()
                .any(|c| c.as_os_str() == "target" || c.as_os_str() == ".git")
        })
        .filter_map(|entry| {
            let path = entry.path();
            if path
                .extension()
                .and_then(|ext| ext.to_str())
                .map(|ext| ext == "rs")
                .unwrap_or(false)
            {
                if let Some(name) = path.file_name().and_then(|n| n.to_str())
                    && (name.ends_with(".instrumented.rs") || name.ends_with(".rs.bak"))
                {
                    return None;
                }
                Some(path.to_owned())
            } else {
                None
            }
        })
        .collect()
}

#[derive(Parser)]
#[command(name = "no-block-pls")]
#[command(about = "Instrument async Rust code to detect blocking operations")]
#[command(
    long_about = "Automatically inserts timing guards between await points in async functions to detect CPU-heavy operations that block the executor."
)]
struct Cli {
    /// Paths to instrument (files or directories)
    #[arg(default_value = ".")]
    paths: Vec<PathBuf>,

    /// Modify files in-place (creates .bak backups)
    #[arg(short = 'i', long = "in-place")]
    in_place: bool,

    /// Restore original files from .bak backups
    #[arg(short = 'r', long = "restore")]
    restore: bool,

    /// Threshold in milliseconds for blocking detection
    #[arg(short = 't', long = "threshold", default_value = "10")]
    threshold_ms: u64,

    /// Quiet mode - suppress informational messages
    #[arg(short = 'q', long = "quiet")]
    quiet: bool,
}

#[derive(Clone, Copy, Debug, ValueEnum)]
enum OutputMode {
    /// Print instrumented code to stdout
    Stdout,
    /// Write instrumented code to files
    InPlace,
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let cli = Cli::parse();

    if cli.restore {
        restore_files(&cli)?;
        return Ok(());
    }

    let mut files = Vec::new();
    for path in &cli.paths {
        if path.is_file() {
            if path.extension().and_then(|e| e.to_str()) == Some("rs") {
                files.push(path.clone());
            }
        } else if path.is_dir() {
            files.extend(find_rust_files(path));
        }
    }

    if files.is_empty() {
        eprintln!("No Rust files found");
        return Ok(());
    }

    if !cli.quiet {
        eprintln!(
            "Processing {} files (threshold: {}ms)...",
            files.len(),
            cli.threshold_ms
        );
    }

    let output_mode = if cli.in_place {
        OutputMode::InPlace
    } else {
        OutputMode::Stdout
    };

    for file in files {
        // Determine if this is a root file
        let is_root = file
            .file_name()
            .and_then(|n| n.to_str())
            .map(|n| n == "lib.rs" || n == "main.rs")
            .unwrap_or(false);

        let content = fs::read_to_string(&file)?;

        let result = if is_root {
            match inject_guard_module(&content, cli.threshold_ms) {
                Ok(output) if output == content => Ok(None),
                Ok(output) => Ok(Some(output)),
                Err(e) => Err(e),
            }
        } else {
            instrument_async_only(&content)
        };

        match result {
            Ok(Some(instrumented)) => match output_mode {
                OutputMode::InPlace => {
                    let backup_path = file.with_extension("rs.bak");
                    fs::copy(&file, &backup_path)?;
                    fs::write(&file, instrumented)?;
                    if !cli.quiet {
                        eprintln!(
                            "✓ Instrumented: {} (backup: {})",
                            file.display(),
                            backup_path.display()
                        );
                    }
                }
                OutputMode::Stdout => {
                    if !cli.quiet {
                        println!("// === {} ===", file.display());
                    }
                    println!("{}", instrumented);
                }
            },
            Ok(None) => {
                // File has no async code, skip it
                if !cli.quiet {
                    eprintln!("⊘ Skipped (no async): {}", file.display());
                }
            }
            Err(e) => {
                eprintln!("✗ Error processing {}: {}", file.display(), e);
            }
        }
    }

    if matches!(output_mode, OutputMode::InPlace) && !cli.quiet {
        eprintln!("\nTo run with instrumentation:");
        eprintln!("  cargo run --release");
        eprintln!("\nTo restore original files:");
        eprintln!("  no-block-pls --restore");
    }

    Ok(())
}

fn restore_files(cli: &Cli) -> Result<(), Box<dyn std::error::Error>> {
    let mut restored = 0;
    let mut errors = 0;

    for path in &cli.paths {
        let dir = if path.is_dir() {
            path.clone()
        } else {
            PathBuf::from(".")
        };

        for entry in WalkDir::new(dir) {
            let entry = entry?;
            let path = entry.path();

            if path.extension().and_then(|e| e.to_str()) == Some("bak")
                && let Some(original) = path.to_str()
                && original.ends_with(".rs.bak")
            {
                let target = path.with_extension("");
                match fs::rename(path, &target) {
                    Ok(()) => {
                        if !cli.quiet {
                            eprintln!("✓ Restored: {}", target.display());
                        }
                        restored += 1;
                    }
                    Err(e) => {
                        eprintln!("✗ Failed to restore {}: {}", path.display(), e);
                        errors += 1;
                    }
                }
            }
        }
    }

    if !cli.quiet {
        eprintln!("\nRestored {} files, {} errors", restored, errors);
    }

    Ok(())
}