use std::sync::Arc;

use axum::extract::FromRequestParts;
use http::{HeaderMap, HeaderValue, StatusCode, request::Parts};
use jsonwebtoken::{DecodingKey, jwk::JwkSet};
use serde::Deserialize;
use services::{User, organization::OrganizationService};
use tokio::sync::RwLock;
use tracing::{info, warn};

use crate::AppState;

#[derive(Clone, Debug, Deserialize)]
struct Claims {
    id: String,
}

impl FromRequestParts<AppState> for User {
    type Rejection = (StatusCode, &'static str);

    async fn from_request_parts(
        parts: &mut Parts,
        state: &AppState,
    ) -> Result<Self, Self::Rejection> {
        let headers = HeaderMap::from_request_parts(parts, state)
            .await
            .map_err(|_| (StatusCode::BAD_REQUEST, "Failed to extract headers"))?;

        let user_id = if state.allow_test_user
            && let Some(user) = headers.get("x-user-id")
        {
            Some(user.to_str().unwrap_or_default().to_string())
        } else if let Some(api_key) = headers.get("x-api-key") {
            Some(from_api_key(api_key, state).await?)
        } else if let Some(bearer) = headers.get("Authorization") {
            Some(from_bearer(bearer, state).await?)
        } else {
            None
        };

        let Some(user_id) = user_id else {
            return Err((StatusCode::UNAUTHORIZED, "Authorization header missing"));
        };

        if let Some(org_id) = headers.get("x-org-id") {
            let Ok(org_id) = org_id.to_str().unwrap_or_default().to_string().parse() else {
                return Err((StatusCode::UNAUTHORIZED, "Authorization Invalid Org"));
            };

            let has_access =
                OrganizationService::has_access(&state.db_connection, &user_id, org_id)
                    .await
                    .map_err(|e| {
                        warn!("{:?}", e);
                        (
                            StatusCode::INTERNAL_SERVER_ERROR,
                            "Failed to check Org access",
                        )
                    })?;

            if !has_access {
                return Err((StatusCode::UNAUTHORIZED, "No Access"));
            }

            return Ok(User::new(user_id, Some(org_id.to_string())));
        }

        Ok(User::new(user_id, None))
    }
}

async fn from_bearer(
    bearer: &HeaderValue,
    state: &AppState,
) -> Result<String, (StatusCode, &'static str)> {
    let raw_jwt = bearer
        .to_str()
        .map_err(|_| (StatusCode::UNAUTHORIZED, "Invalid token"))?["Bearer ".len()..]
        .trim_start();

    let header = jsonwebtoken::decode_header(raw_jwt)
        .map_err(|_| (StatusCode::UNAUTHORIZED, "Invalid token"))?;

    let jwk = state
        .jwks_cache
        .get_key(&state.constants.auth_api_base, header.kid.as_deref())
        .await
        .map_err(|e| {
            warn!("Failed to get key: {:?}", e);
            (StatusCode::INTERNAL_SERVER_ERROR, "Failed to get key")
        })?
        .ok_or((StatusCode::UNAUTHORIZED, "Invalid token"))?;

    let key = DecodingKey::from_jwk(&jwk).map_err(|_| {
        (
            StatusCode::INTERNAL_SERVER_ERROR,
            "Failed to create decoding key",
        )
    })?;

    let mut validation = jsonwebtoken::Validation::new(jsonwebtoken::Algorithm::EdDSA);
    validation.set_audience(&[&state.constants.self_api_base]);

    let claims: Claims = jsonwebtoken::decode(raw_jwt, &key, &validation)
        .map_err(|_| (StatusCode::UNAUTHORIZED, "Invalid token"))?
        .claims;

    Ok(claims.id)
}

async fn from_api_key(
    header: &HeaderValue,
    state: &AppState,
) -> Result<String, (StatusCode, &'static str)> {
    let key = header
        .to_str()
        .map_err(|_| (StatusCode::UNAUTHORIZED, "Invalid token"))?;

    let response = reqwest::Client::new()
        .get(format!(
            "{}/api/auth/api-key/verify",
            state.constants.auth_api_base
        ))
        .header("x-api-key", key)
        .send()
        .await
        .map_err(|e| {
            warn!("{}", e);
            (
                StatusCode::INTERNAL_SERVER_ERROR,
                "failed to verify api key",
            )
        })?;

    if !response.status().is_success() {
        return Err((StatusCode::UNAUTHORIZED, "invalid api key"));
    }

    let text = response.text().await.map_err(|e| {
        warn!("{}", e);
        (StatusCode::UNAUTHORIZED, "invalid api key")
    })?;

    Ok(text)
}

#[derive(Clone)]
pub struct JwksCache {
    inner: Arc<RwLock<Option<JwkSet>>>,
}

impl Default for JwksCache {
    fn default() -> Self {
        Self {
            inner: Arc::new(RwLock::new(None)),
        }
    }
}

impl JwksCache {
    async fn get_key(
        &self,
        auth_api_base: &str,
        kid: Option<&str>,
    ) -> Result<Option<jsonwebtoken::jwk::Jwk>, reqwest::Error> {
        // Fast path: cached lookup
        let guard = self.inner.read().await;
        if let Some(jwks) = &*guard
            && let Some(jwk) = Self::find_jwk(jwks, kid)
        {
            return Ok(Some(jwk.clone()));
        }
        drop(guard);

        // Slow path: single refresh
        let mut guard = self.inner.write().await;

        // Double-check in case another task refreshed already
        if let Some(jwks) = &*guard
            && let Some(jwk) = Self::find_jwk(jwks, kid)
        {
            return Ok(Some(jwk.clone()));
        }

        info!("Refreshing JWKS from {}", auth_api_base);
        let fresh = reqwest::get(format!("{auth_api_base}/api/auth/jwks"))
            .await?
            .json::<JwkSet>()
            .await?;

        let found = Self::find_jwk(&fresh, kid);
        *guard = Some(fresh);
        drop(guard);

        Ok(found)
    }

    fn find_jwk(jwks: &JwkSet, kid: Option<&str>) -> Option<jsonwebtoken::jwk::Jwk> {
        jwks.keys
            .iter()
            .find(|j| j.common.key_id.as_deref() == kid)
            .cloned()
    }
}