neiam /dms

use std::sync::Arc;
use std::time::Duration;
use axum::extract::{DefaultBodyLimit, State};
use axum::http::StatusCode;
use axum::response::Html;
use axum::Router;
use axum::routing::{get, post};
use axum_htmx::HxRequest;
use axum_prometheus::PrometheusMetricLayer;
use log::info;
use metrics_process::Collector;
use minijinja::{context, Environment};
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 dms::{handlers, middleware, AppState, Config, create_db_connection};
use dms::middleware::{web_auth_middleware, UserContext};

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

    let config = Arc::new(Config::from_env().expect("Failed to load configuration"));
    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?;
        info!("Migrations up to date");
    } else {
        info!("AUTO_MIGRATE=false — skipping migrations");
    }

    let collector = Collector::default();
    collector.describe();
    let (prometheus_layer, metric_handle) = PrometheusMetricLayer::pair();

    let mut jinja = Environment::new();
    minijinja_embed::load_templates!(&mut jinja);
    jinja.set_loader(minijinja::path_loader("templates"));

    let app_state = AppState {
        env: jinja,
        db: db.clone(),
        config: config.clone(),
    };

    // The landing page reads `Option<UserContext>` so it works for both
    // anonymous visitors and signed-in users — give it the optional-auth
    // middleware so the cookie populates the extension when present without
    // 302-ing anonymous visitors away.
    let public_routes = Router::new()
        .route("/", get(root))
        .layer(axum::middleware::from_fn_with_state(
            app_state.db.clone(),
            middleware::optional_auth_middleware,
        ));

    let protected_routes = Router::new()
        .route("/profile", get(handlers::show_profile))
        .route("/invites", get(handlers::show_invites).post(handlers::handle_create_invite))
        .route("/api-tokens", get(handlers::show_api_tokens).post(handlers::create_api_token))
        .route("/api-tokens/{id}/delete", post(handlers::delete_api_token))
        .route("/uploads", get(handlers::list_uploads).post(handlers::upload_file))
        .route("/uploads/url", post(handlers::upload_from_url))
        .route("/uploads/{id}", get(handlers::show_upload))
        .route("/uploads/{id}/jobs", post(handlers::enqueue_job))
        .route("/uploads/{id}/tags", post(handlers::add_tags))
        .route("/search", get(handlers::search_uploads))
        .route("/failed", get(handlers::show_failed))
        .route("/failed/retry", post(handlers::retry_all_failed))
        .route("/failed/dlq/requeue", post(handlers::requeue_all_dlqs_html))
        .route("/failed/dlq/{kind}/requeue", post(handlers::requeue_dlq_kind_html))
        .route("/uploads/{id}/jobs/{kind}/retry", post(handlers::retry_one))
        .layer(DefaultBodyLimit::max(config.max_upload_bytes))
        .layer(axum::middleware::from_fn_with_state(
            app_state.db.clone(),
            web_auth_middleware,
        ));

    let api_routes = Router::new()
        .route("/profile", get(api_profile))
        .route("/uploads", get(handlers::api_list_uploads).post(handlers::api_upload_file))
        .route("/uploads/url", post(handlers::api_upload_from_url))
        .route("/uploads/{id}", get(handlers::api_show_upload))
        .route("/uploads/{id}/tags", post(handlers::api_add_tags))
        .route("/search", get(handlers::api_search_uploads))
        .route("/failed", get(handlers::api_list_failed))
        .route("/failed/retry", post(handlers::api_retry_all_failed))
        .route("/failed/dlq/requeue", post(handlers::api_requeue_all_dlqs))
        .route("/failed/dlq/{kind}/requeue", post(handlers::api_requeue_dlq_kind))
        .layer(DefaultBodyLimit::max(config.max_upload_bytes))
        .layer(axum::middleware::from_fn_with_state(
            app_state.db.clone(),
            middleware::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(move || async move {
            collector.collect();
            metric_handle.render()
        }))
        .nest_service("/assets", ServeDir::new("assets/f/"))
        .nest_service("/files", ServeDir::new(&config.storage_root))
        .merge(public_routes)
        .merge(protected_routes)
        .nest("/api", api_routes)
        .with_state(app_state)
        .layer(
            TraceLayer::new_for_http()
                .make_span_with(DefaultMakeSpan::default().include_headers(true))
                .on_failure(
                    |error: ServerErrorsFailureClass, _latency: Duration, _span: &tracing::Span| {
                        tracing::error!("something went wrong {error}")
                    },
                ),
        )
        .layer(prometheus_layer);

    let listener = tokio::net::TcpListener::bind(&config.bind_address).await?;
    info!("Server starting on {}", config.bind_address);

    axum::serve(listener, app).await?;
    Ok(())
}

async fn api_profile(
    user_context: UserContext,
) -> Result<axum::Json<serde_json::Value>, StatusCode> {
    Ok(axum::Json(serde_json::json!({
        "user": {
            "id": user_context.user.id,
            "username": user_context.user.username,
            "email": user_context.user.email,
            "is_active": user_context.user.is_active,
            "created_at": user_context.user.created_at,
            "last_login_at": user_context.user.last_login_at,
        }
    })))
}

async fn root(
    State(state): State<AppState>,
    HxRequest(boosted): HxRequest,
    user_context: Option<UserContext>,
) -> Result<Html<String>, StatusCode> {
    let actual_disk_usage = 0;
    let user = user_context.map(|ctx| ctx.user);

    let template_str = if boosted { "_index.html" } else { "index.html" };

    let template = state.env.get_template(template_str)
        .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;

    let rendered = template
        .render(context! {
            actual_disk_usage => actual_disk_usage.to_string(),
            title => "Home",
            designator => "DMS",
            library_status => "✅",
            library_url => "Not configured",
            user => user,
            registration_open => state.config.registration_enabled,
        })
        .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;

    Ok(Html(rendered))
}