use std::any::TypeId;
use std::collections::{BTreeMap, HashMap};
use std::marker::PhantomData;
use std::mem::take;
use std::process::ExitCode;
use std::sync::Arc;
use async_trait::async_trait;
use clap::{Arg, ArgAction, ArgMatches};
use diode::{App, AppBuilder};
use crate::{CancellationToken, Config, Metrics, RunDaemonsExt, Tracing};
pub trait Command: Send + Sync {
fn command() -> clap::Command
where
Self: Sized;
fn main(
app: Arc<App>,
matches: ArgMatches,
) -> impl std::future::Future<Output = ExitCode> + Send {
let _ = (app, matches);
async move { ExitCode::FAILURE }
}
}
#[async_trait]
trait DynCommand: Send + Sync {
fn command(&self) -> clap::Command;
async fn main(&self, app: Arc<App>, matches: ArgMatches) -> ExitCode;
}
#[async_trait]
impl<T> DynCommand for T
where
T: Command,
{
fn command(&self) -> clap::Command {
T::command()
}
async fn main(&self, app: Arc<App>, matches: ArgMatches) -> ExitCode {
T::main(app, matches).await
}
}
#[derive(Default)]
#[doc(hidden)]
pub struct CommandRegistry {
commands: HashMap<TypeId, Box<dyn DynCommand>>,
}
impl CommandRegistry {
pub fn add_command<T>(&mut self)
where
T: Command + 'static,
{
let type_id = TypeId::of::<T>();
self.commands
.insert(type_id, Box::new(CommandWrapper::<T>(PhantomData)));
}
pub fn has_command<T>(&self) -> bool
where
T: Command + 'static,
{
let type_id = TypeId::of::<T>();
self.commands.contains_key(&type_id)
}
pub fn build_cli(&self) -> clap::Command {
let mut cli = clap::Command::default()
.subcommand_required(true)
.arg(Arg::new("config").long("config").short('c').required(true))
.arg(
Arg::new("config-override")
.long("config-override")
.short('o')
.action(ArgAction::Append),
);
let mut commands = BTreeMap::new();
for command in self.commands.values() {
let subcmd = command.command();
commands.insert(subcmd.get_name().to_owned(), command);
cli = cli.subcommand(subcmd);
}
cli
}
pub async fn run_main(&self, app: Arc<App>, mut matches: ArgMatches) -> ExitCode {
let (name, matches) = matches.remove_subcommand().unwrap();
let command = self
.commands
.values()
.find(|v| v.command().get_name() == name)
.unwrap();
command.main(app, matches).await
}
pub fn len(&self) -> usize {
self.commands.len()
}
pub fn is_empty(&self) -> bool {
self.commands.is_empty()
}
}
struct CommandWrapper<T>(PhantomData<T>)
where
T: Command;
impl<T> Command for CommandWrapper<T>
where
T: Command,
{
fn command() -> clap::Command
where
Self: Sized,
{
T::command()
}
async fn main(app: Arc<App>, matches: ArgMatches) -> ExitCode {
T::main(app, matches).await
}
}
pub trait AddCommandExt {
fn add_command<T>(&mut self) -> &mut Self
where
T: Command + 'static;
fn has_command<T>(&self) -> bool
where
T: Command + 'static;
}
impl AddCommandExt for AppBuilder {
fn add_command<T>(&mut self) -> &mut Self
where
T: Command + 'static,
{
if !self.has_component::<CommandRegistry>() {
self.add_component(CommandRegistry::default());
}
self.get_component_mut::<CommandRegistry>()
.unwrap()
.add_command::<T>();
self
}
fn has_command<T>(&self) -> bool
where
T: Command + 'static,
{
self.get_component_ref::<CommandRegistry>()
.is_some_and(|v| v.has_command::<T>())
}
}
pub trait RunMainExt {
fn run_main(&mut self) -> impl std::future::Future<Output = ExitCode> + Send;
}
impl RunMainExt for AppBuilder {
async fn run_main(&mut self) -> ExitCode {
if !self.has_command::<ServerCommand>() {
self.add_command::<ServerCommand>();
}
if !self.has_command::<ConfigCommand>() {
self.add_command::<ConfigCommand>();
}
let command_registry = take(&mut *self.get_component_mut::<CommandRegistry>().unwrap());
let cli = command_registry.build_cli();
let matches = cli.get_matches();
if !self.has_component::<Config>() {
let config_path = matches.get_one::<String>("config").unwrap();
let mut config = Config::parse_file(config_path).await.unwrap();
let config_override_paths = matches
.get_many::<String>("config-override")
.unwrap_or_default();
for path in config_override_paths {
let config_override = Config::parse_file(path).await.unwrap();
config.merge_from(config_override).unwrap();
}
self.add_component(config);
}
Tracing::build(&*self).unwrap();
Metrics::build(&*self).unwrap();
let app = Arc::new(self.build().await.unwrap());
command_registry.run_main(app, matches).await
}
}
pub struct ServerCommand;
impl Command for ServerCommand {
fn command() -> clap::Command
where
Self: Sized,
{
clap::Command::new("server")
}
async fn main(app: Arc<App>, _matches: ArgMatches) -> ExitCode {
let shutdown = CancellationToken::new();
tokio::spawn({
let shutdown = shutdown.clone();
async move {
tokio::signal::ctrl_c()
.await
.expect("Failed to listen for ctrl_c");
shutdown.cancel();
}
});
if let Err(err) = app.run_daemons(shutdown).await {
panic!("Failed to run server: {err}");
}
ExitCode::SUCCESS
}
}
pub struct ConfigCommand;
impl Command for ConfigCommand {
fn command() -> clap::Command
where
Self: Sized,
{
clap::Command::new("config")
}
async fn main(app: Arc<App>, _matches: ArgMatches) -> ExitCode {
let config = app.get_component_ref::<Config>().unwrap();
println!("{}", serde_json::to_string_pretty(&config.configs).unwrap());
ExitCode::SUCCESS
}
}