use std::sync::LazyLock;

use axum::{extract::Request, http::StatusCode, middleware::Next, response::Response};
use chrono::Utc;
use jsonwebtoken::{Algorithm, DecodingKey, EncodingKey, Validation, decode};
use serde::{Deserialize, Serialize};
use sqlx::FromRow;

#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct Claims {
    pub user_id: i64,
    pub iat: usize,
    pub exp: usize,
}

pub static AUTH_SECRET_KEY: LazyLock<EncodingKey> = LazyLock::new(|| {
    let bytes = std::fs::read("./private_key.pem").unwrap();
    let encoding_key = EncodingKey::from_rsa_pem(&bytes).unwrap();

    encoding_key
});

pub static AUTH_PUBLIC_KEY: LazyLock<DecodingKey> = LazyLock::new(|| {
    let bytes = std::fs::read("./public_key.pem").unwrap();
    let decoding_key = DecodingKey::from_rsa_pem(&bytes).unwrap();

    decoding_key
});

impl Claims {
    pub fn new(user_id: i64) -> Self {
        let now = (Utc::now().timestamp_millis() / 1000) as usize;
        Self {
            user_id,
            iat: now,
            // Should be about 1 year
            exp: now + 31540000,
        }
    }
}

#[derive(Debug, Clone, FromRow)]
pub struct User {
    pub id: i64,
    pub name: String,
    pub hash: String,
}

pub fn verify_token(token: &str) -> Option<Claims> {
    let key = &*AUTH_PUBLIC_KEY;

    decode::<Claims>(token, key, &Validation::new(Algorithm::RS512))
        .map(|token_data| token_data.claims)
        .ok()
}

const AUTH_HEADER_PREFIX: &'static str = "Bearer ";

pub async fn authentication_middleware(
    mut request: Request,
    next: Next,
) -> Result<Response, StatusCode> {
    let Some(Ok(auth_header)) = request
        .headers()
        .get(axum::http::header::AUTHORIZATION)
        .map(|header| header.to_str())
    else {
        return Err(StatusCode::UNAUTHORIZED);
    };

    if auth_header.len() <= AUTH_HEADER_PREFIX.len() {
        return Err(StatusCode::UNAUTHORIZED);
    }

    let token = auth_header.split_at(AUTH_HEADER_PREFIX.len()).1;

    let Some(claims) = verify_token(token) else {
        return Err(StatusCode::UNAUTHORIZED);
    };

    request.extensions_mut().insert(claims);

    return Ok(next.run(request).await);
}