churn/crates/churn-server/src/main.rs

use std::collections::HashMap;
use std::net::SocketAddr;
use std::sync::Arc;

use anyhow::Error;
use axum::extract::State;
use axum::http::StatusCode;
use axum::response::{IntoResponse, Response};
use axum::routing::{get, post};
use axum::{async_trait, Json, Router};
use clap::{Parser, Subcommand};
use serde::{Deserialize, Serialize};
use serde_json::json;
use tokio::sync::Mutex;

#[derive(Parser)]
#[command(author, version, about, long_about = None, subcommand_required = true)]
struct Command {
    #[command(subcommand)]
    command: Option<Commands>,
}

#[derive(Subcommand)]
enum Commands {
    Serve {
        #[arg(env = "SERVICE_HOST", long, default_value = "127.0.0.1:3000")]
        host: SocketAddr,
    },
}

#[derive(Clone)]
struct AgentService(Arc<dyn AgentServiceTrait + Send + Sync + 'static>);

impl std::ops::Deref for AgentService {
    type Target = Arc<dyn AgentServiceTrait + Send + Sync + 'static>;

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

impl Default for AgentService {
    fn default() -> Self {
        Self(Arc::new(DefaultAgentService::default()))
    }
}

struct DefaultAgentService {
    agents: Arc<Mutex<HashMap<String, Agent>>>,
}

impl Default for DefaultAgentService {
    fn default() -> Self {
        Self {
            agents: Arc::default(),
        }
    }
}

#[async_trait]
trait AgentServiceTrait {
    async fn enroll(&self, agent: Agent) -> anyhow::Result<String>;
}

#[async_trait]
impl AgentServiceTrait for DefaultAgentService {
    async fn enroll(&self, agent: Agent) -> anyhow::Result<String> {
        let mut agents = self.agents.lock().await;

        match agents.insert(agent.name.clone(), agent.clone()) {
            Some(_) => {
                tracing::debug!("agents store already contained agent, replaced existing");

                Ok(agent.name)
            }
            None => {
                tracing::debug!("agents store didn't contain agent, inserted");

                Ok(agent.name)
            }
        }
    }
}

#[derive(Clone, Debug, Deserialize, Serialize)]
struct Agent {
    pub name: String,
}

#[derive(Clone)]
struct AppState {
    agent: AgentService,
}

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    dotenv::dotenv().ok();
    tracing_subscriber::fmt::init();

    let cli = Command::parse();

    match cli.command {
        Some(Commands::Serve { host }) => {
            tracing::info!("Starting churn server");

            let app = Router::new()
                .route("/ping", get(ping))
                .nest(
                    "/agent",
                    Router::new()
                        .route("/enroll", post(enroll))
                        .route("/ping", post(agent_ping))
                        .route("/events", post(get_tasks)),
                )
                .with_state(AppState {
                    agent: AgentService::default(),
                });

            tracing::info!("churn server listening on {}", host);
            axum::Server::bind(&host)
                .serve(app.into_make_service())
                .await
                .unwrap();
        }
        None => {}
    }

    Ok(())
}

enum AppError {
    Internal(Error),
}

impl IntoResponse for AppError {
    fn into_response(self) -> Response {
        let (status, error_message) = match self {
            AppError::Internal(e) => {
                tracing::error!("failed with error: {}", e);

                (
                    StatusCode::INTERNAL_SERVER_ERROR,
                    "failed with internal error",
                )
            }
        };

        let body = Json(json!({
            "error": error_message,
        }));

        (status, body).into_response()
    }
}

async fn enroll(
    State(state): State<AppState>,
    Json(agent): Json<Agent>,
) -> Result<Json<Agent>, AppError> {
    let _ = state
        .agent
        .enroll(agent.clone())
        .await
        .map_err(|e| AppError::Internal(e));

    Ok(Json(agent))
}

async fn agent_ping() -> impl IntoResponse {
    todo!()
}

async fn get_tasks() -> impl IntoResponse {
    todo!()
}

async fn ping() -> impl IntoResponse {
    "pong!"
}