use axum::{
    error_handling::HandleErrorLayer,
    http::Uri,
    response::IntoResponse,
    routing::{any, get},
    Router,
};
use axum_oidc::{
    error::MiddlewareError, handle_oidc_redirect, openidconnect::{Audience, ClientId, ClientSecret},
    EmptyAdditionalClaims, OidcAuthLayer, OidcClaims, OidcClient, OidcLoginLayer,
    OidcRpInitiatedLogout,
};
use tokio::net::TcpListener;
use tower::ServiceBuilder;
use tower_sessions::{
    cookie::{time::Duration, SameSite},
    Expiry, MemoryStore, SessionManagerLayer,
};
use tracing::Level;

#[tokio::main]
async fn main() {
    tracing_subscriber::fmt()
        .with_file(true)
        .with_line_number(true)
        .with_max_level(Level::INFO)
        .init();
    dotenvy::dotenv().ok();
    let issuer = std::env::var("ISSUER").expect("ISSUER env variable");
    let client_id = std::env::var("CLIENT_ID").expect("CLIENT_ID env variable");
    let client_secret = std::env::var("CLIENT_SECRET").ok();

    let session_store = MemoryStore::default();
    let session_layer = SessionManagerLayer::new(session_store)
        .with_secure(false)
        .with_same_site(SameSite::Lax)
        .with_expiry(Expiry::OnInactivity(Duration::seconds(120)));

    let oidc_login_service = ServiceBuilder::new()
        .layer(HandleErrorLayer::new(|e: MiddlewareError| async {
            dbg!(&e);
            e.into_response()
        }))
        .layer(OidcLoginLayer::<EmptyAdditionalClaims>::new());

    let mut oidc_client = OidcClient::<EmptyAdditionalClaims>::builder()
        .with_default_http_client()
        .with_redirect_url(Uri::from_static("http://localhost:8080/oidc"))
        .with_client_id(ClientId::new(client_id))
        .add_scope("profile")
        .add_scope("email")
        // Optional: add untrusted audiences. If the `aud` claim contains any of these audiences, the token is rejected.
        .add_untrusted_audience(Audience::new("123456789".to_string()));

    if let Some(client_secret) = client_secret {
        oidc_client = oidc_client.with_client_secret(ClientSecret::new(client_secret));
    }
    let oidc_client = oidc_client.discover(issuer).await.unwrap().build();

    let oidc_auth_service = ServiceBuilder::new()
        .layer(HandleErrorLayer::new(|e: MiddlewareError| async {
            dbg!(&e);
            e.into_response()
        }))
        .layer(OidcAuthLayer::new(oidc_client));

    let app = Router::new()
        .route("/foo", get(authenticated))
        .route("/logout", get(logout))
        .layer(oidc_login_service)
        .route("/bar", get(maybe_authenticated))
        .route("/oidc", any(handle_oidc_redirect::<EmptyAdditionalClaims>))
        .layer(oidc_auth_service)
        .layer(session_layer);

    tracing::info!("Running on http://localhost:8080");
    tracing::info!("Visit http://localhost:8080/bar or http://localhost:8080/foo");

    let listener = TcpListener::bind("[::]:8080").await.unwrap();
    axum::serve(listener, app.into_make_service())
        .await
        .unwrap();
}

async fn authenticated(claims: OidcClaims<EmptyAdditionalClaims>) -> impl IntoResponse {
    format!("Hello {}", claims.subject().as_str())
}

#[axum::debug_handler]
async fn maybe_authenticated(
    claims: Result<OidcClaims<EmptyAdditionalClaims>, axum_oidc::error::ExtractorError>,
) -> impl IntoResponse {
    if let Ok(claims) = claims {
        format!(
            "Hello {}! You are already logged in from another Handler.",
            claims.subject().as_str()
        )
    } else {
        "Hello anon!".to_string()
    }
}

async fn logout(logout: OidcRpInitiatedLogout) -> impl IntoResponse {
    logout.with_post_logout_redirect(Uri::from_static("https://example.com"))
}