use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine};

use axum::extract::Request;
use axum::http::{header::AUTHORIZATION, Uri};
use axum::middleware::Next;
use axum::response::Response;

/// JWT claim that carries the kobo-shelf auth token inside the device's
/// `KoboAccessToken`. The token is embedded in a claim (rather than sent raw)
/// because the device stores its access token as a JWT and treats it as one;
/// a bare opaque string is not a structurally valid access token.
const TOKEN_CLAIM: &str = "kobo_shelf_token";

/// Path segments directly under `/kobo/` that belong to the API surface a
/// device calls against its (token-free) `api_endpoint`. When a request targets
/// one, the auth token is taken from the `Authorization: Bearer` header and
/// spliced into the path so the `/kobo/{auth_token}/...` routes match.
///
/// Image and download URLs are the exception: they are self-authorizing links
/// the server hands the device, which it fetches without an auth header, so
/// they keep the token in the path. Their first segment is a token (UUID), not
/// one of these roots, so they pass through untouched.
const TOKEN_FREE_ROOTS: [&str; 2] = ["v1", "ping"];

/// A structurally valid but identity-less access-token JWT for the token-free
/// `device_auth` response, used when no kobo-shelf token can be recovered from the
/// request. Lets the firmware treat authentication as successful (and proceed
/// to re-initialize) without minting a real session; the device keeps
/// presenting its planted `KoboAccessToken` (which carries the real claim) as
/// `Authorization: Bearer` on subsequent requests.
pub fn session_placeholder_jwt() -> String {
    let header = URL_SAFE_NO_PAD.encode(r#"{"alg":"none","typ":"JWT"}"#);
    let payload = URL_SAFE_NO_PAD.encode(r#"{"sub":"kobo-shelf","exp":4102444800}"#);
    format!("{header}.{payload}.")
}

/// Encode a kobo-shelf auth token as the unsigned JWT the device stores as its
/// `KoboAccessToken` and presents back as `Authorization: Bearer`.
///
/// `alg: "none"` with an empty signature — never cryptographically verified;
/// the far-future `exp` keeps the device from treating it as expired. The
/// token travels in the [`TOKEN_CLAIM`] claim.
pub fn encode_token_jwt(token: &str) -> String {
    let header = URL_SAFE_NO_PAD.encode(r#"{"alg":"none","typ":"JWT"}"#);
    let payload = URL_SAFE_NO_PAD.encode(
        serde_json::json!({ "sub": "kobo-shelf", "exp": 4102444800u64, TOKEN_CLAIM: token })
            .to_string(),
    );
    format!("{header}.{payload}.")
}

/// axum middleware: rewrite token-free `/kobo/...` requests to the canonical
/// `/kobo/{auth_token}/...` form using the `Authorization: Bearer` token.
///
/// Must run **before** routing (apply it as an outer `tower::Layer`, not via
/// `Router::layer`, which runs after `matchit`).
pub async fn inject_bearer_token(mut req: Request, next: Next) -> Response {
    // Diagnostic: on token-free `/kobo` API paths, report whether the firmware
    // actually sent an `Authorization` header and whether we recovered a token
    // from it. This is the load-bearing unknown for header auth on-device.
    if token_free_rest(req.uri().path()).is_some() {
        let has_auth = req.headers().contains_key(AUTHORIZATION);
        tracing::debug!(
            path = %req.uri().path(),
            has_auth_header = has_auth,
            token_recovered = bearer_token(&req).is_some(),
            "kobo header-auth check",
        );
    }
    if let Some(uri) = rewritten_uri(&req) {
        *req.uri_mut() = uri;
    }
    next.run(req).await
}

/// axum middleware: emit a `tracing::error` for any 4xx/5xx response, with the
/// request method, URI, and status. Applied around the whole app so it covers
/// both Kobo and web routes.
pub async fn log_error_responses(req: Request, next: Next) -> Response {
    let method = req.method().clone();
    let uri = req.uri().clone();
    let response = next.run(req).await;
    let status = response.status();
    if status.is_client_error() || status.is_server_error() {
        tracing::error!(%method, %uri, status = status.as_u16(), "request failed");
    }
    response
}

/// If `path` is a token-free `/kobo` API path (one the device calls against its
/// `api_endpoint`), return the portion after `/kobo/`; otherwise `None`.
fn token_free_rest(path: &str) -> Option<&str> {
    let rest = match path.strip_prefix("/kobo/") {
        Some(rest) => rest,
        None if path == "/kobo" => "",
        None => return None,
    };
    let first = rest.split('/').next().unwrap_or("");
    (first.is_empty() || TOKEN_FREE_ROOTS.contains(&first)).then_some(rest)
}

/// Compute the token-injected URI, or `None` when the request should pass
/// through unchanged (not a token-free `/kobo` path, or no usable Bearer token).
fn rewritten_uri(req: &Request) -> Option<Uri> {
    let rest = token_free_rest(req.uri().path())?;

    let token = bearer_token(req)?;
    // Reject anything that could break out of the single path segment.
    if token.is_empty() || token.contains('/') || token.contains('?') || token.contains('#') {
        return None;
    }

    let new_path = format!("/kobo/{token}/{rest}");
    let new_pq = match req.uri().query() {
        Some(query) => format!("{new_path}?{query}"),
        None => new_path,
    };

    let mut parts = req.uri().clone().into_parts();
    parts.path_and_query = new_pq.parse().ok();
    Uri::from_parts(parts).ok()
}

/// Extract the kobo-shelf auth token from an `Authorization: Bearer <jwt>` header.
///
/// The value is the device's `KoboAccessToken` JWT; the token lives in its
/// [`TOKEN_CLAIM`] claim. A value that is not a JWT we minted (no claim) yields
/// `None` and the request passes through unauthenticated.
fn bearer_token(req: &Request) -> Option<String> {
    let value = req.headers().get(AUTHORIZATION)?.to_str().ok()?;
    let jwt = value
        .strip_prefix("Bearer ")
        .or_else(|| value.strip_prefix("bearer "))?
        .trim();
    token_from_jwt(jwt)
}

/// Pull the [`TOKEN_CLAIM`] out of an unsigned JWT's payload segment.
pub fn token_from_jwt(jwt: &str) -> Option<String> {
    let payload_b64 = jwt.split('.').nth(1)?;
    let payload = URL_SAFE_NO_PAD.decode(payload_b64).ok()?;
    let claims: serde_json::Value = serde_json::from_slice(&payload).ok()?;
    claims
        .get(TOKEN_CLAIM)?
        .as_str()
        .map(|token| token.to_string())
}