llm-router/src/main.rs

use axum::{
    body::Body,
    extract::State,
    http::{Request, StatusCode},
    response::{IntoResponse, Response},
    routing::any,
    Router,
};
use clap::Parser;
use futures_util::StreamExt;
use ksni::{Tray, TrayMethods};
use serde::{Deserialize, Serialize};
use std::{
    fs,
    io::{self, Write},
    path::PathBuf,
    process::Command,
    sync::{
        atomic::{AtomicBool, AtomicUsize, Ordering},
        Arc,
    },
    time::Duration,
};

#[derive(Parser, Debug)]
#[command(
    name = "llm-router",
    version,
    about = "LLM Router Proxy with System Tray, Systemd Integration, and Preferred Backend"
)]
struct Args {
    /// Run the proxy service in daemon mode (non-interactive, used by systemd)
    #[arg(long)]
    daemon: 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 {
    listen_port: u16,
    backend_urls: Vec<String>,
    #[serde(default)]
    preferred_backend: Option<String>,
}

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

#[derive(Clone)]
struct AppState {
    backends: Vec<Backend>,
    preferred_backend: Option<String>,
    client: reqwest::Client,
    counter: Arc<AtomicUsize>,
}

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

impl Tray for LLMRouterTray {
    fn id(&self) -> String {
        "llm-router".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 Router ({}/{})", online_count, self.backend_data.len())
    }

    fn tool_tip(&self) -> ksni::ToolTip {
        let mut desc = String::from("Managed 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 Router 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 Router".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");
    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.service");
    path
}

fn uninstall_service() {
    println!("--- Uninstalling LLM Router ---");
    let _ = Command::new("systemctl")
        .args(["--user", "stop", "llm-router.service"])
        .status();
    println!("Stopped systemd service (if running).");
    let _ = Command::new("systemctl")
        .args(["--user", "disable", "llm-router.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");
    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.service"])
        .status()
    {
        Ok(status) if status.success() => {
            println!("Successfully started llm-router.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 Proxy Service
After=network.target

[Service]
Type=simple
ExecStart={} --daemon
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.service"])
        .status()
    {
        Ok(status) if status.success() => {
            println!("Successfully enabled and started llm-router.service as a user daemon!");
        }
        _ => {
            println!("Warning: Could not automatically start systemd service.");
        }
    }
}

fn load_config(is_daemon: bool) -> 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 {
        if is_daemon {
            eprintln!("Error: Configuration file not found at {}.", path.display());
            eprintln!("Please run the application interactively at least once without '--daemon' to complete setup.");
            std::process::exit(1);
        }

        println!("No configuration file found at {}.", path.display());
        println!("--- Interactive Setup Wizard ---");

        print!("Enter port for proxy to listen on [default: 4000]: ");
        io::stdout().flush().unwrap();
        let mut port_str = String::new();
        io::stdin().read_line(&mut port_str).unwrap();
        let listen_port = port_str.trim().parse().unwrap_or(4000);

        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 {
            listen_port,
            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());

        print!(
            "Would you like to install and run this as a background systemd user service? [y/N]: "
        );
        io::stdout().flush().unwrap();
        let mut install_ans = String::new();
        io::stdin().read_line(&mut install_ans).unwrap();
        if install_ans.trim().to_lowercase().starts_with('y') {
            install_systemd_service();
        }

        config
    }
}

#[tokio::main]
async fn main() {
    let args = Args::parse();
    if args.uninstall {
        uninstall_service();
    }
    if args.start_service {
        start_systemd_service();
    }
    let config = load_config(args.daemon);
    let client = reqwest::Client::builder()
        .timeout(Duration::from_secs(120))
        .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 = LLMRouterTray {
        backend_data: initial_backend_data,
        preferred_url: config.preferred_backend.clone(),
    };
    let tray_handle = tray.spawn().await.expect("Failed to start system tray");

    let state = AppState {
        backends: backends.clone(),
        preferred_backend: config.preferred_backend.clone(),
        client: client.clone(),
        counter: Arc::new(AtomicUsize::new(0)),
    };

    let tray_handle_clone = tray_handle.clone();
    let pref_url_for_health = config.preferred_backend.clone();
    tokio::spawn(async move {
        health_check_loop(backends, pref_url_for_health, client, tray_handle_clone).await;
    });

    let app = Router::new()
        .route("/*path", any(proxy_handler))
        .fallback(any(proxy_handler))
        .with_state(state);

    let bind_address = format!("0.0.0.0:{}", config.listen_port);
    let listener = tokio::net::TcpListener::bind(&bind_address).await.unwrap();
    println!(
        "LLM Router running on http://0.0.0.0:{}",
        config.listen_port
    );
    axum::serve(listener, app).await.unwrap();
}

async fn health_check_loop(
    backends: Vec<Backend>,
    preferred_url: Option<String>,
    client: reqwest::Client,
    tray_handle: ksni::Handle<LLMRouterTray>,
) {
    loop {
        for backend in &backends {
            let health_url = format!("{}/v1/models", backend.url);
            let res = client
                .get(&health_url)
                .timeout(Duration::from_secs(2))
                .send()
                .await;
            let is_healthy = res.map_or(false, |r| r.status().is_success());
            let prev = backend.healthy.load(Ordering::Relaxed);

            if prev != is_healthy {
                println!(
                    "Backend [{}] status changed: {}",
                    backend.url,
                    if is_healthy { "ONLINE" } else { "OFFLINE" }
                );
            }
            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_clone = preferred_url.clone();
        tray_handle
            .update(move |tray: &mut LLMRouterTray| {
                tray.backend_data = snapshot;
                tray.preferred_url = pref_url_clone;
            })
            .await;
        tokio::time::sleep(Duration::from_secs(5)).await;
    }
}

fn select_backend(
    backends: &[Backend],
    preferred_backend: &Option<String>,
    counter: &AtomicUsize,
) -> Option<String> {
    let len = backends.len();
    if len == 0 {
        return None;
    }
    if let Some(pref) = preferred_backend {
        if let Some(backend) = backends.iter().find(|b| &b.url == pref) {
            if backend.healthy.load(Ordering::Relaxed) {
                return Some(backend.url.clone());
            }
        }
    }
    let start_idx = counter.fetch_add(1, Ordering::Relaxed);
    for i in 0..len {
        let idx = (start_idx + i) % len;
        if backends[idx].healthy.load(Ordering::Relaxed) {
            return Some(backends[idx].url.clone());
        }
    }
    None
}

async fn proxy_handler(State(state): State<AppState>, req: Request<Body>) -> Response {
    let backend_url =
        match select_backend(&state.backends, &state.preferred_backend, &state.counter) {
            Some(url) => url,
            None => {
                return (
                    StatusCode::SERVICE_UNAVAILABLE,
                    "All LLM backends are offline",
                )
                    .into_response()
            }
        };
    let path = req.uri().path().to_string();
    let query = req
        .uri()
        .query()
        .map(|v| format!("?{}", v))
        .unwrap_or_default();
    let target_url = format!("{}{}{}", backend_url, path, query);
    let method = req.method().clone();
    let headers = req.headers().clone();
    let body_bytes = match axum::body::to_bytes(req.into_body(), usize::MAX).await {
        Ok(bytes) => bytes,
        Err(_) => return (StatusCode::BAD_REQUEST, "Failed to read request body").into_response(),
    };
    let mut forward_req = state.client.request(method, &target_url).body(body_bytes);
    for (key, value) in headers.iter() {
        forward_req = forward_req.header(key, value);
    }
    match forward_req.send().await {
        Ok(resp) => {
            let status = resp.status();
            let mut response_builder = Response::builder().status(status);
            for (key, value) in resp.headers() {
                response_builder = response_builder.header(key, value);
            }
            let stream = resp.bytes_stream().map(|result| {
                result.map_err(|err| std::io::Error::new(std::io::ErrorKind::Other, err))
            });
            let body = Body::from_stream(stream);
            response_builder.body(body).unwrap_or_else(|_| {
                (StatusCode::INTERNAL_SERVER_ERROR, "Response build error").into_response()
            })
        }
        Err(_) => (StatusCode::BAD_GATEWAY, "Failed to reach LLM backend").into_response(),
    }
}