rhaj/devtools/reloadd/src/main.rs

use clap::{Parser};
use notify::{RecursiveMode, Watcher, EventKind, recommended_watcher};
use std::sync::{Arc, Mutex};
use std::process::{Command, Child, Stdio};
use tokio::net::TcpStream;
use tokio::sync::broadcast;
use tokio_tungstenite::accept_async;
use futures_util::{SinkExt, StreamExt};
use tokio::time::{sleep, Duration};
use tokio::io::AsyncWriteExt;
use std::net::TcpListener as StdTcpListener;
use std::net::SocketAddr;

#[derive(Parser, Debug)]
#[command(author, version, about)]
struct Args {
    /// Paths to watch (like src/, templates/, scripts/)
    #[arg(short, long, value_name = "PATH", num_args = 1.., required = true)]
    watch: Vec<String>,

    /// Port for the web server (default: 8080)
    #[arg(short, long, default_value = "8080")]
    port: u16,

    /// Multiple commands to run (format: --run "cargo build" --run "cargo test")
    #[arg(short, long, value_name = "COMMAND", num_args = 1..)]
    run: Vec<String>,

    /// Command to run on change (like: -- run --example server)
    /// This is kept for backward compatibility
    #[arg(last = true)]
    command: Vec<String>,
}

// Run commands in sequence, with the last one potentially being a long-running server
async fn run_commands_in_sequence(commands: &[Vec<String>], server_process: Arc<Mutex<Option<Child>>>) -> bool {
    // Run all commands in sequence
    for (i, command) in commands.iter().enumerate() {
        let is_last = i == commands.len() - 1;
        let program = if command[0] == "cargo" || command[0].ends_with("/cargo") {
            "cargo".to_string()
        } else {
            command[0].clone()
        };
        
        let args = if command[0] == "cargo" || command[0].ends_with("/cargo") {
            command.iter().skip(1).cloned().collect::<Vec<String>>()
        } else {
            command.iter().skip(1).cloned().collect::<Vec<String>>()
        };
        
        let mut cmd = Command::new(&program);
        cmd.args(&args);
        cmd.stdout(Stdio::inherit())
            .stderr(Stdio::inherit());
        
        if is_last {
            // Last command might be a long-running server
            match cmd.spawn() {
                Ok(child) => {
                    *server_process.lock().unwrap() = Some(child);
                    println!("🚀 Started server process: {} {}", program, args.join(" "));
                },
                Err(e) => {
                    eprintln!("❌ Failed to start server process: {}", e);
                    return false;
                }
            }
        } else {
            // For non-last commands, wait for them to complete
            println!("🔄 Running build step: {} {}", program, args.join(" "));
            match cmd.output() {
                Ok(output) => {
                    if !output.status.success() {
                        eprintln!("❌ Command failed: {} {}", program, args.join(" "));
                        return false;
                    }
                    println!("✅ Command completed successfully");
                },
                Err(e) => {
                    eprintln!("❌ Failed to execute command: {}", e);
                    return false;
                }
            }
        }
    }
    
    true
}

// Check if a port is already in use
fn is_port_in_use(port: u16) -> bool {
    StdTcpListener::bind(format!("127.0.0.1:{}", port)).is_err()
}

// Find an available port starting from the given port
fn find_available_port(start_port: u16) -> Option<u16> {
    let mut port = start_port;
    // Try up to 10 ports (start_port through start_port+9)
    for _ in 0..10 {
        if !is_port_in_use(port) {
            return Some(port);
        }
        port += 1;
    }
    None
}

// Generate LiveReload script with the given WebSocket port
fn generate_livereload_script(ws_port: u16) -> String {
    format!(
        r#"<script>
const ws = new WebSocket("ws://localhost:{}");
ws.onmessage = (msg) => {{
  if (msg.data === "reload") location.reload();
}};
</script>"#,
        ws_port
    )
}
#[tokio::main]
async fn main() {
    let args = Args::parse();
    
    // Handle both the new --run and legacy command format
    let commands = if !args.run.is_empty() {
        // New format: each --run is a separate command
        args.run.iter().map(|cmd| {
            // Split the command string into arguments
            cmd.split_whitespace().map(String::from).collect::<Vec<String>>()
        }).collect::<Vec<Vec<String>>>()
    } else if !args.command.is_empty() {
        // Legacy format: single command from the trailing arguments
        println!("Command: {:?}", args.command);
        vec![args.command.clone()]
    } else {
        // No commands provided
        eprintln!("❌ Error: No commands provided. Use --run or trailing arguments.");
        std::process::exit(1);
    };

    // Check if server port is already in use
    let server_port = args.port;
    if is_port_in_use(server_port) {
        eprintln!("❌ Error: Port {} is already in use. Stop any running instances before starting a new one.", server_port);
        std::process::exit(1);
    }

    // Find an available WebSocket port
    let ws_port = match find_available_port(35729) {
        Some(port) => port,
        None => {
            eprintln!("❌ Error: Could not find an available WebSocket port. Please free up some ports and try again.");
            std::process::exit(1);
        }
    };

    let (tx, _) = broadcast::channel::<()>(10);
    let tx_ws = tx.clone();
    let server_process = Arc::new(Mutex::new(None::<Child>));

    // Validate paths before starting the watcher
    let mut invalid_paths = Vec::new();
    for path in &args.watch {
        if !std::path::Path::new(path).exists() {
            invalid_paths.push(path.clone());
        }
    }

    if !invalid_paths.is_empty() {
        eprintln!("❌ Error: The following watch paths do not exist:");
        for path in invalid_paths {
            eprintln!("   - {}", path);
        }
        eprintln!("Please provide valid paths to watch.");
        std::process::exit(1);
    }

    // Start WebSocket reload server
    match start_websocket_server(tx_ws.clone(), ws_port).await {
        Ok(_) => {},
        Err(e) => {
            eprintln!("❌ Failed to start WebSocket server: {}", e);
            std::process::exit(1);
        }
    }
    
    // Output the LiveReload script for users to add to their projects
    println!("📋 Add this script to your HTML for live reloading:");
    println!("{}", generate_livereload_script(ws_port));

    // 🚀 Run all commands in sequence
    if !run_commands_in_sequence(&commands, Arc::clone(&server_process)).await {
        eprintln!("❌ Command execution failed.");
        std::process::exit(1);
    }

    println!("🔁 Watching paths: {:?}", args.watch);
    println!("🌐 Connect browser to ws://localhost:{}", ws_port);
    println!("🖥️  Web server running on http://localhost:{}", server_port);

    // Create a runtime handle for the watcher to use
    let rt_handle = tokio::runtime::Handle::current();
    
    // Clone necessary values before moving into the watcher thread
    let watch_paths = args.watch.clone();
    let commands_clone = commands.clone();
    let server_process_clone = Arc::clone(&server_process);
    let tx_clone = tx.clone();

    // Create a dedicated thread for the file watcher
    let watcher_thread = std::thread::spawn(move || {
        // Create a new watcher
        let mut watcher = match recommended_watcher(move |res: notify::Result<notify::Event>| {
            if let Ok(event) = res {
                if matches!(event.kind, EventKind::Modify(_) | EventKind::Create(_) | EventKind::Remove(_)) {
                    println!("📦 Change detected, restarting...");
                    
                    // Kill previous process
                    if let Some(mut child) = server_process_clone.lock().unwrap().take() {
                        let _ = child.kill();
                    }
                    
                    // Run new process (only the primary command gets restarted)
                    if let Some(first_command) = commands.first() {
                        let program = if first_command[0] == "cargo" || first_command[0].ends_with("/cargo") {
                            "cargo".to_string()
                        } else {
                            first_command[0].clone()
                        };
                        
                        let args = if first_command[0] == "cargo" || first_command[0].ends_with("/cargo") {
                            first_command.iter().skip(1).cloned().collect::<Vec<String>>()
                        } else {
                            first_command.iter().skip(1).cloned().collect::<Vec<String>>()
                        };
                        
                        let mut cmd = Command::new(&program);
                        cmd.args(&args);
                        cmd.stdout(Stdio::inherit())
                            .stderr(Stdio::inherit());
                        
                        match cmd.spawn() {
                            Ok(child) => {
                                *server_process_clone.lock().unwrap() = Some(child);
                                
                                // Immediately send reload signal without waiting for server
                                let tx = tx_clone.clone();
                                rt_handle.spawn(async move {
                                    // Small delay to give the server a moment to start
                                    sleep(Duration::from_millis(500)).await;
                                    // Notify browser to reload
                                    let _ = tx.send(());
                                });
                            },
                            Err(e) => {
                                eprintln!("❌ Failed to spawn process: {}", e);
                            }
                        }
                    }
                }
            } else if let Err(e) = res {
                eprintln!("❌ Watch error: {}", e);
            }
        }) {
            Ok(w) => w,
            Err(e) => {
                eprintln!("❌ Failed to create watcher: {}", e);
                std::process::exit(1);
            }
        };
        
        // Add watches
        for path in &watch_paths {
            match watcher.watch(path.as_ref(), RecursiveMode::Recursive) {
                Ok(_) => println!("👁️  Watching path: {}", path),
                Err(e) => {
                    eprintln!("❌ Failed to watch path '{}': {}", path, e);
                    std::process::exit(1);
                }
            }
        }
        
        // Keep the thread alive
        loop {
            std::thread::sleep(std::time::Duration::from_secs(3600));
        }
    });
    
    // Wait for the watcher thread to finish (it shouldn't unless there's an error)
    match watcher_thread.join() {
        Ok(_) => {},
        Err(e) => {
            eprintln!("❌ Watcher thread panicked: {:?}", e);
            std::process::exit(1);
        }
    }
}

async fn start_websocket_server(tx: broadcast::Sender<()>, ws_port: u16) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
    let addr = format!("127.0.0.1:{}", ws_port);
    let listener = tokio::net::TcpListener::bind(&addr).await?;
    println!("WebSocket server started on ws://localhost:{}", ws_port);

    // Spawn a task to handle WebSocket connections
    tokio::spawn(async move {
        while let Ok((stream, addr)) = listener.accept().await {
            println!("New WebSocket connection from: {}", addr);
            let tx = tx.clone();
            let mut rx = tx.subscribe();

            tokio::spawn(async move {
                match accept_async(stream).await {
                    Ok(ws_stream) => {
                        let (mut write, _) = ws_stream.split();
                        
                        while rx.recv().await.is_ok() {
                            if let Err(e) = write.send(tokio_tungstenite::tungstenite::Message::Text("reload".into())).await {
                                eprintln!("❌ Error sending reload message to client {}: {}", addr, e);
                                break;
                            }
                        }
                    },
                    Err(e) => {
                        eprintln!("❌ Failed to accept WebSocket connection from {}: {}", addr, e);
                    }
                }
            });
        }
    });

    Ok(())
}