use crate::Config;
use crate::tls::TlsConfig;
use axum::*;
use axum_server::tls_rustls::RustlsAcceptor;
use http::{HeaderName, HeaderValue, Method};
use std::net::SocketAddr;
use std::str::FromStr;
use std::sync::Arc;

pub struct HttpResponse {
    pub status: u16,
    pub body: Vec<u8>,
}

impl HttpResponse {
    fn respond(self) -> response::Response {
        response::Response::builder()
            .status(self.status)
            .header("Content-Type", "application/json")
            .body(body::Body::from(self.body))
            .expect("failed to encode response")
    }
}

pub type HttpRespondCb = Box<dyn FnOnce(HttpResponse) + 'static + Send>;

pub enum HttpRequest {
    HealthGet,
    BootstrapGet {
        space: bytes::Bytes,
    },
    BootstrapPut {
        space: bytes::Bytes,
        agent: bytes::Bytes,
        body: bytes::Bytes,
    },
}

type HSend = async_channel::Sender<(HttpRequest, HttpRespondCb)>;
type HRecv = async_channel::Receiver<(HttpRequest, HttpRespondCb)>;

#[derive(Clone)]
pub struct HttpReceiver(HRecv);

impl HttpReceiver {
    pub fn recv(&self) -> Option<(HttpRequest, HttpRespondCb)> {
        self.0.recv_blocking().ok()
    }
}

pub struct ServerConfig {
    pub addrs: Vec<std::net::SocketAddr>,
    pub worker_thread_count: usize,
    pub tls_config: Option<TlsConfig>,
    #[cfg(feature = "iroh-relay")]
    pub quic_bind_addr: Option<std::net::SocketAddr>,
}

pub struct Server {
    t_join: Option<std::thread::JoinHandle<()>>,
    addrs: Vec<std::net::SocketAddr>,
    receiver: HttpReceiver,
    h_send: HSend,
    tls_reload_handle: Option<tokio::task::AbortHandle>,
    shutdown: Option<axum_server::Handle>,
    auth_tracker: crate::auth::AuthTokenTracker,
    #[cfg(feature = "iroh-relay")]
    relay_allowlist: Option<crate::RelayAllowlist>,
    /// Held to keep the QAD server task alive; dropped on shutdown.
    #[cfg(feature = "iroh-relay")]
    _qad_server: Option<iroh_relay::server::Server>,
}

impl Drop for Server {
    fn drop(&mut self) {
        self.h_send.close();
        if let Some(shutdown) = self.shutdown.take() {
            shutdown.shutdown();
        }
        if let Some(t_join) = self.t_join.take() {
            let _ = t_join.join();
        }
        if let Some(tls_reload_handle) = self.tls_reload_handle.take() {
            tls_reload_handle.abort();
        }
    }
}

impl Server {
    pub fn new(
        config: Arc<Config>,
        server_config: ServerConfig,
    ) -> std::io::Result<Self> {
        let (s_ready, r_ready) = tokio::sync::oneshot::channel();
        let t_join = std::thread::spawn(move || {
            tokio_thread(config, server_config, s_ready)
        });
        match r_ready.blocking_recv() {
            Ok(Ok(Ready {
                h_send,
                addrs,
                receiver,
                tls_reload_handle,
                shutdown,
                auth_tracker,
                #[cfg(feature = "iroh-relay")]
                relay_allowlist,
                #[cfg(feature = "iroh-relay")]
                qad_server,
            })) => Ok(Self {
                t_join: Some(t_join),
                addrs,
                receiver,
                h_send,
                tls_reload_handle,
                shutdown: Some(shutdown),
                auth_tracker,
                #[cfg(feature = "iroh-relay")]
                relay_allowlist,
                #[cfg(feature = "iroh-relay")]
                _qad_server: qad_server,
            }),
            Ok(Err(err)) => Err(err),
            Err(_) => Err(std::io::Error::other("failed to bind server")),
        }
    }

    pub fn server_addrs(&self) -> &[std::net::SocketAddr] {
        self.addrs.as_slice()
    }

    pub fn receiver(&self) -> &HttpReceiver {
        &self.receiver
    }

    pub fn auth_tracker(&self) -> &crate::auth::AuthTokenTracker {
        &self.auth_tracker
    }

    #[cfg(feature = "iroh-relay")]
    pub fn relay_allowlist(&self) -> Option<&crate::RelayAllowlist> {
        self.relay_allowlist.as_ref()
    }
}

struct Ready {
    h_send: HSend,
    addrs: Vec<std::net::SocketAddr>,
    receiver: HttpReceiver,
    tls_reload_handle: Option<tokio::task::AbortHandle>,
    shutdown: axum_server::Handle,
    auth_tracker: crate::auth::AuthTokenTracker,
    #[cfg(feature = "iroh-relay")]
    relay_allowlist: Option<crate::RelayAllowlist>,
    #[cfg(feature = "iroh-relay")]
    qad_server: Option<iroh_relay::server::Server>,
}

#[derive(Clone)]
pub struct AppState {
    pub h_send: HSend,

    // Feature-independent authentication
    pub auth_tracker: crate::auth::AuthTokenTracker,
    pub auth_config: Arc<crate::auth::AuthConfig>,
    pub auth_failures: opentelemetry::metrics::Counter<u64>,

    // SBD-specific (keep for SBD websockets)
    #[cfg(feature = "sbd")]
    pub sbd_state: Option<crate::sbd::SbdState>,

    // Iroh relay allowlist: Some when authentication hook server is configured
    #[cfg(feature = "iroh-relay")]
    pub relay_allowlist: Option<crate::RelayAllowlist>,
}

type BoxFut<'a, T> =
    std::pin::Pin<Box<dyn std::future::Future<Output = T> + Send + 'a>>;

fn tokio_thread(
    config: Arc<Config>,
    server_config: ServerConfig,
    ready: tokio::sync::oneshot::Sender<std::io::Result<Ready>>,
) {
    tracing::trace!(?config, "Starting tokio thread");

    let allowed_headers =
        match ["Authorization", "Content-Type", "Content-Length", "Accept"]
            .iter()
            .map(|h| HeaderName::from_str(h))
            .collect::<Result<Vec<_>, _>>()
        {
            Ok(values) => tower_http::cors::AllowHeaders::list(values),
            Err(err) => {
                if ready.send(Err(std::io::Error::other(err))).is_err() {
                    tracing::error!("Failed to send ready error");
                }
                return;
            }
        };

    let auth_enabled = config.auth.authentication_hook_server.is_some();

    // When specific origins are listed, use them and allow credentials (if auth is on).
    // When no origins are listed and auth is on, mirror the request origin back — this
    // satisfies the CORS spec (credentials cannot be paired with `*`) while still
    // accepting requests from any origin, including browsers.
    // When auth is off, allow any origin without credentials.
    let (origin, allow_credentials) = match &config.allowed_origins {
        Some(origins) if !origins.is_empty() => {
            match origins
                .iter()
                .map(|o| HeaderValue::from_str(o.as_str()))
                .collect::<Result<Vec<_>, _>>()
            {
                Ok(values) => {
                    (tower_http::cors::AllowOrigin::list(values), auth_enabled)
                }
                Err(err) => {
                    if ready.send(Err(std::io::Error::other(err))).is_err() {
                        tracing::error!("Failed to send ready error");
                    }
                    return;
                }
            }
        }
        _ if auth_enabled => {
            (tower_http::cors::AllowOrigin::mirror_request(), true)
        }
        _ => (tower_http::cors::AllowOrigin::any(), false),
    };

    tokio::runtime::Builder::new_multi_thread()
        .enable_all()
        .build()
        .unwrap()
        .block_on(async move {
            let (h_send, h_recv) =
                async_channel::bounded(server_config.worker_thread_count);

            #[cfg(feature = "sbd")]
            let sbd_config = Arc::new(config.sbd.clone());

            #[cfg(feature = "sbd")]
            let ip_rate = Arc::new(sbd_server::IpRate::new(sbd_config.clone()));

            #[cfg(feature = "sbd")]
            let sbd_server_meter = opentelemetry::global::meter("sbd-server");

            #[cfg(feature = "sbd")]
            let c_slot = if config.no_relay_server {
                None
            } else {
                Some(sbd_server::CSlot::new(sbd_config.clone(), ip_rate.clone(), sbd_server_meter.clone()))
            };

            let (rustls_config, tls_reload_handle) = if let Some(tls_config) = server_config.tls_config {
                let rustls_config = tls_config
                    .create_tls_config()
                    .await
                    .expect("Failed to create TLS config");

                let tls_reload_handle = tokio::task::spawn(tls_config.reload_task(rustls_config.clone())).abort_handle();

                (Some(rustls_config), Some(tls_reload_handle))
            } else {
                (None, None)
            };

            // Initialize feature-independent authentication
            let meter = opentelemetry::global::meter("bootstrap-auth");
            let auth_failures = meter
                .u64_counter("bootstrap.auth_failures")
                .with_description("Number of failed authentication attempts")
                .with_unit("count")
                .build();
            let auth_tracker = crate::auth::AuthTokenTracker::default();
            let auth_config = Arc::new(config.auth.clone());

            let app = Router::<AppState>::new()
                .route("/authenticate", routing::put(handle_auth))
                .route("/health", routing::get(handle_health_get))
                .route("/bootstrap/{space}", routing::get(handle_boot_get))
                .route(
                    "/bootstrap/{space}/{agent}",
                    routing::put(handle_boot_put),
                );

            let mut app = app
                .layer(tower_http::cors::CorsLayer::new()
                    .allow_methods(tower_http::cors::AllowMethods::list([Method::GET, Method::PUT, Method::OPTIONS]))
                    .allow_headers(allowed_headers)
                    .allow_origin(origin)
                    .allow_credentials(allow_credentials)
                );

            #[cfg(feature = "iroh-relay")]
            let relay_allowlist: Option<crate::RelayAllowlist> =
                if config.auth.authentication_hook_server.is_some() {
                    Some(crate::RelayAllowlist::default())
                } else {
                    None
                };

            if !config.no_relay_server {
                #[cfg(feature = "sbd")]
                {
                    app = app.route("/{pub_key}", routing::get(crate::sbd::handle_sbd));
                }
                #[cfg(feature = "iroh-relay")]
                {
                    app = app.route(
                        "/relay/register",
                        routing::put(handle_relay_register),
                    );

                    let relay_state =
                        if let Some(allowlist) = relay_allowlist.clone() {
                            crate::iroh_relay_axum::create_relay_state_with_allowlist(allowlist)
                        } else {
                            crate::iroh_relay_axum::create_relay_state()
                        };
                    let relay_router = Router::new()
                        .route("/relay", routing::get(crate::iroh_relay_axum::relay_handler))
                        .route("/ping", routing::get(crate::iroh_relay_axum::relay_probe_handler))
                        .route("/generate_204", routing::get(crate::iroh_relay_axum::captive_portal_handler))
                        .with_state(relay_state);

                    app = app.merge(relay_router);
                }
            }

            // Clone auth_tracker before moving it into AppState so we can return it
            let auth_tracker_for_ready = auth_tracker.clone();
            #[cfg(feature = "iroh-relay")]
            let relay_allowlist_for_ready = relay_allowlist.clone();

            let app: Router = app
                .layer(extract::DefaultBodyLimit::max(1024))
                .with_state(AppState {
                    h_send: h_send.clone(),
                    auth_tracker,
                    auth_config,
                    auth_failures,
                    #[cfg(feature = "sbd")]
                    sbd_state: if config.no_relay_server {
                        None
                    } else {
                        {
                            Some(crate::sbd::SbdState {
                                config: sbd_config.clone(),
                                ip_rate: ip_rate.clone(),
                                c_slot: c_slot.as_ref().expect("Missing c_slot with SBD enabled").weak(),
                            })
                        }
                    },
                    #[cfg(feature = "iroh-relay")]
                    relay_allowlist,
                });

            let receiver = HttpReceiver(h_recv);

            let mut addrs = Vec::with_capacity(server_config.addrs.len());
            let mut servers: Vec<BoxFut<'static, std::io::Result<()>>> =
                Vec::with_capacity(server_config.addrs.len());

            let shutdown_handle = axum_server::Handle::new();

            for addr in server_config.addrs {
                tracing::info!("Binding to: {}", addr);

                let listener = match tokio::task::spawn_blocking(move || {
                    std::net::TcpListener::bind(addr)
                })
                .await
                .expect("Failed to run bind task")
                {
                    Ok(listener) => listener,
                    Err(err) => {
                        let _ = ready.send(Err(err));
                        return;
                    }
                };

                match listener.local_addr() {
                    Ok(addr) => {
                        tracing::info!("Bound with local address: {}", addr);
                        addrs.push(addr)
                    },
                    Err(err) => {
                        let _ = ready.send(Err(err));
                        return;
                    }
                }

                let app = app.clone();
                let shutdown_handle = shutdown_handle.clone();
                if let Some(tls_config) = &rustls_config {
                    let acceptor = RustlsAcceptor::new(tls_config.clone());

                    #[cfg(feature = "sbd")]
                    let acceptor = {
                        acceptor.acceptor(crate::sbd::SbdAcceptor::new(
                                sbd_config.clone(),
                                ip_rate.clone(),
                            ))
                    };

                    let s = axum_server::Server::from_tcp(listener)
                        .acceptor(acceptor)
                        .handle(shutdown_handle)
                        .serve(app.into_make_service_with_connect_info::<SocketAddr>());

                    servers.push(Box::pin(s));
                } else {
                    let server = axum_server::Server::from_tcp(listener);

                    #[cfg(feature = "sbd")]
                    let server = {
                        server.acceptor(crate::sbd::SbdAcceptor::new(
                            sbd_config.clone(),
                            ip_rate.clone(),
                        ))
                    };

                    let s = std::future::IntoFuture::into_future(
                        server
                            .handle(shutdown_handle)
                            .serve(app.into_make_service_with_connect_info::<SocketAddr>()),
                    );
                    servers.push(Box::pin(s));
                };
            }

            // Spawn the QUIC Address Discovery (QAD) server if configured
            #[cfg(feature = "iroh-relay")]
            let qad_server = if let Some(quic_bind_addr) = server_config.quic_bind_addr {
                let result = if let (Some(cert_path), Some(key_path)) = (&config.tls_cert, &config.tls_key) {
                    crate::qad::spawn_qad_server(quic_bind_addr, cert_path, key_path).await
                } else {
                    tracing::info!("No TLS cert/key configured, using self-signed certificate for QAD");
                    crate::qad::spawn_qad_server_self_signed(quic_bind_addr).await
                };
                match result {
                    Ok(server) => {
                        tracing::info!(
                            addr = ?server.quic_addr(),
                            "QAD (QUIC Address Discovery) server started"
                        );
                        Some(server)
                    }
                    Err(err) => {
                        let _ = ready.send(Err(std::io::Error::other(format!(
                            "failed to start QAD server: {err}"
                        ))));
                        return;
                    }
                }
            } else {
                None
            };

            tracing::info!("Sending ready signal");

            if ready
                .send(Ok(Ready {
                    h_send,
                    addrs,
                    receiver,
                    tls_reload_handle,
                    shutdown: shutdown_handle,
                    auth_tracker: auth_tracker_for_ready,
                    #[cfg(feature = "iroh-relay")]
                    relay_allowlist: relay_allowlist_for_ready,
                    #[cfg(feature = "iroh-relay")]
                    qad_server,
                }))
                .is_err()
            {
                return;
            }

            let _ = futures::future::join_all(servers).await;
        });
}

async fn handle_auth(
    extract::State(state): extract::State<AppState>,
    body: bytes::Bytes,
) -> axum::response::Response {
    match crate::auth::process_authenticate(
        &state.auth_config,
        &state.auth_tracker,
        state.auth_failures,
        body,
    )
    .await
    {
        Ok(token) => axum::response::IntoResponse::into_response(axum::Json(
            serde_json::json!({
                "authToken": *token,
            }),
        )),
        Err(crate::auth::AuthenticateError::Unauthorized) => {
            tracing::debug!("/authenticate: UNAUTHORIZED");
            axum::response::IntoResponse::into_response((
                axum::http::StatusCode::UNAUTHORIZED,
                "Unauthorized",
            ))
        }
        Err(crate::auth::AuthenticateError::Pending) => {
            tracing::debug!("/authenticate: PENDING");
            axum::response::IntoResponse::into_response(
                axum::http::StatusCode::ACCEPTED,
            )
        }
        Err(crate::auth::AuthenticateError::HookServerError(err)) => {
            tracing::debug!(?err, "/authenticate: BAD_GATEWAY");
            axum::response::IntoResponse::into_response((
                axum::http::StatusCode::BAD_GATEWAY,
                format!("BAD_GATEWAY: {err:?}"),
            ))
        }
        Err(crate::auth::AuthenticateError::OtherError(err)) => {
            tracing::warn!(?err, "/authenticate: INTERNAL_SERVER_ERROR");
            axum::response::IntoResponse::into_response((
                axum::http::StatusCode::INTERNAL_SERVER_ERROR,
                format!("INTERNAL_SERVER_ERROR: {err:?}"),
            ))
        }
    }
}

async fn handle_dispatch(
    h_send: &HSend,
    req: HttpRequest,
) -> response::Response {
    let (s, r) = tokio::sync::oneshot::channel();
    let s = Box::new(move |res| {
        let _ = s.send(res);
    });
    tokio::time::timeout(std::time::Duration::from_secs(10), async move {
        let _ = h_send.send((req, s)).await;
        match r.await {
            Ok(r) => r.respond(),
            Err(_) => HttpResponse {
                status: 500,
                body: b"{\"error\":\"request dropped\"}".to_vec(),
            }
            .respond(),
        }
    })
    .await
    .unwrap_or_else(|_| {
        HttpResponse {
            status: 500,
            body: b"{\"error\":\"internal timeout\"}".to_vec(),
        }
        .respond()
    })
}

async fn handle_health_get(
    extract::State(state): extract::State<AppState>,
) -> response::Response {
    // NOTE - This health call is currently not protected by auth token.
    //        Perhaps in the future, this should be configurable so
    //        infrastructure maintainers can weigh the trade-offs.

    handle_dispatch(&state.h_send, HttpRequest::HealthGet).await
}

async fn handle_boot_get(
    extract::Path(space): extract::Path<String>,
    headers: axum::http::HeaderMap,
    extract::State(state): extract::State<AppState>,
) -> response::Response {
    // Check authentication (feature-independent)
    let token: Option<Arc<str>> = headers
        .get("Authorization")
        .and_then(|t| t.to_str().ok())
        .and_then(|t| t.strip_prefix("Bearer "))
        .map(<Arc<str>>::from);

    if !state.auth_tracker.is_valid(&token, &state.auth_config) {
        return axum::response::IntoResponse::into_response((
            axum::http::StatusCode::UNAUTHORIZED,
            "Unauthorized",
        ));
    }

    let space = match b64_to_bytes(&space) {
        Ok(space) => space,
        Err(err) => return *err,
    };
    handle_dispatch(&state.h_send, HttpRequest::BootstrapGet { space }).await
}

async fn handle_boot_put(
    extract::Path((space, agent)): extract::Path<(String, String)>,
    headers: axum::http::HeaderMap,
    extract::State(state): extract::State<AppState>,
    body: bytes::Bytes,
) -> response::Response<body::Body> {
    // Check authentication (feature-independent)
    let token: Option<Arc<str>> = headers
        .get("Authorization")
        .and_then(|t| t.to_str().ok())
        .and_then(|t| t.strip_prefix("Bearer "))
        .map(<Arc<str>>::from);

    if !state.auth_tracker.is_valid(&token, &state.auth_config) {
        return axum::response::IntoResponse::into_response((
            axum::http::StatusCode::UNAUTHORIZED,
            "Unauthorized",
        ));
    }

    let space = match b64_to_bytes(&space) {
        Ok(space) => space,
        Err(err) => return *err,
    };
    let agent = match b64_to_bytes(&agent) {
        Ok(agent) => agent,
        Err(err) => return *err,
    };
    handle_dispatch(
        &state.h_send,
        HttpRequest::BootstrapPut { space, agent, body },
    )
    .await
}

#[cfg(feature = "iroh-relay")]
async fn handle_relay_register(
    extract::State(state): extract::State<AppState>,
    headers: axum::http::HeaderMap,
    body: bytes::Bytes,
) -> response::Response {
    tracing::debug!(body_len = body.len(), "Relay register request received");

    // Validate bearer token
    let token: Option<Arc<str>> = headers
        .get("Authorization")
        .and_then(|t| t.to_str().ok())
        .and_then(|t| t.strip_prefix("Bearer "))
        .map(<Arc<str>>::from);

    if !state.auth_tracker.is_valid(&token, &state.auth_config) {
        tracing::warn!("Relay register: bearer token invalid or missing");
        return axum::response::IntoResponse::into_response((
            axum::http::StatusCode::UNAUTHORIZED,
            "Unauthorized",
        ));
    }

    // Expect exactly 32 bytes (iroh public key)
    let key_bytes: &[u8; 32] = match body.as_ref().try_into() {
        Ok(b) => b,
        Err(_) => {
            tracing::warn!(
                body_len = body.len(),
                "Relay register: expected 32 bytes"
            );
            return axum::response::IntoResponse::into_response((
                axum::http::StatusCode::BAD_REQUEST,
                "Expected exactly 32 bytes (iroh public key)",
            ));
        }
    };

    let key = match iroh_base::PublicKey::from_bytes(key_bytes) {
        Ok(k) => k,
        Err(_) => {
            tracing::warn!("Relay register: invalid public key bytes");
            return axum::response::IntoResponse::into_response((
                axum::http::StatusCode::BAD_REQUEST,
                "Invalid public key bytes",
            ));
        }
    };

    // Register the key in the allowlist (if auth is enabled).
    // token must be Some here: is_valid returned true, and the allowlist is
    // only present when an auth hook server is configured, which requires a
    // bearer token for is_valid to succeed. Guard defensively anyway.
    if let Some(allowlist) = &state.relay_allowlist {
        let Some(token) = token else {
            return axum::response::IntoResponse::into_response((
                axum::http::StatusCode::INTERNAL_SERVER_ERROR,
                "Internal server error",
            ));
        };
        allowlist.register(key, token);
        tracing::info!(
            key = %key.fmt_short(),
            "Relay register: key added to allowlist"
        );
    } else {
        tracing::warn!(
            "Relay register: no allowlist configured, registration has no effect"
        );
    }

    axum::response::IntoResponse::into_response(axum::Json(serde_json::json!(
        {}
    )))
}

fn b64_to_bytes(
    s: &str,
) -> Result<bytes::Bytes, Box<response::Response<body::Body>>> {
    use base64::prelude::*;
    Ok(bytes::Bytes::copy_from_slice(
        &match BASE64_URL_SAFE_NO_PAD.decode(s) {
            Ok(b) => b,
            Err(err) => {
                return Err(Box::new(
                    HttpResponse {
                        status: 400,
                        body: err.to_string().into_bytes(),
                    }
                    .respond(),
                ));
            }
        },
    ))
}
