aiecho 0.1.4

//! AI-LAN Service Discovery CLI
//!
//! Command-line tools for service discovery:
//! - ai-discover-agent: Run the discovery server agent
//! - ai-discover scan: Scan for services on the network

use clap::{Parser, Subcommand};
use std::path::PathBuf;
use std::time::Duration;
use tracing::{debug, error, info};

use aiecho::{discover_services, ClientConfig, DiscoveryScanner, DiscoveryServer};

/// AI-LAN Service Discovery System
///
/// A lightweight, zero-config LAN service discovery mechanism for AI agents.
#[derive(Parser)]
#[command(name = "aiecho")]
#[command(version = "0.1.0")]
#[command(about = "AI-LAN Service Discovery System", long_about = None)]
struct Cli {
    #[command(subcommand)]
    command: Commands,
}

#[derive(Subcommand)]
enum Commands {
    /// Run the discovery agent (service side)
    Agent {
        /// Root directory to scan for .echo files
        #[arg(short, long, default_value = ".", value_name = "DIR")]
        root_path: PathBuf,

        /// Enable verbose logging
        #[arg(short, long)]
        verbose: bool,

        /// Override UDP discovery port
        #[arg(long)]
        udp_port: Option<u16>,
    },

    /// Scan for services on the network
    Scan {
        /// Output format: json, yaml, table
        #[arg(short, long, default_value = "json")]
        output: String,

        /// Scan timeout in seconds
        #[arg(short, long, default_value = "2.0")]
        timeout: f64,

        /// Skip fetching service manifests
        #[arg(long)]
        no_manifest: bool,

        /// Output to file instead of stdout
        #[arg(short, long, value_name = "FILE")]
        output_file: Option<PathBuf>,

        /// Enable verbose logging
        #[arg(short, long)]
        verbose: bool,
    },

    /// Listen for service changes in real-time
    Listen {
        /// Output services JSON file to watch
        #[arg(short, long, value_name = "FILE")]
        output_file: PathBuf,

        /// Auto-scan interval in seconds
        #[arg(short, long, default_value = "30")]
        interval: u32,

        /// Skip fetching service manifests
        #[arg(long)]
        no_manifest: bool,

        /// Enable verbose logging
        #[arg(short, long)]
        verbose: bool,
    },
}

#[tokio::main]
async fn main() {
    let cli = Cli::parse();

    // Setup logging
    let verbose = match &cli.command {
        Commands::Agent { verbose, .. } => *verbose,
        Commands::Scan { verbose, .. } => *verbose,
        Commands::Listen { verbose, .. } => *verbose,
    };

    tracing_subscriber::fmt()
        .with_max_level(if verbose {
            tracing::Level::DEBUG
        } else {
            tracing::Level::INFO
        })
        .init();

    // Execute command
    match cli.command {
        Commands::Agent {
            root_path,
            verbose: _,
            udp_port,
        } => {
            run_agent(root_path, udp_port).await;
        }
        Commands::Scan {
            output,
            timeout,
            no_manifest,
            output_file,
            verbose: _,
        } => {
            run_scan(output, timeout, no_manifest, output_file).await;
        }
        Commands::Listen {
            output_file,
            interval,
            no_manifest,
            verbose: _,
        } => {
            run_listen(output_file, interval, no_manifest).await;
        }
    }
}

async fn run_agent(root_path: PathBuf, udp_port: Option<u16>) {
    // Scan for .echo files
    info!("Scanning for .echo files in {}", root_path.display());

    // Discover services
    let services = discover_services(&root_path);

    if services.is_empty() {
        info!("No services found");
        return;
    }

    info!("Found {} service(s)", services.len());

    // Start servers for each service
    let mut servers = Vec::new();
    for (echo_path, mut service_config) in services {
        // Override UDP port if specified
        if let Some(port) = udp_port {
            service_config.udp_port = port;
        }

        info!("Starting discovery agent: {}", service_config.service_name);
        info!("  Service ID: {}", service_config.service_id);
        info!("  HTTP Port: {}", service_config.http_port);
        info!("  UDP Port: {}", service_config.udp_port);
        info!("  From: {}", echo_path.display());

        let mut server = DiscoveryServer::new(service_config);
        if let Err(e) = server.start().await {
            error!("Failed to start server: {}", e);
            continue;
        }

        servers.push(server);
    }

    if servers.is_empty() {
        error!("No servers started");
        return;
    }

    info!("All agents started. Press Ctrl+C to stop.");

    // Keep running
    tokio::signal::ctrl_c()
        .await
        .expect("Failed to listen for ctrl+c");

    info!("Stopping agents...");
    for mut server in servers {
        if let Err(e) = server.stop().await {
            error!("Error stopping server: {}", e);
        }
    }

    info!("All agents stopped.");
}

async fn run_scan(output: String, timeout: f64, no_manifest: bool, output_file: Option<PathBuf>) {
    info!("Scanning for services...");

    let config = ClientConfig {
        timeout,
        fetch_manifest: !no_manifest,
        output_format: output.clone(),
        ..Default::default()
    };

    let scanner = DiscoveryScanner::new(config);

    match scanner.scan(Some(!no_manifest)).await {
        Ok(services) => {
            if services.is_empty() {
                info!("No services found.");
                return;
            }

            info!("Found {} service(s)", services.len());

            // Format output
            match output.as_str() {
                "json" => {
                    let result: Vec<serde_json::Value> = services
                        .iter()
                        .map(|s| {
                            serde_json::json!({
                                "ip": s.ip(),
                                "port": s.port(),
                                "base_url": s.base_url(),
                                "manifest": s.manifest(),
                            })
                        })
                        .collect();

                    let json = serde_json::to_string_pretty(&result).unwrap();

                    if let Some(path) = output_file {
                        std::fs::write(&path, &json).unwrap();
                        info!("Output written to {}", path.display());
                    } else {
                        println!("{}", json);
                    }
                }
                "table" => {
                    println!("\n=== Discovered Services ===\n");
                    for s in &services {
                        println!("  {}:{}", s.ip(), s.port());
                        println!();
                    }
                }
                _ => {
                    error!("Unsupported output format: {}", output);
                }
            }
        }
        Err(e) => {
            error!("Scan failed: {}", e);
            std::process::exit(1);
        }
    }
}

async fn run_listen(output_file: PathBuf, interval: u32, no_manifest: bool) {
    info!("Listening for service changes...");
    info!("  Output file: {}", output_file.display());
    info!("  Auto-scan interval: {}s", interval);

    let config = ClientConfig {
        scan_interval: interval,
        fetch_manifest: !no_manifest,
        ..Default::default()
    };

    let scanner = DiscoveryScanner::new(config);

    // Initial scan
    info!("Running initial scan...");
    match scanner.scan(Some(!no_manifest)).await {
        Ok(services) => {
            if !services.is_empty() {
                info!("Found {} service(s)", services.len());
            }
        }
        Err(e) => {
            error!("Initial scan failed: {}", e);
        }
    }

    // Keep running with periodic scans
    loop {
        tokio::time::sleep(Duration::from_secs(interval as u64)).await;

        match scanner.scan(Some(!no_manifest)).await {
            Ok(services) => {
                if !services.is_empty() {
                    let result: Vec<serde_json::Value> = services
                        .iter()
                        .map(|s| {
                            serde_json::json!({
                                "ip": s.ip(),
                                "port": s.port(),
                                "manifest": s.manifest(),
                            })
                        })
                        .collect();

                    let json = serde_json::to_string_pretty(&result).unwrap();
                    let _ = std::fs::write(&output_file, &json);
                }
            }
            Err(e) => {
                debug!("Periodic scan failed: {}", e);
            }
        }
    }
}