use opaque_ke::{
    rand::rngs::OsRng, CredentialFinalization, CredentialRequest,
    RegistrationRequest, ServerLogin, ServerLoginParameters,
    ServerRegistration,
};
use uuid::Uuid;
use crate::models::opaque::{LoginSession, NKodeCipherSuite, NKodeServerSetup, PasswordFile, RegisterSession};
use crate::repository::opaque::repos::{OpaqueDatabaseRepo, AuthRepoError, OpaqueSessionRepo};

pub struct RegCache {
    pub session_id: Uuid,
    pub identifier: Vec<u8>,
}

pub struct LoginCache {
    pub session_id: Uuid,
    pub identifiers: Vec<u8>,
    pub server_login: ServerLogin<NKodeCipherSuite>,
}

pub trait CredKind {
    fn has<R: OpaqueDatabaseRepo>(repo: &R, id: &[u8]) -> bool;
    fn get_pf<R: OpaqueDatabaseRepo>(repo: &R, id: &[u8]) -> Result<PasswordFile, AuthRepoError>;
    fn put_pf<R: OpaqueDatabaseRepo>(repo: &mut R, id: &[u8], pf: PasswordFile) -> Result<(), AuthRepoError>;
    fn prereq_for_register<R: OpaqueDatabaseRepo>(_repo: &R, _id: &[u8]) -> Result<(), AuthRepoError> {
        Ok(())
    }
}

pub struct Key;
pub struct Code;

impl CredKind for Key {
    fn has<R: OpaqueDatabaseRepo>(repo: &R, id: &[u8]) -> bool {
        repo.has_key(id)
    }
    fn get_pf<R: OpaqueDatabaseRepo>(repo: &R, id: &[u8]) -> Result<PasswordFile, AuthRepoError> {
        repo.get_key_passcode_file(id)
    }
    fn put_pf<R: OpaqueDatabaseRepo>(repo: &mut R, id: &[u8], pf: PasswordFile) -> Result<(), AuthRepoError> {
        repo.new_key(id, pf)
    }
}

impl CredKind for Code {
    fn has<R: OpaqueDatabaseRepo>(repo: &R, id: &[u8]) -> bool {
        repo.has_code(id)
    }
    fn get_pf<R: OpaqueDatabaseRepo>(repo: &R, id: &[u8]) -> Result<PasswordFile, AuthRepoError> {
        repo.get_code_passcode_file(id)
    }
    fn put_pf<R: OpaqueDatabaseRepo>(repo: &mut R, id: &[u8], pf: PasswordFile) -> Result<(), AuthRepoError> {
        repo.new_code(id, pf)
    }
    fn prereq_for_register<R: OpaqueDatabaseRepo>(repo: &R, id: &[u8]) -> Result<(), AuthRepoError> {
        if repo.has_key(id) {
            Ok(())
        } else {
            Err(AuthRepoError::KeyNotRegistered)
        }
    }
}

pub struct OpaqueAuth<R: OpaqueDatabaseRepo, S: OpaqueSessionRepo> {
    server_setup: NKodeServerSetup,
    user_repo: R,
    session: S,
}

impl<R: OpaqueDatabaseRepo, S: OpaqueSessionRepo> OpaqueAuth<R, S> {
    pub fn new(server_setup: NKodeServerSetup, user_repo: R, session: S) -> Self {
        Self { server_setup, user_repo, session }
    }

    pub async fn reg_start<K: CredKind>(
        &mut self,
        identifier: &[u8],
        request: RegistrationRequest<NKodeCipherSuite>,
    ) -> Result<RegisterSession, String> {
        K::prereq_for_register(&self.user_repo, identifier)
            .map_err(|e| format!("registration prereq failed: {e:?}"))?;
        let start = ServerRegistration::<NKodeCipherSuite>::start(
            &self.server_setup,
            request,
            identifier,
        ).map_err(|e| format!("opaque reg start: {e:?}"))?;
        let cache = self.session
            .new_reg_session(identifier)
            .map_err(|e| format!("reg cache: {e}"))?;

        Ok(RegisterSession { session_id: cache.session_id, response: start.message })
    }

    pub async fn reg_finish<K: CredKind>(
        &mut self,
        session_id: &Uuid,
        password_file: PasswordFile,
    ) -> Result<(), String> {
        let sess = self.session
            .get_reg_session(session_id)
            .map_err(|e| format!("get reg session: {e}"))?;
        K::prereq_for_register(&self.user_repo, sess.identifier.as_slice())
            .map_err(|e| format!("registration prereq failed: {e:?}"))?;
        K::put_pf(&mut self.user_repo, sess.identifier.as_slice(), password_file)
            .map_err(|e| format!("repo write: {e:?}"))?;
        self.session
            .clear_reg_session(session_id)
            .map_err(|e| format!("clear reg session: {e}"))
    }

    pub async fn login_start<K: CredKind>(
        &mut self,
        identifier: &[u8],
        request: CredentialRequest<NKodeCipherSuite>,
    ) -> Result<LoginSession, String> {
        let password_file = K::get_pf(&self.user_repo, identifier)
            .map_err(|e| format!("repo read: {e:?}"))?;

        let password_file =
            ServerRegistration::<NKodeCipherSuite>::deserialize(password_file.as_slice())
                .map_err(|e| format!("pf deserialize: {e:?}"))?;

        let mut server_rng = OsRng;
        let start = ServerLogin::start(
            &mut server_rng,
            &self.server_setup,
            Some(password_file),
            request,
            identifier,
            ServerLoginParameters::default(),
        ).map_err(|e| format!("opaque login start: {e:?}"))?;

        let cache = self.session
            .new_login_session(identifier, start.state)
            .map_err(|e| format!("login cache: {e}"))?;

        Ok(LoginSession { session_id: cache.session_id, response: start.message })
    }

    pub async fn login_finish<K: CredKind>(
        &mut self,
        session_id: &Uuid,
        finalize: CredentialFinalization<NKodeCipherSuite>,
    ) -> Result<Vec<u8>, String> {
        let cache = self.session
            .get_login_session(session_id)
            .map_err(|e| format!("get login session: {e}"))?;

        let finish = cache.server_login
            .finish(finalize, ServerLoginParameters::default())
            .map_err(|e| format!("opaque login finish: {e:?}"))?;

        self.session
            .clear_login_session(session_id)
            .map_err(|e| format!("clear login session: {e}"))?;

        Ok(finish.session_key.to_vec())
    }
}