neiam /bismuth

use std::sync::Arc;
use std::time::Duration;

use axum::routing::{get, post};
use axum::Router;
use log::info;
use metrics_process::Collector;
use tower_http::classify::ServerErrorsFailureClass;
use tower_http::services::ServeDir;
use tower_http::trace::{DefaultMakeSpan, TraceLayer};
use tracing_subscriber::layer::SubscriberExt;
use tracing_subscriber::util::SubscriberInitExt;

use axum_prometheus::PrometheusMetricLayer;
use bismuth::middleware::{optional_auth_middleware, web_auth_middleware};
use bismuth::{build_environment, create_db_connection, handlers, AppState, Config};

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    tracing_subscriber::registry()
        .with(tracing_subscriber::EnvFilter::try_from_default_env().unwrap_or_else(|_| {
            "bismuth=debug,webservice=debug,tower_http=info".into()
        }))
        .with(tracing_subscriber::fmt::layer())
        .init();

    let config = Arc::new(Config::from_env());
    info!("configuration loaded");

    let db = Arc::new(create_db_connection(&config.database_url).await?);
    info!("database connection established");

    if config.auto_migrate {
        use migration::MigratorTrait;
        info!("running pending migrations (set AUTO_MIGRATE=false to skip)");
        migration::Migrator::up(db.as_ref(), None).await?;
    }

    // Prometheus request metrics + process (RSS/CPU/fd) metrics at /metrics.
    let collector = Arc::new(Collector::default());
    collector.describe();
    let (prometheus_layer, metric_handle) = PrometheusMetricLayer::pair();

    let state = AppState {
        env: build_environment(),
        db: db.clone(),
        config: config.clone(),
    };

    // Pages that render for anyone but know who's signed in: the editor
    // (landing page) and a snippet's share page, which offers its owner
    // edit/delete. These get the non-redirecting optional-auth layer.
    let public = Router::new()
        .route("/", get(handlers::editor))
        .route("/s/{slug}", get(handlers::show_snippet))
        .layer(axum::middleware::from_fn_with_state(db.clone(), optional_auth_middleware));

    // Stateless rendering and the embeddable/image routes: no session needed.
    let open = Router::new()
        .route("/preview", post(handlers::preview))
        .route("/export", post(handlers::export))
        .route("/s/{slug}/card.svg", get(handlers::snippet_svg))
        .route("/s/{slug}/card.png", get(handlers::snippet_png))
        .route("/e/{slug}", get(handlers::embed));

    // Everything requiring a login. The web-auth layer 302s anonymous
    // visitors to /login.
    let protected = Router::new()
        .route("/profile", get(handlers::show_profile))
        .route("/invites", get(handlers::show_invites).post(handlers::handle_create_invite))
        .route("/snippets", get(handlers::list_snippets).post(handlers::create_snippet))
        .route("/snippets/{slug}", post(handlers::update_snippet))
        .route("/snippets/{slug}/edit", get(handlers::edit_snippet))
        .route("/snippets/{slug}/delete", post(handlers::delete_snippet))
        .layer(axum::middleware::from_fn_with_state(db.clone(), web_auth_middleware));

    let app = Router::new()
        .route("/login", get(handlers::show_login))
        .route("/register", get(handlers::show_register))
        .route("/auth/login", post(handlers::handle_login))
        .route("/auth/register", post(handlers::handle_register))
        .route("/auth/logout", post(handlers::handle_logout))
        .route(
            "/metrics",
            get({
                let collector = collector.clone();
                let handle = metric_handle.clone();
                move || {
                    let collector = collector.clone();
                    let handle = handle.clone();
                    async move {
                        collector.collect();
                        handle.render()
                    }
                }
            }),
        )
        .nest_service("/assets", ServeDir::new("assets/f"))
        .merge(public)
        .merge(open)
        .merge(protected)
        .with_state(state)
        .layer(
            TraceLayer::new_for_http()
                .make_span_with(DefaultMakeSpan::default().include_headers(false))
                .on_failure(
                    |error: ServerErrorsFailureClass, _latency: Duration, _span: &tracing::Span| {
                        tracing::error!("request failed: {error}");
                    },
                ),
        )
        .layer(prometheus_layer);

    let listener = tokio::net::TcpListener::bind(&config.bind_address).await?;
    info!("server listening on {}", config.bind_address);
    axum::serve(listener, app).await?;
    Ok(())
}