llm-router-companion/src/main.rs

use clap::Parser;
use ksni::{Tray, TrayMethods};
use serde::{Deserialize, Serialize};
use std::{
    fs,
    io::{self, Write},
    path::PathBuf,
    process::Command,
    sync::{
        Arc,
        atomic::{AtomicBool, Ordering},
    },
    time::Duration,
};

#[derive(Parser, Debug)]
#[command(
    name = "llm-router-tray",
    version,
    about = "LLM Server Status Tray Monitor with Systemd Integration and Preferred Backend"
)]
struct Args {
    /// Interactively setup configuration and install the systemd user service
    #[arg(long)]
    install: bool,

    /// Start the installed systemd user service
    #[arg(long)]
    start_service: bool,

    /// Uninstall the service, remove configuration, and clean up systemd files
    #[arg(long)]
    uninstall: bool,
}

#[derive(Serialize, Deserialize, Clone)]
struct Config {
    backend_urls: Vec<String>,
    #[serde(default)]
    preferred_backend: Option<String>,
}

#[derive(Clone)]
struct Backend {
    url: String,
    healthy: Arc<AtomicBool>,
}

#[derive(Debug)]
struct LLMTray {
    backend_data: Vec<(String, bool)>,
    preferred_url: Option<String>,
}

impl Tray for LLMTray {
    fn id(&self) -> String {
        "llm-router-tray".into()
    }

    fn icon_name(&self) -> String {
        "cpu".into()
    }

    fn title(&self) -> String {
        let online_count = self.backend_data.iter().filter(|(_, h)| *h).count();
        format!("LLM Monitor ({}/{})", online_count, self.backend_data.len())
    }

    fn tool_tip(&self) -> ksni::ToolTip {
        let mut desc = String::from("Monitored LLM Endpoints:\n");
        for (url, healthy) in &self.backend_data {
            let status = if *healthy {
                "🟢 Online"
            } else {
                "🔴 Offline"
            };
            let is_pref = self.preferred_url.as_ref() == Some(url);
            let pref_tag = if is_pref { " [Preferred]" } else { "" };
            desc.push_str(&format!("• {}{}: {}\n", url, pref_tag, status));
        }
        ksni::ToolTip {
            title: "LLM Server Status".into(),
            description: desc.trim().into(),
            ..Default::default()
        }
    }

    fn menu(&self) -> Vec<ksni::menu::MenuItem<Self>> {
        use ksni::menu::*;
        let mut items = vec![];
        items.push(
            StandardItem {
                label: "Live Endpoint Statuses:".into(),
                enabled: false,
                ..Default::default()
            }
            .into(),
        );
        items.push(MenuItem::Separator);

        for (url, healthy) in &self.backend_data {
            let status_str = if *healthy {
                "🟢 Online"
            } else {
                "🔴 Offline"
            };
            let is_pref = self.preferred_url.as_ref() == Some(url);
            let label = if is_pref {
                format!("★ {} [{}] (Preferred)", url, status_str)
            } else {
                format!("{} [{}]", url, status_str)
            };
            items.push(
                StandardItem {
                    label,
                    enabled: false,
                    ..Default::default()
                }
                .into(),
            );
        }

        items.push(MenuItem::Separator);
        items.push(
            StandardItem {
                label: "Quit Monitor".into(),
                icon_name: "application-exit".into(),
                activate: Box::new(|_| std::process::exit(0)),
                ..Default::default()
            }
            .into(),
        );
        items
    }
}

fn get_config_path() -> PathBuf {
    let mut path = dirs::config_dir().unwrap_or_else(|| PathBuf::from("."));
    path.push("llm-router-tray");
    fs::create_dir_all(&path).unwrap();
    path.push("config.toml");
    path
}

fn get_service_path() -> PathBuf {
    let mut path = dirs::config_dir().unwrap_or_else(|| PathBuf::from("."));
    path.push("systemd");
    path.push("user");
    path.push("llm-router-tray.service");
    path
}

fn uninstall_service() {
    println!("--- Uninstalling LLM Router Tray ---");
    let _ = Command::new("systemctl")
        .args(["--user", "stop", "llm-router-tray.service"])
        .status();
    println!("Stopped systemd service (if running).");

    let _ = Command::new("systemctl")
        .args(["--user", "disable", "llm-router-tray.service"])
        .status();
    println!("Disabled systemd service.");

    let service_path = get_service_path();
    if service_path.exists() {
        if let Err(e) = fs::remove_file(&service_path) {
            println!("Failed to remove service file: {}", e);
        } else {
            println!("Deleted service file: {}", service_path.display());
        }
    }

    let _ = Command::new("systemctl")
        .args(["--user", "daemon-reload"])
        .status();

    let mut config_dir = dirs::config_dir().unwrap_or_else(|| PathBuf::from("."));
    config_dir.push("llm-router-tray");
    if config_dir.exists() {
        if let Err(e) = fs::remove_dir_all(&config_dir) {
            println!("Failed to remove config folder: {}", e);
        } else {
            println!("Removed configuration folder: {}", config_dir.display());
        }
    }
    println!("\nUninstallation complete!");
    std::process::exit(0);
}

fn start_systemd_service() {
    match Command::new("systemctl")
        .args(["--user", "start", "llm-router-tray.service"])
        .status()
    {
        Ok(status) if status.success() => {
            println!("Successfully started llm-router-tray.service!");
        }
        _ => {
            println!("Failed to start systemd service. Ensure it has been installed first.");
        }
    }
    std::process::exit(0);
}

fn install_systemd_service() {
    let exe_path = match std::env::current_exe() {
        Ok(path) => path,
        Err(e) => {
            println!("Failed to get current executable path: {}", e);
            return;
        }
    };
    let service_path = get_service_path();
    if let Some(parent) = service_path.parent() {
        let _ = fs::create_dir_all(parent);
    }

    let service_content = format!(
        r#"[Unit]
Description=LLM Router Tray Monitor Service
After=network.target graphical-session.target

[Service]
Type=simple
ExecStart={}
Restart=on-failure
RestartSec=5

[Install]
WantedBy=default.target
"#,
        exe_path.display()
    );

    if let Err(e) = fs::write(&service_path, service_content) {
        println!("Failed to write systemd service file: {}", e);
        return;
    }
    println!("Created systemd service at: {}", service_path.display());

    let _ = Command::new("systemctl")
        .args(["--user", "daemon-reload"])
        .status();

    match Command::new("systemctl")
        .args(["--user", "enable", "--now", "llm-router-tray.service"])
        .status()
    {
        Ok(status) if status.success() => {
            println!("Successfully enabled and started llm-router-tray.service as a user daemon!");
        }
        _ => {
            println!("Warning: Could not automatically start systemd service.");
        }
    }
}

fn install_interactive() {
    let path = get_config_path();
    println!("--- LLM Router Tray Interactive Setup Wizard ---");
    let mut backend_urls = Vec::new();
    println!("Enter your LLM server backend URLs (leave blank and press Enter when finished):");
    loop {
        print!("Backend URL [e.g., http://192.168.1.50:11434]: ");
        io::stdout().flush().unwrap();
        let mut url = String::new();
        io::stdin().read_line(&mut url).unwrap();
        let url = url.trim().to_string();

        if url.is_empty() {
            if backend_urls.is_empty() {
                println!("Error: You must provide at least one backend server URL.");
                continue;
            }
            break;
        }
        backend_urls.push(url);
    }

    let mut preferred_backend = None;
    if backend_urls.len() > 1 {
        print!("Enter a preferred backend URL from the list above (leave blank for none): ");
        io::stdout().flush().unwrap();
        let mut pref = String::new();
        io::stdin().read_line(&mut pref).unwrap();
        let pref = pref.trim().to_string();
        if !pref.is_empty() {
            if backend_urls.contains(&pref) {
                preferred_backend = Some(pref);
            } else {
                println!(
                    "Warning: Provided preferred URL was not in the backend list. Ignoring preference."
                );
            }
        }
    }
    let config = Config {
        backend_urls,
        preferred_backend,
    };
    let toml_content = toml::to_string_pretty(&config).expect("Failed to serialize config");
    fs::write(&path, toml_content).expect("Failed to write config file");
    println!("\nConfiguration successfully saved to {}\n", path.display());
    install_systemd_service();
    std::process::exit(0);
}

fn load_config() -> Config {
    let path = get_config_path();
    if path.exists() {
        let content = fs::read_to_string(&path).expect("Failed to read config file");
        toml::from_str(&content).expect("Failed to parse config TOML")
    } else {
        println!("No configuration file found at {}.", path.display());
        println!("Running interactive setup wizard...");
        install_interactive();
        unreachable!()
    }
}

#[tokio::main]
async fn main() {
    let args = Args::parse();
    if args.uninstall {
        uninstall_service();
    }
    if args.start_service {
        start_systemd_service();
    }
    if args.install {
        install_interactive();
    }
    let config = load_config();
    let client = reqwest::Client::builder()
        .timeout(Duration::from_secs(2))
        .build()
        .unwrap();
    let mut backends = Vec::new();
    let mut initial_backend_data = Vec::new();
    for url in config.backend_urls {
        backends.push(Backend {
            url: url.clone(),
            healthy: Arc::new(AtomicBool::new(false)),
        });
        initial_backend_data.push((url, false));
    }
    let tray = LLMTray {
        backend_data: initial_backend_data,
        preferred_url: config.preferred_backend.clone(),
    };
    let tray_handle = tray.spawn().await.expect("Failed to start system tray");
    let preferred_url_clone = config.preferred_backend.clone();

    tokio::spawn(async move {
        loop {
            for backend in &backends {
                let health_url = format!("{}/v1/models", backend.url);
                let res = client.get(&health_url).send().await;
                let is_healthy = res.map_or(false, |r| r.status().is_success());
                backend.healthy.store(is_healthy, Ordering::Relaxed);
            }
            let snapshot: Vec<(String, bool)> = backends
                .iter()
                .map(|b| (b.url.clone(), b.healthy.load(Ordering::Relaxed)))
                .collect();
            let pref_url_ref = preferred_url_clone.clone();
            tray_handle
                .update(move |tray: &mut LLMTray| {
                    tray.backend_data = snapshot;
                    tray.preferred_url = pref_url_ref;
                })
                .await;
            tokio::time::sleep(Duration::from_secs(5)).await;
        }
    });
    std::future::pending::<()>().await;
}