agent-relay 0.1.0

use agent_relay::git::{self, SecureRelay};
use agent_relay::Relay;
use clap::{Parser, Subcommand};
use colored::Colorize;
use std::path::PathBuf;

#[derive(Parser)]
#[command(
    name = "agent-relay",
    version,
    about = "Agent-to-agent messaging for AI coding tools"
)]
struct Cli {
    /// Path to the relay directory (default: .relay in current dir)
    #[arg(long, default_value = ".relay")]
    dir: PathBuf,

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

#[derive(Subcommand)]
enum Commands {
    /// Register this agent with the relay
    Register {
        /// Agent name (e.g. "claude", "gemini", "gpt", "human")
        #[arg(long, short)]
        agent: String,
        /// Session identifier (unique per agent instance)
        #[arg(long, short)]
        session: String,
    },

    /// Unregister an agent
    Unregister {
        /// Session ID to remove
        #[arg(long, short)]
        session: String,
    },

    /// Send a message (broadcast or DM)
    Send {
        /// Message content
        message: String,
        /// Session ID of the sender
        #[arg(long, short = 'f', default_value = "cli")]
        from: String,
        /// Agent name of the sender
        #[arg(long, short, default_value = "human")]
        agent: String,
        /// Send to a specific session (omit for broadcast)
        #[arg(long, short)]
        to: Option<String>,
    },

    /// Check inbox
    Inbox {
        /// Session ID to check inbox for
        #[arg(long, short, default_value = "cli")]
        session: String,
        /// Max messages to show
        #[arg(long, short, default_value_t = 20)]
        limit: usize,
    },

    /// List active agents
    Agents,

    /// Show unread count
    Unread {
        /// Session ID
        #[arg(long, short, default_value = "cli")]
        session: String,
    },

    /// Clean up old messages and dead agents
    Cleanup {
        /// Max message age in seconds (default: 3600 = 1 hour)
        #[arg(long, default_value_t = 3600)]
        max_age: u64,
    },

    /// Initialize secure mode: detect git identity, SSH key, GitHub collaborators
    Init,

    /// Send a signed message (uses your SSH key)
    SecureSend {
        /// Message content
        message: String,
        /// Session ID of the sender
        #[arg(long, short = 'f', default_value = "cli")]
        from: String,
        /// Send to a specific session (omit for broadcast)
        #[arg(long, short)]
        to: Option<String>,
    },

    /// Read inbox with signature verification
    SecureInbox {
        /// Session ID to check inbox for
        #[arg(long, short, default_value = "cli")]
        session: String,
        /// Max messages to show
        #[arg(long, short, default_value_t = 20)]
        limit: usize,
    },

    /// List GitHub collaborators who can message this repo
    Collaborators,

    /// Watch for new messages (poll every N seconds)
    Watch {
        /// Session ID to watch for
        #[arg(long, short, default_value = "cli")]
        session: String,
        /// Poll interval in seconds
        #[arg(long, short, default_value_t = 2)]
        interval: u64,
        /// Command to run when new messages arrive (optional)
        #[arg(long, short)]
        exec: Option<String>,
    },
}

fn main() {
    let cli = Cli::parse();
    let dir = cli.dir;
    let relay = Relay::new(dir.clone());

    match cli.command {
        Commands::Register { agent, session } => {
            let reg = relay.register(&agent, &session, std::process::id());
            println!(
                "{} Registered {} as {} (pid {})",
                "OK".green().bold(),
                session.cyan(),
                agent.yellow(),
                reg.pid
            );
        }

        Commands::Unregister { session } => {
            relay.unregister(&session);
            println!("{} Unregistered {}", "OK".green().bold(), session.cyan());
        }

        Commands::Send {
            message,
            from,
            agent,
            to,
        } => {
            let msg = relay.send(&from, &agent, to.as_deref(), &message);
            let target = to.as_deref().unwrap_or("all");
            println!(
                "{} [{}] {} -> {}: {}",
                "SENT".green().bold(),
                msg.id.dimmed(),
                agent.yellow(),
                target.cyan(),
                message
            );
        }

        Commands::Inbox { session, limit } => {
            let msgs = relay.inbox(&session, limit);
            if msgs.is_empty() {
                println!("{}", "No messages.".dimmed());
                return;
            }
            for (msg, is_new) in &msgs {
                let marker = if *is_new {
                    "NEW".green().bold().to_string()
                } else {
                    "read".dimmed().to_string()
                };
                let ts = format_timestamp(msg.timestamp);
                let target = msg
                    .to_session
                    .as_deref()
                    .map(|t| format!(" -> {}", t.cyan()))
                    .unwrap_or_default();
                println!(
                    "[{}] {} {}{}: {}  {}",
                    marker,
                    msg.from_agent.yellow(),
                    format!("({})", msg.from_session).dimmed(),
                    target,
                    msg.content,
                    ts.dimmed()
                );
            }
        }

        Commands::Agents => {
            relay.cleanup_dead();
            let agents = relay.agents();
            if agents.is_empty() {
                println!("{}", "No agents registered.".dimmed());
                return;
            }
            println!(
                "{:<20} {:<15} {:<8} {}",
                "SESSION".bold(),
                "AGENT".bold(),
                "PID".bold(),
                "HEARTBEAT".bold()
            );
            for a in &agents {
                let ts = format_timestamp(a.last_heartbeat);
                println!(
                    "{:<20} {:<15} {:<8} {}",
                    a.session_id.cyan(),
                    a.agent_id.yellow(),
                    a.pid,
                    ts.dimmed()
                );
            }
        }

        Commands::Unread { session } => {
            let count = relay.unread_count(&session);
            if count == 0 {
                println!("{}", "No unread messages.".dimmed());
            } else {
                println!(
                    "{} {} unread message{}",
                    "INBOX".green().bold(),
                    count,
                    if count == 1 { "" } else { "s" }
                );
            }
        }

        Commands::Cleanup { max_age } => {
            let dead = relay.cleanup_dead();
            let old = relay.cleanup_old(max_age);
            println!(
                "{} Removed {} dead agent{}, {} old message{}",
                "OK".green().bold(),
                dead,
                if dead == 1 { "" } else { "s" },
                old,
                if old == 1 { "" } else { "s" }
            );
        }

        Commands::Init => {
            let dir = dir.clone();
            match SecureRelay::from_git_repo(dir) {
                Ok(sr) => {
                    println!(
                        "{} Git identity: {} <{}>",
                        "OK".green().bold(),
                        sr.identity.name.yellow(),
                        sr.identity.email.cyan()
                    );
                    match &sr.ssh_key {
                        Some(key) => println!(
                            "{} SSH key: {}",
                            "OK".green().bold(),
                            key.display().to_string().dimmed()
                        ),
                        None => println!(
                            "{} No SSH key found — messages will be unsigned",
                            "WARN".yellow().bold()
                        ),
                    }

                    // Try to fetch GitHub collaborators and build allowed_signers
                    match sr.init_allowed_signers() {
                        Ok(count) => println!(
                            "{} Imported {} SSH key{} from GitHub collaborators",
                            "OK".green().bold(),
                            count,
                            if count == 1 { "" } else { "s" }
                        ),
                        Err(e) => println!(
                            "{} Could not fetch collaborators: {} (messaging still works, just unsigned)",
                            "WARN".yellow().bold(),
                            e.dimmed()
                        ),
                    }
                }
                Err(e) => {
                    eprintln!("{} {}", "ERROR".red().bold(), e);
                    std::process::exit(1);
                }
            }
        }

        Commands::SecureSend { message, from, to } => {
            let dir = dir.clone();
            match SecureRelay::from_git_repo(dir) {
                Ok(sr) => match sr.send_signed(&from, to.as_deref(), &message) {
                    Ok(signed) => {
                        let target = to.as_deref().unwrap_or("all");
                        let sig_status = if signed.signature.is_some() {
                            "SIGNED".green().bold().to_string()
                        } else {
                            "UNSIGNED".yellow().bold().to_string()
                        };
                        println!(
                            "{} [{}] {} -> {}: {}",
                            sig_status,
                            signed.message.id.dimmed(),
                            sr.identity.name.yellow(),
                            target.cyan(),
                            message
                        );
                    }
                    Err(e) => {
                        eprintln!("{} {}", "ERROR".red().bold(), e);
                        std::process::exit(1);
                    }
                },
                Err(e) => {
                    eprintln!("{} {}", "ERROR".red().bold(), e);
                    std::process::exit(1);
                }
            }
        }

        Commands::SecureInbox { session, limit } => {
            let dir = dir.clone();
            match SecureRelay::from_git_repo(dir) {
                Ok(sr) => {
                    let msgs = sr.inbox_verified(&session, limit);
                    if msgs.is_empty() {
                        println!("{}", "No messages.".dimmed());
                        return;
                    }
                    for signed in &msgs {
                        let verified = match signed.verified {
                            Some(true) => format!("{}", "VERIFIED".green().bold()),
                            Some(false) => format!("{}", "INVALID".red().bold()),
                            None => {
                                if signed.signature.is_some() {
                                    format!("{}", "SIGNED".yellow().bold())
                                } else {
                                    format!("{}", "UNSIGNED".dimmed())
                                }
                            }
                        };
                        let ts = format_timestamp(signed.message.timestamp);
                        let target = signed
                            .message
                            .to_session
                            .as_deref()
                            .map(|t| format!(" -> {}", t.cyan()))
                            .unwrap_or_default();
                        println!(
                            "[{}] {} <{}>{}: {}  {}",
                            verified,
                            signed.git_identity.name.yellow(),
                            signed.git_identity.email.dimmed(),
                            target,
                            signed.message.content,
                            ts.dimmed()
                        );
                    }
                }
                Err(e) => {
                    eprintln!("{} {}", "ERROR".red().bold(), e);
                    std::process::exit(1);
                }
            }
        }

        Commands::Collaborators => match git::parse_github_remote() {
            Some((owner, repo)) => {
                println!("Collaborators for {}/{}:", owner.cyan(), repo.cyan());
                match git::github_collaborators(&owner, &repo) {
                    Ok(collabs) => {
                        for c in &collabs {
                            let perm = if c.can_push {
                                "push".green()
                            } else {
                                "read".dimmed()
                            };
                            println!("  {} [{}]", c.username.yellow(), perm);
                        }
                    }
                    Err(e) => {
                        eprintln!("{} {}", "ERROR".red().bold(), e);
                        std::process::exit(1);
                    }
                }
            }
            None => {
                eprintln!(
                    "{} Not a GitHub repo (no origin remote found)",
                    "ERROR".red().bold()
                );
                std::process::exit(1);
            }
        },

        Commands::Watch {
            session,
            interval,
            exec,
        } => {
            println!(
                "Watching for messages to {} (every {}s)...",
                session.cyan(),
                interval
            );
            loop {
                let count = relay.poll(&session);
                if count > 0 {
                    println!(
                        "{} {} new message{}",
                        "PING".green().bold(),
                        count,
                        if count == 1 { "" } else { "s" }
                    );
                    if let Some(ref cmd) = exec {
                        let _ = std::process::Command::new("sh").arg("-c").arg(cmd).spawn();
                    }
                }
                std::thread::sleep(std::time::Duration::from_secs(interval));
            }
        }
    }
}

fn format_timestamp(ts: u64) -> String {
    use chrono::{DateTime, Utc};
    let dt = DateTime::<Utc>::from_timestamp(ts as i64, 0);
    match dt {
        Some(d) => d.format("%H:%M:%S").to_string(),
        None => ts.to_string(),
    }
}