use crate::helpers::json_error_response;
use axum::extract::Request;
use axum::http::{HeaderMap, StatusCode};
use axum::middleware::Next;
use axum::response::IntoResponse;
use jwt_compact::alg::{Hs256, Hs256Key};
use jwt_compact::{AlgorithmExt, Claims, Token, UntrustedToken, ValidationError};
use serde::{Deserialize, Serialize};
use std::env;
use tracing::log;

#[derive(Clone, Debug)]
pub struct Did(pub Option<String>);

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum AuthScheme {
    Bearer,
    DPoP,
}

#[derive(Serialize, Deserialize)]
pub struct TokenClaims {
    pub sub: String,
}

pub async fn extract_did(mut req: Request, next: Next) -> impl IntoResponse {
    let auth = extract_auth(req.headers());

    match auth {
        Ok(auth_opt) => {
            match auth_opt {
                None => json_error_response(StatusCode::BAD_REQUEST, "TokenRequired", "")
                    .expect("Error creating an error response"),
                Some((scheme, token_str)) => {
                    // For Bearer, validate JWT and extract DID from `sub`.
                    // For DPoP, we currently only pass through and do not validate here; insert None DID.
                    match scheme {
                        AuthScheme::Bearer => {
                            let token = UntrustedToken::new(&token_str);
                            if token.is_err() {
                                return json_error_response(
                                    StatusCode::BAD_REQUEST,
                                    "TokenRequired",
                                    "",
                                )
                                .expect("Error creating an error response");
                            }
                            let parsed_token = token.expect("Already checked for error");
                            let claims: Result<Claims<TokenClaims>, ValidationError> =
                                parsed_token.deserialize_claims_unchecked();
                            if claims.is_err() {
                                return json_error_response(
                                    StatusCode::BAD_REQUEST,
                                    "TokenRequired",
                                    "",
                                )
                                .expect("Error creating an error response");
                            }

                            let key = Hs256Key::new(
                                env::var("PDS_JWT_SECRET")
                                    .expect("PDS_JWT_SECRET not set in the pds.env"),
                            );
                            let token: Result<Token<TokenClaims>, ValidationError> =
                                Hs256.validator(&key).validate(&parsed_token);
                            if token.is_err() {
                                return json_error_response(
                                    StatusCode::BAD_REQUEST,
                                    "InvalidToken",
                                    "",
                                )
                                .expect("Error creating an error response");
                            }
                            let token = token.expect("Already checked for error,");
                            req.extensions_mut()
                                .insert(Did(Some(token.claims().custom.sub.clone())));
                        }
                        AuthScheme::DPoP => {
                            //Not going to worry about oauth email update for now, just always forward to the PDS
                            req.extensions_mut().insert(Did(None));
                        }
                    }

                    next.run(req).await
                }
            }
        }
        Err(err) => {
            log::error!("Error extracting token: {err}");
            json_error_response(StatusCode::BAD_REQUEST, "InvalidToken", "")
                .expect("Error creating an error response")
        }
    }
}

fn extract_auth(headers: &HeaderMap) -> Result<Option<(AuthScheme, String)>, String> {
    match headers.get(axum::http::header::AUTHORIZATION) {
        None => Ok(None),
        Some(hv) => {
            match hv.to_str() {
                Err(_) => Err("Authorization header is not valid".into()),
                Ok(s) => {
                    // Accept forms like: "Bearer <token>" or "DPoP <token>" (case-sensitive for the scheme here)
                    let mut parts = s.splitn(2, ' ');
                    match (parts.next(), parts.next()) {
                    (Some("Bearer"), Some(tok)) if !tok.is_empty() =>
                        Ok(Some((AuthScheme::Bearer, tok.to_string()))),
                    (Some("DPoP"), Some(tok)) if !tok.is_empty() =>
                        Ok(Some((AuthScheme::DPoP, tok.to_string()))),
                    _ => Err("Authorization header must be in format 'Bearer <token>' or 'DPoP <token>'".into()),
                }
                }
            }
        }
    }
}