pub mod accounts;
pub mod books;
pub mod graphql_with_subscriptions;
pub mod inventory;
pub mod monolith;
pub mod products;
pub mod reviews;

use std::sync::{atomic::AtomicUsize, Arc};

use async_graphql_axum::{GraphQL, GraphQLSubscription};
use axum::{
    extract::Request,
    http::StatusCode,
    middleware::{self, Next},
    response::{IntoResponse, Response},
    routing::{get, post_service},
    Router,
};
use tokio::{
    net::TcpListener,
    sync::oneshot::{self, Sender},
    task::JoinHandle,
};

async fn delay_middleware(req: Request, next: Next) -> Response {
    let delay_ms: Option<u64> = std::env::var("SUBGRAPH_DELAY_MS")
        .ok()
        .and_then(|s| s.parse().ok())
        .filter(|d| *d != 0);

    if let Some(delay_ms) = delay_ms {
        tokio::time::sleep(tokio::time::Duration::from_millis(delay_ms)).await;
    }

    next.run(req).await
}

async fn add_subgraph_header(req: Request, next: Next) -> Response {
    let path = req.uri().path();
    let subgraph_name = path.trim_start_matches('/').to_string();

    let mut response = next.run(req).await;

    if !subgraph_name.is_empty() && subgraph_name != "health" {
        if let Ok(header_value) = subgraph_name.parse() {
            response.headers_mut().insert("x-subgraph", header_value);
        }
    }

    response
}

async fn health_check_handler() -> impl IntoResponse {
    StatusCode::OK
}

pub fn start_subgraphs_server(port: Option<u16>) -> (JoinHandle<()>, Sender<()>) {
    let (shutdown_tx, shutdown_rx) = oneshot::channel::<()>();
    let host = std::env::var("HOST").unwrap_or("0.0.0.0".to_owned());
    let port = port
        .map(|v| v.to_string())
        .unwrap_or(std::env::var("PORT").unwrap_or("4200".to_owned()));

    let (mut app, _) = subgraphs_app(HTTPStreamingSubscriptionProtocol::default());
    app = app.route("/health", get(health_check_handler));

    println!("Starting server on http://{}:{}", host, port);

    let server_handle = tokio::spawn(async move {
        axum::serve(
            TcpListener::bind(&format!("{}:{}", host, port))
                .await
                .unwrap(),
            app,
        )
        .with_graceful_shutdown(async {
            shutdown_rx.await.ok();
            println!("Graceful shutdown signal received.");
        })
        .await
        .expect("failed to start subgraphs server");
    });

    (server_handle, shutdown_tx)
}

/// The protocol to use for GraphQL subscriptions over HTTP streaming.
/// It is purely the streaming HTTP protocol, other subscription protocols
/// are handled automatically through HTTP negotiation (like websocket upgrades
/// or http callbacks).
#[derive(Clone, Default)]
pub enum HTTPStreamingSubscriptionProtocol {
    #[default]
    PreferMultipartFallbackSse,
    MultipartOnly,
    SseOnly,
}

pub fn subgraphs_app(
    subscriptions_protocol: HTTPStreamingSubscriptionProtocol,
) -> (Router<()>, Arc<AtomicUsize>) {
    let (reviews_schema, active_subscriptions) = reviews::get_subgraph();
    let router = Router::new()
        .route(
            "/accounts",
            post_service(GraphQL::new(accounts::get_subgraph())),
        )
        .route("/books", post_service(GraphQL::new(books::get_subgraph())))
        .route(
            "/inventory",
            post_service(GraphQL::new(inventory::get_subgraph())),
        )
        .route(
            "/products",
            post_service(GraphQL::new(products::get_subgraph())),
        )
        .route_service(
            "/reviews/ws",
            GraphQLSubscription::new(reviews_schema.clone()),
        )
        .route(
            "/reviews",
            post_service(graphql_with_subscriptions::GraphQL::new(
                reviews_schema,
                subscriptions_protocol,
            )),
        )
        .route(
            "/monolith",
            post_service(GraphQL::new(monolith::get_schema())),
        )
        .route_layer(middleware::from_fn(add_subgraph_header))
        .route_layer(middleware::from_fn(delay_middleware));
    (router, active_subscriptions)
}
