neiam /devendra

//! devendra-controller: serve the web UI and author desired state into the
//! iroh document agents sync from.

use std::path::PathBuf;
use std::sync::Arc;

use anyhow::{Context, Result};
use clap::{Parser, Subcommand};
use minijinja::Environment;
use tracing_subscriber::{EnvFilter, layer::SubscriberExt, util::SubscriberInitExt};

use devendra_controller::{AppState, Config, agent_ticket, open_or_create_doc, publish, web};
use devendra_node::Node;

#[derive(Parser, Debug)]
#[command(
    name = "devendra-controller",
    about = "Web UI + iroh authoring for devendra"
)]
struct Cli {
    /// Directory holding desired state (`hosts/*.toml`, `files/`).
    #[arg(
        long,
        env = "DEVENDRA_CONFIG_DIR",
        default_value = "./config",
        global = true
    )]
    config_dir: PathBuf,

    /// Directory for iroh blob/doc stores + the persisted namespace id.
    #[arg(
        long,
        env = "DEVENDRA_DATA_DIR",
        default_value = "./data",
        global = true
    )]
    data_dir: PathBuf,

    /// Web UI bind address.
    #[arg(
        long,
        env = "DEVENDRA_BIND",
        default_value = "127.0.0.1:3900",
        global = true
    )]
    bind: String,

    /// Fixed UDP port for the iroh endpoint. Keep this stable so agents keep
    /// reaching the controller across restarts.
    #[arg(
        long,
        env = "DEVENDRA_IROH_PORT",
        default_value = "11204",
        global = true
    )]
    iroh_port: u16,

    #[command(subcommand)]
    cmd: Cmd,
}

#[derive(Subcommand, Debug)]
enum Cmd {
    /// Publish current config, then serve the web UI.
    Serve,
    /// Print the agent join ticket and exit.
    Ticket,
    /// Publish current config to the document and exit.
    Publish,
}

#[tokio::main]
async fn main() -> Result<()> {
    tracing_subscriber::registry()
        .with(EnvFilter::try_from_default_env().unwrap_or_else(|_| {
            "devendra_controller=info,devendra_node=info,tower_http=info".into()
        }))
        // Logs go to stderr so subcommands like `ticket`/`publish` keep stdout
        // clean for machine capture (e.g. `DEVENDRA_TICKET=$(… ticket)`).
        .with(tracing_subscriber::fmt::layer().with_writer(std::io::stderr))
        .init();

    let cli = Cli::parse();
    let config = Arc::new(Config {
        config_dir: cli.config_dir,
        data_dir: cli.data_dir,
        bind_address: cli.bind,
        iroh_port: cli.iroh_port,
    });

    let node = Arc::new(Node::spawn(&config.data_dir, Some(config.iroh_port)).await?);
    let doc = open_or_create_doc(&node, &config).await?;

    match cli.cmd {
        Cmd::Ticket => {
            let ticket = agent_ticket(&node, &doc).await?;
            println!("{ticket}");
        }
        Cmd::Publish => {
            let outcomes = publish::publish_all(&node, &doc, &config.config_dir).await?;
            for o in &outcomes {
                println!(
                    "{}: generation {} ({})",
                    o.host_id,
                    o.generation,
                    if o.changed { "updated" } else { "unchanged" }
                );
            }
            if outcomes.is_empty() {
                println!("no hosts found under {}", config.hosts_dir().display());
            }
        }
        Cmd::Serve => {
            // Make the on-disk config live before accepting requests.
            let outcomes = publish::publish_all(&node, &doc, &config.config_dir).await?;
            tracing::info!(hosts = outcomes.len(), "published initial desired state");

            let state = AppState {
                env: build_env(),
                node: node.clone(),
                doc: doc.clone(),
                config: config.clone(),
            };
            let app = web::router(state);

            let listener = tokio::net::TcpListener::bind(&config.bind_address)
                .await
                .with_context(|| format!("binding {}", config.bind_address))?;
            tracing::info!(addr = %config.bind_address, "controller web UI up");
            tracing::info!(
                "agent ticket available at http://{}/ticket",
                config.bind_address
            );
            axum::serve(listener, app)
                .with_graceful_shutdown(async {
                    let _ = tokio::signal::ctrl_c().await;
                    tracing::info!("shutting down");
                })
                .await?;
        }
    }

    // Flush stores + close the network so no half-written blob is left behind.
    node.shutdown().await?;
    Ok(())
}

/// In debug builds, load templates live from disk for fast iteration; in
/// release, use the copies embedded at build time by `build.rs`.
fn build_env() -> Environment<'static> {
    let mut env = Environment::new();
    #[cfg(debug_assertions)]
    {
        env.set_loader(minijinja::path_loader(concat!(
            env!("CARGO_MANIFEST_DIR"),
            "/templates"
        )));
    }
    #[cfg(not(debug_assertions))]
    {
        minijinja_embed::load_templates!(&mut env);
    }
    env
}