neiam /fieri-forever

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

use axum::extract::State;
use axum::http::StatusCode;
use axum::response::Html;
use axum::routing::{get, post};
use axum::Router;
use log::info;
use metrics_process::Collector;
use minijinja::context;
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 fieri_eternal::middleware::{optional_auth_middleware, web_auth_middleware, UserContext};
use fieri_eternal::live::{self, Live};
use fieri_eternal::osm::OsmClient;
use fieri_eternal::{
    add_config_globals, 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(|_| {
            "fieri_eternal=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();

    // Relay outing changes from Postgres NOTIFY to this replica's SSE streams.
    let live = Live::new();
    live.spawn_listener(live::pool(&db));

    let mut env = build_environment();
    add_config_globals(&mut env, &config);
    let state = AppState {
        env,
        db: db.clone(),
        config: config.clone(),
        osm: Arc::new(OsmClient::new(&config.geocode_url, &config.poi_url, &config.osm_api_url)?),
        live,
    };

    // Landing page: reads Option<UserContext>, so anonymous and signed-in
    // visitors both work — it gets the non-redirecting optional-auth layer.
    let public = Router::new().route("/", get(root)).layer(
        axum::middleware::from_fn_with_state(db.clone(), optional_auth_middleware),
    );

    // 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("/areas", get(handlers::areas::list_areas).post(handlers::areas::create_area))
        .route("/areas/new", get(handlers::areas::new_area))
        .route("/areas/{id}", get(handlers::areas::show_area))
        .route("/areas/{id}/restaurants/new", get(handlers::restaurants::new_restaurant))
        .route(
            "/areas/{id}/plan",
            get(handlers::lunch::show_planner).post(handlers::lunch::make_plan),
        )
        .route("/areas/{id}/history", get(handlers::lunch::show_history))
        .route("/areas/{id}/rename", post(handlers::areas::rename_area))
        .route("/areas/{id}/members", get(handlers::members::show_members))
        .route("/areas/{id}/members/invite", post(handlers::members::invite))
        .route("/areas/{id}/members/{user_id}/role", post(handlers::members::set_role))
        .route("/areas/{id}/members/{user_id}/remove", post(handlers::members::remove_member))
        .route(
            "/areas/{id}/invitations/{user_id}/cancel",
            post(handlers::members::cancel_invitation),
        )
        .route("/invitations/{area_id}/{answer}", post(handlers::members::answer_invitation))
        .route("/friends", get(handlers::friends::show_friends).post(handlers::friends::request))
        .route("/friends/{user_id}/{action}", post(handlers::friends::answer))
        .route("/areas/{id}/refresh", post(handlers::areas::refresh_area))
        .route("/cuisines", get(handlers::cuisines::list_cuisines))
        .route("/cuisines/{name}", get(handlers::cuisines::show_cuisine))
        .route("/areas/{id}/outings", post(handlers::outings::create_outing))
        .route("/outings/{id}", get(handlers::outings::show_outing))
        .route("/outings/{id}/panel", get(handlers::outings::panel))
        .route("/outings/{id}/events", get(handlers::outings::events))
        .route("/outings/{id}/join", post(handlers::outings::join_outing))
        .route("/outings/{id}/leave", post(handlers::outings::leave_outing))
        .route("/outings/{id}/vote", post(handlers::outings::vote))
        .route("/outings/{id}/reroll", post(handlers::outings::reroll))
        .route("/outings/{id}/decide", post(handlers::outings::decide))
        .route("/outings/{id}/cancel", post(handlers::outings::cancel))
        .route("/places/search", get(handlers::places::search))
        .route("/places/nearby", get(handlers::places::nearby))
        .route("/restaurants", post(handlers::restaurants::create_restaurant))
        .route(
            "/restaurants/{id}",
            get(handlers::restaurants::show_restaurant).post(handlers::restaurants::update_restaurant),
        )
        .route("/restaurants/{id}/rank", post(handlers::restaurants::set_rank))
        .route("/restaurants/{id}/archive", post(handlers::restaurants::toggle_archive))
        .route("/restaurants/{id}/refresh", post(handlers::restaurants::refresh_restaurant))
        .route("/restaurants/{id}/similar", get(handlers::restaurants::similar))
        .route("/lunches", post(handlers::lunch::log_lunch))
        .route("/lunches/{id}/delete", post(handlers::lunch::delete_lunch))
        .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(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).with_graceful_shutdown(shutdown_signal()).await?;
    info!("shut down");
    Ok(())
}

/// Resolves on SIGTERM (what Kubernetes and podman send) or Ctrl-C. As PID 1
/// in a container the process gets no default signal handling, so without
/// this it ignores SIGTERM and is killed after the grace period instead.
///
/// axum then stops accepting and lets in-flight requests finish — but live
/// vote streams (SSE) never finish on their own, so give everything a few
/// seconds and then exit. Browsers' EventSource reconnects by itself, to
/// whichever pod is up.
async fn shutdown_signal() {
    let ctrl_c = async {
        let _ = tokio::signal::ctrl_c().await;
    };
    #[cfg(unix)]
    let term = async {
        match tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate()) {
            Ok(mut s) => {
                s.recv().await;
            }
            Err(_) => std::future::pending::<()>().await,
        }
    };
    #[cfg(not(unix))]
    let term = std::future::pending::<()>();
    tokio::select! {
        _ = ctrl_c => {},
        _ = term => {},
    }
    info!("shutting down: draining requests for up to {}s", SHUTDOWN_GRACE.as_secs());
    tokio::spawn(async {
        tokio::time::sleep(SHUTDOWN_GRACE).await;
        info!("grace period over; exiting with streams still open");
        std::process::exit(0);
    });
}

/// Well inside Kubernetes' default 30 s terminationGracePeriodSeconds.
const SHUTDOWN_GRACE: Duration = Duration::from_secs(5);

/// Landing page. htmx `hx-boost` swaps the `#content` region on navigation,
/// so this can just render the whole page every time.
async fn root(
    State(state): State<AppState>,
    user: Option<UserContext>,
) -> Result<Html<String>, StatusCode> {
    let user = user.map(|c| c.user);
    // Signed in, the landing page is your areas, plus anything waiting on
    // you: area invitations and friend requests.
    let (areas, invitations, friend_requests) = match &user {
        Some(u) => (
            handlers::areas::my_areas(state.db.as_ref(), u.id).await?,
            handlers::members::pending_for(state.db.as_ref(), u.id).await?,
            handlers::friends::incoming_count(state.db.as_ref(), u.id).await?,
        ),
        None => (vec![], vec![], 0),
    };
    let template = state
        .env
        .get_template("index.html")
        .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
    let body = template
        .render(context! {
            title => "Home",
            user => user,
            areas => areas,
            invitations => invitations,
            friend_requests => friend_requests,
            registration_open => state.config.registration_enabled,
        })
        .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
    Ok(Html(body))
}