use anyhow::Error;
use futures_util::Future;

use giterated_models::instance::Instance;

use giterated_models::user::User;

use giterated_models::authenticated::{
    AuthenticatedPayload, AuthenticationSource, UserTokenMetadata,
};
use jsonwebtoken::{decode, Algorithm, DecodingKey, TokenData, Validation};
use rsa::{
    pkcs1::DecodeRsaPublicKey,
    pss::{Signature, VerifyingKey},
    sha2::Sha256,
    signature::Verifier,
    RsaPublicKey,
};
use serde::{de::DeserializeOwned, Serialize};
use std::{fmt::Debug, ops::Deref};

use crate::connection::wrapper::ConnectionState;

pub struct NetworkMessage(pub Vec<u8>);

impl Deref for NetworkMessage {
    type Target = [u8];

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

pub struct AuthenticatedUser(pub User);

#[derive(Debug, thiserror::Error)]
pub enum UserAuthenticationError {
    #[error("user authentication missing")]
    Missing,
    // #[error("{0}")]
    // InstanceAuthentication(#[from] Error),
    #[error("user token was invalid")]
    InvalidToken,
    #[error("an error has occured")]
    Other(#[from] Error),
}

pub struct AuthenticatedInstance(Instance);

impl AuthenticatedInstance {
    pub fn inner(&self) -> &Instance {
        &self.0
    }
}

#[async_trait::async_trait]
pub trait FromMessage<S: Send + Sync>: Sized + Send + Sync {
    async fn from_message(message: &NetworkMessage, state: &S) -> Result<Self, Error>;
}

#[async_trait::async_trait]
impl FromMessage<ConnectionState> for AuthenticatedUser {
    async fn from_message(
        network_message: &NetworkMessage,
        state: &ConnectionState,
    ) -> Result<Self, Error> {
        let message: AuthenticatedPayload =
            serde_json::from_slice(&network_message).map_err(|e| Error::from(e))?;

        let (auth_user, auth_token) = message
            .source
            .iter()
            .filter_map(|auth| {
                if let AuthenticationSource::User { user, token } = auth {
                    Some((user, token))
                } else {
                    None
                }
            })
            .next()
            .ok_or_else(|| UserAuthenticationError::Missing)?;

        let authenticated_instance =
            AuthenticatedInstance::from_message(network_message, state).await?;

        let public_key_raw = state.public_key(&auth_user.instance).await?;
        let verification_key = DecodingKey::from_rsa_pem(public_key_raw.as_bytes()).unwrap();

        let data: TokenData<UserTokenMetadata> = decode(
            auth_token.as_ref(),
            &verification_key,
            &Validation::new(Algorithm::RS256),
        )
        .unwrap();

        if data.claims.user != *auth_user
            || data.claims.generated_for != *authenticated_instance.inner()
        {
            Err(Error::from(UserAuthenticationError::InvalidToken))
        } else {
            Ok(AuthenticatedUser(data.claims.user))
        }
    }
}

#[async_trait::async_trait]
impl FromMessage<ConnectionState> for AuthenticatedInstance {
    async fn from_message(
        network_message: &NetworkMessage,
        state: &ConnectionState,
    ) -> Result<Self, Error> {
        let message: AuthenticatedPayload =
            serde_json::from_slice(&network_message).map_err(|e| Error::from(e))?;

        info!("Authenticated payload: {:?}", message);

        let (instance, signature) = message
            .source
            .iter()
            .filter_map(|auth: &AuthenticationSource| {
                if let AuthenticationSource::Instance {
                    instance,
                    signature,
                } = auth
                {
                    Some((instance, signature))
                } else {
                    None
                }
            })
            .next()
            // TODO: Instance authentication error
            .ok_or_else(|| UserAuthenticationError::Missing)?;

        info!("Instance: {}", instance.clone().to_string());

        info!("Instance public key: {}", state.public_key(instance).await?);

        let public_key = RsaPublicKey::from_pkcs1_pem(&state.public_key(instance).await?).unwrap();

        let verifying_key: VerifyingKey<Sha256> = VerifyingKey::new(public_key);

        verifying_key.verify(
            &message.payload,
            &Signature::try_from(signature.as_ref()).unwrap(),
        )?;

        Ok(AuthenticatedInstance(instance.clone()))
    }
}

#[async_trait::async_trait]
impl<S, T> FromMessage<S> for Option<T>
where
    T: FromMessage<S>,
    S: Send + Sync + 'static,
{
    async fn from_message(message: &NetworkMessage, state: &S) -> Result<Self, Error> {
        Ok(T::from_message(message, state).await.ok())
    }
}

#[async_trait::async_trait]
pub trait MessageHandler<T, S, R> {
    async fn handle_message(self, message: &NetworkMessage, state: &S) -> Result<R, Error>;
}
#[async_trait::async_trait]
impl<F, R, S, T1, T, E> MessageHandler<(T1,), S, R> for T
where
    T: FnOnce(T1) -> F + Clone + Send + 'static,
    F: Future<Output = Result<R, E>> + Send,
    T1: FromMessage<S> + Send,
    S: Send + Sync,
    E: std::error::Error + Send + Sync + 'static,
{
    async fn handle_message(self, message: &NetworkMessage, state: &S) -> Result<R, Error> {
        let value = T1::from_message(message, state).await?;
        self(value).await.map_err(|e| Error::from(e))
    }
}

#[async_trait::async_trait]
impl<F, R, S, T1, T2, T, E> MessageHandler<(T1, T2), S, R> for T
where
    T: FnOnce(T1, T2) -> F + Clone + Send + 'static,
    F: Future<Output = Result<R, E>> + Send,
    T1: FromMessage<S> + Send,
    T2: FromMessage<S> + Send,
    S: Send + Sync,
    E: std::error::Error + Send + Sync + 'static,
{
    async fn handle_message(self, message: &NetworkMessage, state: &S) -> Result<R, Error> {
        let value = T1::from_message(message, state).await?;
        let value_2 = T2::from_message(message, state).await?;
        self(value, value_2).await.map_err(|e| Error::from(e))
    }
}

#[async_trait::async_trait]
impl<F, R, S, T1, T2, T3, T, E> MessageHandler<(T1, T2, T3), S, R> for T
where
    T: FnOnce(T1, T2, T3) -> F + Clone + Send + 'static,
    F: Future<Output = Result<R, E>> + Send,
    T1: FromMessage<S> + Send,
    T2: FromMessage<S> + Send,
    T3: FromMessage<S> + Send,
    S: Send + Sync,
    E: std::error::Error + Send + Sync + 'static,
{
    async fn handle_message(self, message: &NetworkMessage, state: &S) -> Result<R, Error> {
        let value = T1::from_message(message, state).await?;
        let value_2 = T2::from_message(message, state).await?;
        let value_3 = T3::from_message(message, state).await?;

        self(value, value_2, value_3)
            .await
            .map_err(|e| Error::from(e))
    }
}

pub struct State<T>(pub T);

#[async_trait::async_trait]
impl<T> FromMessage<T> for State<T>
where
    T: Clone + Send + Sync,
{
    async fn from_message(_: &NetworkMessage, state: &T) -> Result<Self, Error> {
        Ok(Self(state.clone()))
    }
}

// Temp
#[async_trait::async_trait]
impl<T, S> FromMessage<S> for Message<T>
where
    T: DeserializeOwned + Send + Sync + Serialize + Debug,
    S: Clone + Send + Sync,
{
    async fn from_message(message: &NetworkMessage, _: &S) -> Result<Self, Error> {
        let payload: AuthenticatedPayload = serde_json::from_slice(&message)?;
        let payload = bincode::deserialize(&payload.payload)?;

        info!("Deserialized payload: {:#?}", payload);

        Ok(Message(payload))
    }
}

pub struct Message<T: Serialize + DeserializeOwned>(pub T);

/// Handshake-specific message type.
///
/// Uses basic serde_json-based deserialization to maintain the highest
/// level of compatibility across versions.
pub struct HandshakeMessage<T: Serialize + DeserializeOwned>(pub T);

#[async_trait::async_trait]
impl<T, S> FromMessage<S> for HandshakeMessage<T>
where
    T: DeserializeOwned + Send + Sync + Serialize,
    S: Clone + Send + Sync,
{
    async fn from_message(message: &NetworkMessage, _: &S) -> Result<Self, Error> {
        Ok(HandshakeMessage(serde_json::from_slice(&message.0)?))
    }
}