use std::sync::Arc;

use anyhow::anyhow;
use cas_types::{REQUEST_ID_HEADER, SESSION_ID_HEADER};
use error_printer::{ErrorPrinter, OptionPrinter};
use http::{Extensions, StatusCode};
use reqwest::header::{AUTHORIZATION, HeaderValue};
use reqwest::{Request, Response};
use reqwest_middleware::{ClientBuilder, ClientWithMiddleware, Middleware, Next};
use tokio::sync::Mutex;
use tracing::{Instrument, info, info_span, warn};
use utils::auth::{AuthConfig, TokenProvider};

use crate::{CasClientError, error};

#[cfg(not(target_family = "wasm"))]
fn reqwest_client(user_agent: &str) -> Result<reqwest::Client, CasClientError> {
    use xet_runtime::{XetRuntime, xet_config};

    let client = XetRuntime::get_or_create_reqwest_client(|| {
        let mut builder = reqwest::Client::builder()
            .pool_idle_timeout(xet_config().client.idle_connection_timeout)
            .pool_max_idle_per_host(xet_config().client.max_idle_connections)
            .connect_timeout(xet_config().client.connect_timeout)
            .read_timeout(xet_config().client.read_timeout)
            // Explicitly NOT setting .timeout() to disable transfer-level timeout.
            // We rely on packet-level timeouts (read_timeout) which reset when data
            // is received, allowing slow but progressing transfers to complete.
            .http1_only(); // high throughput parallel I/O has been shown to bottleneck with http2

        if !user_agent.is_empty() {
            builder = builder.user_agent(user_agent);
        }

        builder.build()
    })?;

    info!(
        idle_timeout=?xet_config().client.idle_connection_timeout,
        max_idle_connections=xet_config().client.max_idle_connections,
        user_agent=?if user_agent.is_empty() { None } else { Some(user_agent) },
        "HTTP client configured"
    );

    Ok(client)
}

#[cfg(target_family = "wasm")]
fn reqwest_client(user_agent: &str) -> Result<reqwest::Client, CasClientError> {
    // For WASM, create a new client with the specified user_agent
    // Note: we could cache this, but user_agent can vary, so we create per-call
    let mut builder = reqwest::Client::builder();
    if !user_agent.is_empty() {
        builder = builder.user_agent(user_agent);
    }
    Ok(builder.build()?)
}

/// Builds authenticated HTTP Client to talk to CAS.
pub fn build_auth_http_client(
    auth_config: &Option<AuthConfig>,
    session_id: &str,
    user_agent: &str,
) -> Result<ClientWithMiddleware, CasClientError> {
    let auth_middleware = auth_config.as_ref().map(AuthMiddleware::from).info_none("CAS auth disabled");
    let logging_middleware = Some(LoggingMiddleware);
    let session_middleware = (!session_id.is_empty()).then(|| SessionMiddleware(session_id.to_owned()));
    Ok(ClientBuilder::new(reqwest_client(user_agent)?)
        .maybe_with(auth_middleware)
        .maybe_with(logging_middleware)
        .maybe_with(session_middleware)
        .build())
}

/// Builds HTTP Client to talk to CAS.
pub fn build_http_client(session_id: &str, user_agent: &str) -> Result<ClientWithMiddleware, CasClientError> {
    build_auth_http_client(&None, session_id, user_agent)
}

/// Helper trait to allow the reqwest_middleware client to optionally add a middleware.
trait OptionalMiddleware {
    fn maybe_with<M: Middleware>(self, middleware: Option<M>) -> Self;
}

impl OptionalMiddleware for ClientBuilder {
    fn maybe_with<M: Middleware>(self, middleware: Option<M>) -> Self {
        match middleware {
            Some(m) => self.with(m),
            None => self,
        }
    }
}

#[derive(Copy, Clone)]
pub struct Api(pub &'static str);

/// Adds logging middleware that logs requests and responses.
pub struct LoggingMiddleware;

#[cfg_attr(not(target_family = "wasm"), async_trait::async_trait)]
#[cfg_attr(target_family = "wasm", async_trait::async_trait(?Send))]
impl Middleware for LoggingMiddleware {
    async fn handle(
        &self,
        req: Request,
        extensions: &mut http::Extensions,
        next: Next<'_>,
    ) -> reqwest_middleware::Result<Response> {
        let api = extensions.get::<Api>().map(|a| a.0);
        next.run(req, extensions)
            .instrument(info_span!("client::request", api))
            .await
            .inspect(|res| {
                let status_code = res.status().as_u16();
                let request_id = request_id_from_response(res);
                info!(request_id, status_code, "Received CAS response");
            })
            .inspect_err(|err| {
                warn!(?err, "Request error");
            })
    }
}

/// AuthMiddleware is a thread-safe middleware that adds a CAS auth token to outbound requests.
/// If the token it holds is expired, it will automatically be refreshed.
pub struct AuthMiddleware {
    token_provider: Arc<Mutex<TokenProvider>>,
}

impl AuthMiddleware {
    /// Fetches a token from our TokenProvider. This locks the TokenProvider as we might need
    /// to refresh the token if it has expired.
    ///
    /// In the common case, this lock is held only to read the underlying token stored
    /// in memory. However, in the event of an expired token (e.g. once every 15 min),
    /// we will need to hold the lock while making a call to refresh the token
    /// (e.g. to a remote service). During this time, no other CAS requests can proceed
    /// from this client until the token has been fetched. This is expected/ok since we
    /// don't have a valid token and thus any calls would fail.
    async fn get_token(&self) -> Result<String, anyhow::Error> {
        let mut provider = self.token_provider.lock().await;
        provider
            .get_valid_token()
            .await
            .map_err(|err| {
                warn!(?err, "Token refresh failed");
                anyhow!("couldn't get token: {err:?}")
            })
            .inspect(|_token| {
                info!("Token refresh successful for CAS authentication");
            })
    }
}

impl From<&AuthConfig> for AuthMiddleware {
    fn from(cfg: &AuthConfig) -> Self {
        Self {
            token_provider: Arc::new(Mutex::new(TokenProvider::new(cfg))),
        }
    }
}

#[cfg_attr(not(target_family = "wasm"), async_trait::async_trait)]
#[cfg_attr(target_family = "wasm", async_trait::async_trait(?Send))]
impl Middleware for AuthMiddleware {
    async fn handle(
        &self,
        mut req: Request,
        extensions: &mut http::Extensions,
        next: Next<'_>,
    ) -> reqwest_middleware::Result<Response> {
        let token = self.get_token().await.map_err(reqwest_middleware::Error::Middleware)?;

        let headers = req.headers_mut();
        headers.insert(AUTHORIZATION, HeaderValue::from_str(&format!("Bearer {token}")).unwrap());
        next.run(req, extensions).await
    }
}

pub struct SessionMiddleware(String);

// WASM compatibility; note the use of the pattern:
//
// #[cfg_attr(not(target_family = "wasm"), async_trait::async_trait)]
// #[cfg_attr(target_family = "wasm", async_trait::async_trait(?Send))]
// instead of just #[async_trait::async_trait]
// this makes the use of the async_trait wasm compatible to not enforce `Send` bounds when
// compiling for wasm, while those Send bounds are useful in non-wasm mode.

#[cfg_attr(not(target_family = "wasm"), async_trait::async_trait)]
#[cfg_attr(target_family = "wasm", async_trait::async_trait(?Send))]
impl Middleware for SessionMiddleware {
    async fn handle(
        &self,
        mut req: Request,
        extensions: &mut Extensions,
        next: Next<'_>,
    ) -> reqwest_middleware::Result<Response> {
        req.headers_mut()
            .insert(SESSION_ID_HEADER, HeaderValue::from_str(&self.0).unwrap());
        next.run(req, extensions).await
    }
}

/// Helper trait to log the different types of errors that come back from a request to CAS,
/// transforming the implementation into some new error type.
pub trait ResponseErrorLogger<T> {
    fn process_error(self, api: &str) -> T;
}

/// Add ResponseErrorLogger to Result<Response> for our requests.
/// This logs an error if one occurred before receiving a response or
/// if the status code indicates a failure.
/// As a result of these checks, the response is also transformed into a
/// cas_client::error::Result instead of the raw reqwest_middleware::Result.
impl ResponseErrorLogger<error::Result<Response>> for reqwest_middleware::Result<Response> {
    fn process_error(self, api: &str) -> error::Result<Response> {
        let res = self
            .map_err(CasClientError::from)
            .log_error(format!("error invoking {api} api"))?;
        let request_id = request_id_from_response(&res);
        let error_message = format!("{api} api failed: request id: {request_id}");
        let status = res.status();
        let res = res.error_for_status().map_err(CasClientError::from);
        match (api, status) {
            ("get_reconstruction", StatusCode::RANGE_NOT_SATISFIABLE) => res.debug_error(&error_message),
            // not all status codes mean fatal error
            _ => res.info_error(&error_message),
        }
    }
}

pub fn request_id_from_response(res: &Response) -> &str {
    res.headers()
        .get(REQUEST_ID_HEADER)
        .and_then(|h| h.to_str().ok())
        .unwrap_or_default()
}
