#![allow(non_snake_case)]
use std::collections::HashMap;
use std::sync::Arc;
use std::time::Duration;

use axum::debug_handler;
use axum::extract::State;
use axum::{
    extract::{Json, Path},
    http::StatusCode,
    response::IntoResponse,
};
use std::time::Instant;
use tower_cookies::Cookies;
use tower_sessions::Session;
use tracing::*;

/*
 * Webauthn RS auth handlers.
 * These files use webauthn to process the data received from each route, and are closely tied to axum
 */

// 1. Import the prelude - this contains everything needed for the server to function.
use webauthn_rs::prelude::*;

use crate::config::{COOKIE_NAME, SESSION_ACTIVE_TIME};
use crate::ServerState;

// 2. The first step a client (user) will carry out is requesting a credential to be
// registered. We need to provide a challenge for this. The work flow will be:
//
//          ┌───────────────┐     ┌───────────────┐      ┌───────────────┐
//          │ Authenticator │     │    Browser    │      │     Site      │
//          └───────────────┘     └───────────────┘      └───────────────┘
//                  │                     │                      │
//                  │                     │     1. Start Reg     │
//                  │                     │─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─▶│
//                  │                     │                      │
//                  │                     │     2. Challenge     │
//                  │                     │◀ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ┤
//                  │                     │                      │
//                  │  3. Select Token    │                      │
//             ─ ─ ─│◀ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─│                      │
//  4. Verify │     │                     │                      │
//                  │  4. Yield PubKey    │                      │
//            └ ─ ─▶│─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─▶                      │
//                  │                     │                      │
//                  │                     │  5. Send Reg Opts    │
//                  │                     │─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─▶│─ ─ ─
//                  │                     │                      │     │ 5. Verify
//                  │                     │                      │         PubKey
//                  │                     │                      │◀─ ─ ┘
//                  │                     │                      │─ ─ ─
//                  │                     │                      │     │ 6. Persist
//                  │                     │                      │       Credential
//                  │                     │                      │◀─ ─ ┘
//                  │                     │                      │
//                  │                     │                      │
//
// In this step, we are responding to the start reg(istration) request, and providing
// the challenge to the browser.

// TODO - Improve error handling and messages
fn auth_user(cookie_content:String,session_table:&mut HashMap<Uuid, Instant>) -> bool {
    let Ok(uuid) = Uuid::parse_str(&cookie_content) else {
        return false;
    };
    let Some(expire) = session_table.get(&uuid) else {
        return false;
    };
    if *expire <= Instant::now() {
        return false;
    }
    session_table.insert(uuid, Instant::now()+Duration::from_secs(*SESSION_ACTIVE_TIME));
    tracing::info!("valid cookie {}",uuid);
    return true;

}
#[debug_handler]
pub async fn start_register(
    State(state): State<Arc<ServerState>>,
    cookies: Cookies,
    session: Session,
    Path(username): Path<String>,
) -> Result<Json<CreationChallengeResponse>, String> {
    tracing::info!("Start register");
    // todo!("Auth User");
    let Some(cookie_content) = cookies.get(&COOKIE_NAME) else {
        return Err(WebauthnError::AuthenticationFailure.to_string());
    };
    let mut session_table = state.session.lock().await;
    if !auth_user(cookie_content.value().to_string(), &mut session_table){
        return Err(WebauthnError::AuthenticationFailure.to_string());
    }
    // Remove any previous registrations that may have occurred from the session.
    let _ = session
        .remove::<(String, Uuid, PasskeyRegistration)>("reg_state")
        .await
        .map_err(|_| WebauthnError::AuthenticationFailure.to_string())?;
    let pool = &state.db;
    let username = username.clone();
    let (user_id, exclude_credentials): (Uuid, Option<Vec<CredentialID>>) =
        match sqlx::query!("SELECT key FROM users WHERE name = $1;", username)
            .fetch_optional(pool)
            .await
            .map_err(|_| WebauthnError::AuthenticationFailure.to_string())?
        {
            Some(record) => {
                let uid = record.KEY.clone();
                let records = sqlx::query!(
                    "SELECT credential FROM CREDENTIALS WHERE user_id = $1;",
                    uid
                )
                .fetch_all(pool)
                .await
                .map_err(|_| WebauthnError::AuthenticationFailure.to_string())?;

                (
                    Uuid::parse_str(&record.KEY)
                        .map_err(|_| WebauthnError::AuthenticationFailure.to_string())?,
                    Some(
                        records
                            .iter()
                            .map(|record| serde_json::from_str::<Passkey>(&record.CREDENTIAL))
                            .collect::<Result<Vec<Passkey>, _>>()
                            .map_err(|_| WebauthnError::CredentialPersistenceError.to_string())?
                            .iter()
                            .map(|passkey| passkey.cred_id().clone())
                            .collect(),
                    ),
                )
            }
            None => (Uuid::new_v4(), None),
        };

    let res = match state.webauthn.start_passkey_registration(
        user_id,
        &username,
        &username,
        exclude_credentials,
    ) {
        Ok((ccr, reg_state)) => {
            // Note that due to the session store in use being a server side memory store, this is
            // safe to store the reg_state into the session since it is not client controlled and
            // not open to replay attacks. If this was a cookie store, this would be UNSAFE.
            session
                .insert("reg_state", (username, user_id, reg_state))
                .await
                .map_err(|_| WebauthnError::ChallengePersistenceError.to_string())?;
            Json(ccr)
        }
        Err(e) => {
            debug!("challenge_register -> {:?}", e);
            return Err(WebauthnError::ChallengePersistenceError.to_string());
        }
    };
    Ok(res)
}

// 3. The browser has completed it's steps and the user has created a public key
// on their device. Now we have the registration options sent to us, and we need
// to verify these and persist them.

pub async fn finish_register(
    State(state): State<Arc<ServerState>>,
    session: Session,
    Json(reg): Json<RegisterPublicKeyCredential>,
) -> Result<impl IntoResponse, String> {
    info!("Confirming registration....");
    let pool = &state.db;
    let Ok(Some((user_name, user_id, reg_state))) = session
        .get::<(String, Uuid, PasskeyRegistration)>("reg_state")
        .await
    else {
        return Err(WebauthnError::AuthenticationFailure.to_string()); //Corrupt Session
    };

    let _ = session.remove::<(Uuid, PasskeyAuthentication)>("reg_state").await;

    let res = match state.webauthn.finish_passkey_registration(&reg, &reg_state) {
        Ok(key) => {
            info!("Passkey is okay");
            let uid = &user_id.to_string();
            let username = &user_name.to_string();
            // Check if the user_id already exists
            let record = sqlx::query!(
                "SELECT COUNT(KEY) AS count FROM users WHERE KEY = $1;",
                uid
            )
            .fetch_one(pool)
            .await
            .map_err(|_| WebauthnError::AuthenticationFailure.to_string())?;

            // If the user doesn't exist, insert them into the users table
            if record.count == 0
                && sqlx::query!(
                    "INSERT INTO users(KEY, NAME) VALUES($1, $2);",
                    uid,
                    username
                )
                .execute(pool)
                .await
                .map_err(|_| WebauthnError::UserNotPresent.to_string())?
                .rows_affected()
                    != 1
            {
                return Err(WebauthnError::AuthenticationFailure.to_string());
            }

            // Serialise the key
            let serialised_key = serde_json::ser::to_string(&key)
                .map_err(|_| WebauthnError::CredentialPersistenceError.to_string())?;

            // Insert the key into the auth table
            if sqlx::query!(
                "INSERT INTO CREDENTIALS(user_id, credential) VALUES($1, $2);",
                uid,
                serialised_key
            )
            .execute(pool)
            .await
            .map_err(|_| WebauthnError::UserNotPresent.to_string())?
            .rows_affected()
                != 1
            {
                return Err(WebauthnError::AuthenticationFailure.to_string());
            }

            StatusCode::OK
        }
        Err(e) => {
            debug!("challenge_register -> {:?}", e);
            StatusCode::BAD_REQUEST
        }
    };

    Ok(res)
}

// 4. Now that our public key has been registered, we can authenticate a user and verify
// that they are the holder of that security token. The work flow is similar to registration.
//
//          ┌───────────────┐     ┌───────────────┐      ┌───────────────┐
//          │ Authenticator │     │    Browser    │      │     Site      │
//          └───────────────┘     └───────────────┘      └───────────────┘
//                  │                     │                      │
//                  │                     │     1. Start Auth    │
//                  │                     │─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─▶│
//                  │                     │                      │
//                  │                     │     2. Challenge     │
//                  │                     │◀ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ┤
//                  │                     │                      │
//                  │  3. Select Token    │                      │
//             ─ ─ ─│◀ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─│                      │
//  4. Verify │     │                     │                      │
//                  │    4. Yield Sig     │                      │
//            └ ─ ─▶│─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─▶                      │
//                  │                     │    5. Send Auth      │
//                  │                     │        Opts          │
//                  │                     │─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─▶│─ ─ ─
//                  │                     │                      │     │ 5. Verify
//                  │                     │                      │          Sig
//                  │                     │                      │◀─ ─ ┘
//                  │                     │                      │
//                  │                     │                      │
//
// The user indicates the wish to start authentication and we need to provide a challenge.

pub async fn start_authentication(
    State(state): State<Arc<ServerState>>,
    session: Session,
    Path(user_name): Path<String>,
) -> Result<impl IntoResponse, String> {
    info!("Start Authentication");
    let pool = &state.db;
    // Remove any previous authentication that may have occurred from the session.
    let _ = session.remove::<(Uuid, PasskeyAuthentication)>("auth_state");
    let user_id = sqlx::query!("SELECT KEY FROM users WHERE NAME = $1;", user_name)
        .fetch_one(pool)
        .await
        .map_err(|_| WebauthnError::UserNotPresent.to_string())?
        .KEY;

    let records = sqlx::query!(
        "SELECT credential FROM CREDENTIALS WHERE user_id = $1;",
        user_id
    )
    .fetch_all(pool)
    .await
    .map_err(|_| WebauthnError::AuthenticationFailure.to_string())?;

    if records.is_empty() {
        return Err(WebauthnError::CredentialNotFound.to_string());
    }

    let allow_credentials: Vec<Passkey> = records
        .iter()
        .map(|record| serde_json::de::from_str::<Passkey>(&record.CREDENTIAL))
        .collect::<Result<Vec<Passkey>, _>>()
        .map_err(|_| WebauthnError::CredentialPersistenceError.to_string())?;

    let res = match state
        .webauthn
        .start_passkey_authentication(&allow_credentials)
    {
        Ok((rcr, auth_state)) => {
            // Note that due to the session store in use being a server side memory store, this is
            // safe to store the auth_state into the session since it is not client controlled and
            // not open to replay attacks. If this was a cookie store, this would be UNSAFE.
            session
                .insert("auth_state", (&user_id, auth_state))
                .await
                .map_err(|_| WebauthnError::AuthenticationFailure.to_string())?;
            Json(rcr)
        }
        Err(e) => {
            debug!("challenge_authenticate -> {:?}", e);
            return Err(WebauthnError::MismatchedChallenge.to_string());
        }
    };
    Ok(res)
}

// 5. The browser and user have completed their part of the processing. Only in the
// case that the webauthn authenticate call returns Ok, is authentication considered
// a success. If the browser does not complete this call, or *any* error occurs,
// this is an authentication failure.

pub async fn finish_authentication(
    State(state): State<Arc<ServerState>>,
    session: Session,
    Json(auth): Json<PublicKeyCredential>,
) -> Result<impl IntoResponse, String> {
    let pool = &state.db;
    let (user_id, auth_state): (Uuid, PasskeyAuthentication) = session
        .get("auth_state")
        .await
        .unwrap()
        .ok_or(WebauthnError::AuthenticationFailure.to_string())?;

    let _ = session.remove::<(Uuid, PasskeyAuthentication)>("auth_state");

    let res = match state
        .webauthn
        .finish_passkey_authentication(&auth, &auth_state)
    {
        Ok(auth_result) => {
            let uid = user_id.clone().to_string();
            let records = sqlx::query!(
                "SELECT CREDENTIAL FROM CREDENTIALS WHERE USER_ID = $1;",
                uid
            )
            .fetch_all(pool)
            .await
            .map_err(|_| WebauthnError::AuthenticatorDataMissingExtension.to_string())?;
            if records.is_empty() {
                return Err(WebauthnError::UserNotPresent.to_string());
            }

            for record in records {
                let mut credential = serde_json::from_str::<Passkey>(&record.CREDENTIAL)
                    .map_err(|_| WebauthnError::CredentialExistCheckError.to_string())?;

                if credential.cred_id() == auth_result.cred_id() {
                    credential.update_credential(&auth_result);

                    let credential = serde_json::to_string(&credential)
                        .map_err(|_| WebauthnError::CredentialPersistenceError.to_string())?;

                    if sqlx::query!(
                        "UPDATE CREDENTIALS SET credential = $1 WHERE user_id = $2 AND credential = $3;",
                        credential,
                        uid,
                        record.CREDENTIAL
                    )
                    .execute(pool)
                    .await.is_ok_and(|e|e.rows_affected() !=1 )
                    {
                        return Err(WebauthnError::AuthenticationFailure.to_string());
                    }

                    break;
                }
            }

            let user_name = sqlx::query!(
                "SELECT NAME FROM users WHERE KEY = $1;",
                uid
            )
            .fetch_one(pool)
            .await
            .map_err(|_| WebauthnError::AuthenticationFailure.to_string())?;
            // Add our own values to the session
            session
                .insert("user_id", user_id)
                .await
                .map_err(|_| WebauthnError::InvalidUserUniqueId.to_string())?;
            session
                .insert("user_name", user_name.NAME)
                .await
                .map_err(|_| WebauthnError::InvalidUsername.to_string())?;

            StatusCode::OK
        }
        Err(e) => {
            debug!("challenge_register -> {:?}", e);
            StatusCode::BAD_REQUEST
        }
    };

    info!("Authentication Successful!");
    Ok(res)
}