use std::{collections::HashMap, ops::Deref};

use anyhow::Error;
use futures_util::Future;
use giterated_models::model::{
    authenticated::{Authenticated, AuthenticationSource, UserTokenMetadata},
    instance::Instance,
    user::User,
};
use jsonwebtoken::{decode, Algorithm, DecodingKey, TokenData, Validation};
use rsa::{
    pkcs1::{DecodeRsaPrivateKey, DecodeRsaPublicKey},
    pss::{Signature, VerifyingKey},
    sha2::Sha256,
    signature::Verifier,
    RsaPublicKey,
};
use serde::{de::DeserializeOwned, Serialize};
use serde_json::Value;

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: Authenticated<HashMap<String, Value>> =
            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 = 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: Authenticated<Value> =
            serde_json::from_slice(&network_message).map_err(|e| Error::from(e))?;

        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)?;

        let public_key = {
            let cached_keys = state.cached_keys.read().await;

            if let Some(key) = cached_keys.get(&instance) {
                key.clone()
            } else {
                drop(cached_keys);
                let mut cached_keys = state.cached_keys.write().await;
                let key = public_key(instance).await?;
                let public_key = RsaPublicKey::from_pkcs1_pem(&key).unwrap();
                cached_keys.insert(instance.clone(), public_key.clone());
                public_key
            }
        };

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

        let message_json = serde_json::to_vec(&message.message).unwrap();

        verifying_key.verify(
            &message_json,
            &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,
    S: Clone + Send + Sync,
{
    async fn from_message(message: &NetworkMessage, _: &S) -> Result<Self, Error> {
        Ok(Message(serde_json::from_slice(&message)?))
    }
}

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

async fn public_key(instance: &Instance) -> Result<String, Error> {
    let key = reqwest::get(format!("https://{}/.giterated/pubkey.pem", instance.url))
        .await?
        .text()
        .await?;

    Ok(key)
}