use crate::{
    collection::Collection, decoder::RecordSchema, indexer::mapping::Mapping, record::Record,
};

#[cfg(feature = "native")]
use crate::config::{self, Config};

use cipherstash_grpc::{
    self,
    api::{
        self,
        api_client::ApiClient,
        collections,
        documents::{
            self, put_reply::PutError, streaming_put_reply::StreamingPutError,
            streaming_put_request::Payload as StreamingPutPayload, GetReply, PutReply,
            StreamingPutBegin, StreamingPutDocument, StreamingPutReply, StreamingPutRequest,
        },
        queries::{
            constraint, query_reply::QueryError, Constraint, Exact, Filter, Ordering, Query,
            QueryReply, QueryRequest, Range,
        },
    },
    tonic::{
        self,
        codegen::{Body, StdError},
        metadata::AsciiMetadataKey,
        IntoRequest, Request,
    },
};
use envelopers::EnvelopeCipher;
use futures::{
    channel::mpsc::unbounded,
    stream::{self, Stream, StreamExt},
    try_join, TryFutureExt,
};
use hmac::{Hmac, Mac};
use rand_chacha::{rand_core::SeedableRng, ChaChaRng};
use sha2::Sha256;
#[cfg(feature = "native")]
use std::path::PathBuf;
use uuid::Uuid;

pub use self::rpc::{
    AggregateType, CollectionDecryptResult, DecryptedCollection, Direction, QueryInput,
};
use self::rpc::{DecryptedCollectionMetadata, DecryptedIndex};

pub mod client_data;
mod rpc;

use client_data::{GetAccessToken, GetCipherEngine, GetNamingKey};

fn get_ref_id(name: &str, naming_key: &[u8]) -> Result<Vec<u8>, rpc::errors::RefIDError> {
    let mut mac: Hmac<Sha256> = Hmac::new_from_slice(naming_key)?;
    mac.update(name.as_bytes());
    let ref_id = mac.finalize().into_bytes().to_vec();

    Ok(ref_id)
}

#[cfg(feature = "native")]
#[derive(thiserror::Error, Debug)]
pub enum LoginError {
    #[error("{0}")]
    ConfigError(#[from] config::errors::ConfigError),
}

/// Log into a CipherStash workspace.
///
/// If `workspace` is given, a profile will be created for that workspace.
/// This is useful for logging into a workspace for the first time.
///
/// If a workspace is not provided, this command will load an existing profile and then log in using that profile.
///
/// The profile name defaults to "default", but can be overridden with the `CS_PROFILE_NAME` env var or the `defaultProfile` opt in `{cs_config_path}/config.json`.
///
/// Returns the path to the directory of the profile that was either created or loaded.
///
/// Note: if there's an existing `Config` struct around in memory, it can potentially be stale after calling `login` and will need to be reloaded.
#[cfg(feature = "native")]
pub async fn login(
    workspace: Option<String>,
    region: Option<String>,
    profile_name: Option<String>,
    force: bool,
) -> Result<PathBuf, LoginError> {
    if let (Some(workspace), Some(region)) = (workspace, region) {
        let profile_dir_path =
            Config::create_profile(workspace, region, profile_name, force).await?;
        Ok(profile_dir_path)
    } else {
        let config = Config::load(config::LoadOpts {
            profile_name,
            ..Default::default()
        })?;

        config.refresh_token().await?;

        Ok(config.profile_dir_path())
    }
}

/// The interface between the client and the data service.
///
/// This struct manages encrypting and decrypting records and metadata, encrypting queries and
/// inserting and retrieving records.
///
/// It communicates with the data service through a gRPC client.
pub struct Stash<Client, Data> {
    data: Data,
    client: ApiClient<Client>,
    rng: ChaChaRng,
}

#[cfg(feature = "native")]
#[derive(thiserror::Error, Debug)]
pub enum CreateClientError {
    #[error("InvalidUri: {0}")]
    InvalidUri(String),
}

#[cfg(feature = "native")]
use self::client_data::NativeClientData;

#[cfg(feature = "native")]
impl Stash<tonic::transport::Channel, NativeClientData> {
    /// Try and create a native client from a resolved [`Config`]
    pub fn try_from_config(config: &Config) -> Result<Self, CreateClientError> {
        let data = NativeClientData {
            cipher_engine: EnvelopeCipher::init(config.key_provider()),
            access_token: config.data_service_credentials_provider(),
            naming_key: config.naming_key_provider(),
        };

        let conn = tonic::transport::Channel::from_shared(config.endpoint())
            .map_err(|e| CreateClientError::InvalidUri(e.to_string()))?
            // Create the connection lazily to avoid awaiting here. The client will connect
            // to the data service on the first call made by the client.
            .connect_lazy();

        Ok(Self::new(conn, data))
    }
}

#[cfg(feature = "web-worker")]
pub type StashClient = Stash<cipherstash_grpc::web_client::Client, BrowserClientData>;

#[cfg(feature = "native")]
pub type StashClient = Stash<tonic::transport::Channel, NativeClientData>;

#[cfg(feature = "web-worker")]
use crate::stash::client_data::BrowserClientData;

#[cfg(feature = "web-worker")]
impl Stash<cipherstash_grpc::web_client::Client, BrowserClientData> {
    pub fn new_web_client(
        base_uri: String,
        token: &str,
        key: &str,
    ) -> Result<Self, Box<dyn std::error::Error>> {
        use crate::data_service_credentials::DataServiceCredentials;
        use envelopers::SimpleKeyProvider;

        let client = cipherstash_grpc::web_client::Client::new(base_uri);
        let provider: SimpleKeyProvider = SimpleKeyProvider::init(
            hex::decode(&key)?
                .try_into()
                .map_err(|e: Vec<u8>| format!("Failed to convert key: len: {}", e.len()))?,
        );

        let data = BrowserClientData {
            cipher_engine: EnvelopeCipher::init(Box::new(provider)),
            access_token: DataServiceCredentials::new(token),
        };

        Ok(Self::new(client, data))
    }
}

impl<Client, Data> Stash<Client, Data>
where
    Client: tonic::client::GrpcService<tonic::body::BoxBody>,
    Client::ResponseBody: Body + Send + 'static,
    Client::Error: Into<StdError>,
    <Client::ResponseBody as Body>::Error: Into<StdError> + Send,

    Data: GetAccessToken + GetCipherEngine,
{
    /// Create a new [`Stash`] instance from a gRPC client, encryption key provider and access token provider.
    pub fn new(client: Client, data: Data) -> Self {
        let client = ApiClient::new(client);

        Self {
            client,
            data,
            rng: ChaChaRng::from_entropy(),
        }
    }

    async fn with_authentication<T>(
        &self,
        request: impl IntoRequest<T>,
    ) -> Result<Request<T>, rpc::errors::AuthorizationError> {
        let mut req = request.into_request();

        req.metadata_mut().insert(
            AsciiMetadataKey::from_static("authorization"),
            self.data
                .get_access_token()
                .await
                .map_err(|e| {
                    rpc::errors::AuthorizationError::new(format!(
                        "Failed to get access token: {}",
                        e
                    ))
                })?
                .as_header()
                .parse()
                .map_err(|_| {
                    rpc::errors::AuthorizationError::new("Invalid authorization header")
                })?,
        );

        Ok(req)
    }

    pub async fn delete_collection(
        &mut self,
        collection_ref: &[u8],
    ) -> Result<(), rpc::DeleteCollectionError> {
        self.client
            .delete_collection(
                self.with_authentication(collections::DeleteRequest {
                    r#ref: collection_ref.to_vec(),
                })
                .await?,
            )
            .await?
            .into_inner();

        Ok(())
    }

    pub async fn list_collections(
        &mut self,
    ) -> Result<rpc::ListCollectionsResponse, rpc::ListCollectionsError> {
        let res = self
            .client
            .collection_list(
                self.with_authentication(collections::ListRequest {})
                    .await?,
            )
            .await?
            .into_inner();

        let cipher = self.data.get_cipher_engine();

        let mut collections = vec![];

        for info in res.collections.into_iter() {
            collections.push(rpc::CollectionDecryptResult::decrypt(cipher, info).await)
        }

        Ok(rpc::ListCollectionsResponse { collections })
    }

    pub async fn get_all(
        &mut self,
        collection: &Collection,
        ids: Vec<Uuid>,
    ) -> Result<Vec<Record>, rpc::GetAllError> {
        let response = self
            .client
            .get_all(documents::GetAllRequest {
                collection_id: collection.id.as_bytes().to_vec(),
                ids: ids.into_iter().map(|id| id.as_bytes().to_vec()).collect(),
            })
            .await?
            .into_inner();

        let mut records: Vec<Record> = Vec::new();

        for document in response.documents.into_iter() {
            let record = Record::decrypt(self.data.get_cipher_engine(), document).await?;
            records.push(record);
        }

        Ok(records)
    }

    pub async fn delete_record(
        &mut self,
        collection: &Collection,
        id: Uuid,
    ) -> Result<(), rpc::DeleteRecordError> {
        self.client
            .delete(
                self.with_authentication(documents::DeleteRequest {
                    id: id.as_bytes().to_vec(),
                    collection_id: collection.id.as_bytes().to_vec(),
                })
                .await?,
            )
            .await?
            .into_inner();

        Ok(())
    }

    pub async fn put_record(
        &mut self,
        collection: &mut Collection,
        record: &Record,
    ) -> Result<(), rpc::PutRecordError> {
        let vectors = collection.indexer.encrypt(record)?;

        let PutReply { error } = self
            .client
            .put(
                self.with_authentication(documents::PutRequest {
                    vectors: vectors.into_iter().map(|x| x.into()).collect(),
                    collection_id: collection.id.as_bytes().to_vec(),
                    source: Some(record.encrypt(self.data.get_cipher_engine()).await?),
                    first_schema_version: collection.first_schema_version(),
                    last_schema_version: collection.last_schema_version(),
                })
                .await?,
            )
            .await?
            .into_inner();

        match PutError::from_i32(error).ok_or(rpc::PutRecordError::FailedToParseResponse(error))? {
            PutError::NoError => Ok(()),
            PutError::ErrUnknownSchemaVersion => Err(rpc::PutRecordError::UnknownSchemaVersion),
            PutError::ErrIncompleteSchemaVersionCoverage => {
                Err(rpc::PutRecordError::IncompleteSchemaVersionCoverage)
            }
        }
    }

    pub async fn put_all(
        &mut self,
        collection: &mut Collection,
        records: impl IntoIterator<Item = Record> + 'static,
    ) -> Result<usize, rpc::PutStreamError> {
        self.put_stream(collection, stream::iter(records)).await
    }

    pub async fn put_stream(
        &mut self,
        collection: &mut Collection,
        record_stream: impl Stream<Item = Record> + 'static,
    ) -> Result<usize, rpc::PutStreamError> {
        // Because it's not possible to send most of the `Stash` struct between threads and
        // `put_stream` requires futures that are Send this method has been written as two futures.
        //
        // The first future takes the record stream and encrypts the records and then sends them
        // using the following futures::channel to the second future that sends them over the wire.
        //
        // It would be possible to make the `Stash` struct Send + Sync however it would be a lot of
        // work for not a lot of gain just yet.
        let (encrypted_record_send, encrypted_record_recv) = unbounded::<StreamingPutRequest>();

        let mut req = Request::new(
            stream::iter([
                // Start the streaming put request by specifying the collection id / schema
                // versions and what not
                documents::StreamingPutRequest {
                    payload: Some(StreamingPutPayload::Begin(StreamingPutBegin {
                        collection_id: collection.id.as_bytes().to_vec(),
                        first_schema_version: collection.first_schema_version(),
                        last_schema_version: collection.last_schema_version(),
                    })),
                },
            ])
            // After sending the begin message - chain the encrypted records coming from the other
            // future.
            .chain(encrypted_record_recv),
        );

        req = self.with_authentication(req).await?;

        // Future for consuming the input stream, encrypting the records and sending them to the
        // "response future" to be inserted into the database.
        let send_message_fut = async {
            let mut stream = Box::pin(record_stream);

            while let Some(record) = stream.next().await {
                let vectors = collection.indexer.encrypt(&record)?;

                // Encrypt and send the document to the "response future".
                // If this errors it will close the future and drop and close the channel.
                encrypted_record_send.unbounded_send(documents::StreamingPutRequest {
                    payload: Some(StreamingPutPayload::Document(StreamingPutDocument {
                        source: Some(record.encrypt(self.data.get_cipher_engine()).await?),
                        vectors: vectors.into_iter().map(|x| x.into()).collect(),
                    })),
                })?;
            }

            // Close the channel to make sure the encrypted_record_recv stream completes.
            encrypted_record_send.close_channel();

            // There are no more records left in the input stream so return without error.
            Ok::<(), rpc::PutStreamError>(())
        };

        let response_fut = self
            .client
            .put_stream(req)
            .map_err(rpc::PutStreamError::from);

        let StreamingPutReply {
            error,
            num_inserted,
        } =
            // Try and consume both futures until completion.
            // If either of the futures error out this method will return Err.
            try_join!(response_fut, send_message_fut)?.0.into_inner();

        match StreamingPutError::from_i32(error)
            .ok_or(rpc::PutStreamError::FailedToParseResponse(error))?
        {
            StreamingPutError::NoError => Ok(num_inserted as _),
            StreamingPutError::ErrUnknownSchemaVersion => {
                Err(rpc::PutStreamError::UnknownSchemaVersion)
            }
            StreamingPutError::ErrIncompleteSchemaVersionCoverage => {
                Err(rpc::PutStreamError::IncompleteSchemaVersionCoverage)
            }
        }
    }

    pub async fn get_record(
        &mut self,
        collection: &Collection,
        id: Uuid,
    ) -> Result<Option<Record>, rpc::GetRecordError> {
        let GetReply { source } = self
            .client
            .get(
                self.with_authentication(documents::GetRequest {
                    id: id.as_bytes().to_vec(),
                    collection_id: collection.id.as_bytes().to_vec(),
                })
                .await?,
            )
            .await?
            .into_inner();

        Ok(if let Some(doc) = source {
            Some(Record::decrypt(self.data.get_cipher_engine(), doc).await?)
        } else {
            None
        })
    }

    pub async fn query(
        &mut self,
        collection: &mut Collection,
        query: QueryInput,
    ) -> Result<rpc::QueryResponse, rpc::QueryError> {
        let constraints = collection
            .indexer
            .encrypt_query(query.query)?
            .into_iter()
            .map(|c| Constraint {
                index_id: c.index_id.as_bytes().to_vec(),
                condition: match c.condition {
                    crate::indexer::vector::ConstraintCondition::Exact { term } => {
                        Some(constraint::Condition::Exact(Exact {
                            term: term.0.into_iter().map(|t| t.to_bytes()).collect(),
                        }))
                    }
                    crate::indexer::vector::ConstraintCondition::Range { upper, lower } => {
                        Some(constraint::Condition::Range(Range {
                            upper: upper.0.into_iter().map(|t| t.to_bytes()).collect(),
                            lower: lower.0.into_iter().map(|t| t.to_bytes()).collect(),
                        }))
                    }
                    crate::indexer::vector::ConstraintCondition::Filter { bits } => {
                        Some(constraint::Condition::Filter(Filter {
                            bits: bits.into_iter().map(|x| x.into()).collect(),
                        }))
                    }
                },
            })
            .collect();

        let aggregates = query.aggregates.into_iter().map(|a| a.into()).collect();

        let ordering: Vec<Ordering> = query
            .ordering
            .into_iter()
            .map(|o| {
                collection
                    .indexer
                    .get_index_id(&o.index_name)
                    .ok_or(rpc::QueryError::InvalidIndexName(o.index_name))
                    .map(|id| Ordering {
                        index_id: id.as_bytes().to_vec(),
                        direction: o.direction.into(),
                    })
            })
            .collect::<Result<_, _>>()?;

        let enc_query = Query {
            limit: query.limit,
            constraints,
            aggregates,
            ordering,
            skip_results: query.skip_results,
            offset: query.offset,
        };

        let query_reply: QueryReply = self
            .client
            .query(
                self.with_authentication(QueryRequest {
                    collection_id: collection.id.as_bytes().to_vec(),
                    query: Some(enc_query),
                    schema_version: collection.last_schema_version(),
                })
                .await?,
            )
            .await?
            .into_inner();

        match QueryError::from_i32(query_reply.error)
            .ok_or(rpc::QueryError::FailedToParseResponse(query_reply.error))?
        {
            QueryError::NoError => {
                let mut records = Vec::new();
                for document in query_reply.records {
                    let record = Record::decrypt(self.data.get_cipher_engine(), document).await?;
                    records.push(record);
                }

                let aggregates = query_reply
                    .aggregates
                    .into_iter()
                    .map(|a| a.into())
                    .collect();

                Ok(rpc::QueryResponse {
                    aggregates,
                    records,
                })
            }
            QueryError::ErrUnknownSchemaVersion => Err(rpc::QueryError::UnknownSchemaVersion),
            QueryError::ErrObsoleteSchemaVersion => Err(rpc::QueryError::ErrObsoleteSchemaVersion),
        }
    }
}

impl<Client, Data> Stash<Client, Data>
where
    Client: tonic::client::GrpcService<tonic::body::BoxBody>,
    Client::ResponseBody: Body + Send + 'static,
    Client::Error: Into<StdError>,
    <Client::ResponseBody as Body>::Error: Into<StdError> + Send,

    Data: GetAccessToken + GetCipherEngine + GetNamingKey,
{
    pub async fn collection_info(
        &mut self,
        name: String,
    ) -> Result<CollectionDecryptResult, rpc::CollectionInfoError> {
        let r#ref = get_ref_id(&name, &self.data.get_naming_key().await?)?;

        let info_reply = self
            .client
            .collection_info(
                self.with_authentication(collections::InfoRequest {
                    r#ref,
                    metadata: Vec::new(),
                })
                .await?,
            )
            .await?
            .into_inner();

        let result =
            rpc::CollectionDecryptResult::decrypt(self.data.get_cipher_engine(), info_reply).await;

        Ok(result)
    }

    pub async fn create_collection(
        &mut self,
        name: String,
        schema: RecordSchema,
        indexes: impl IntoIterator<Item = (String, Mapping)>,
    ) -> Result<Collection, rpc::CreateCollectionError> {
        let rng = &mut self.rng;
        let cipher = self.data.get_cipher_engine();

        let mut encrypted_indexes: Vec<api::indexes::Index> = vec![];

        for (name, mapping) in indexes.into_iter() {
            let index = DecryptedIndex::new_from_mapping(rng, name, mapping)?;
            encrypted_indexes.push(index.encrypt(cipher).await?);
        }

        let r#ref = get_ref_id(&name, &self.data.get_naming_key().await?)?;

        let metadata = DecryptedCollectionMetadata {
            name,
            record_type: schema,
        };

        let info_reply = self
            .client
            .create_collection(
                self.with_authentication(collections::CreateRequest {
                    indexes: encrypted_indexes,
                    r#ref,
                    metadata: metadata.encrypt(cipher).await?,
                    schema_version_drop_delay: 0,
                })
                .await?,
            )
            .await?
            .into_inner();

        let result = rpc::CollectionDecryptResult::decrypt(cipher, info_reply).await;

        Ok(result.try_into()?)
    }
}

#[cfg(test)]
mod tests {
    use envelopers::{EnvelopeCipher, KeyProvider, SimpleKeyProvider};
    use std::future::Future;
    use uuid::uuid;

    use crate::async_value::StaticValue;
    use crate::data_service_credentials::DataServiceCredentials;
    use crate::record;

    use super::{
        client_data::NativeClientData,
        rpc::test_client::{TestClient, TestClientBuilder},
        *,
    };

    #[test]
    fn test_get_ref_id_returns_correct() {
        let naming_key: [u8; 32] = [1; 32];
        let result = get_ref_id("movies", &naming_key).unwrap();

        let expected: Vec<u8> = vec![
            153, 4, 166, 247, 23, 159, 65, 142, 187, 193, 188, 84, 36, 77, 72, 45, 143, 129, 37,
            18, 2, 3, 243, 245, 200, 27, 239, 45, 28, 38, 195, 31,
        ];

        assert_eq!(result, expected);
    }

    async fn create_test_client<F: Future<Output = TestClientBuilder>>(
        callback: impl FnOnce(TestClientBuilder, EnvelopeCipher<Box<dyn KeyProvider>>) -> F,
    ) -> Stash<TestClient, NativeClientData> {
        let callback_provider: SimpleKeyProvider = SimpleKeyProvider::init([1; 16]);

        let callback_cipher_engine: EnvelopeCipher<Box<dyn KeyProvider>> =
            EnvelopeCipher::init(Box::new(callback_provider));

        let provider: SimpleKeyProvider = SimpleKeyProvider::init([1; 16]);

        let cipher_engine: EnvelopeCipher<Box<dyn KeyProvider>> =
            EnvelopeCipher::init(Box::new(provider));

        let access_token = Box::new(StaticValue::new(DataServiceCredentials::new("abcd")));
        let naming_key = Box::new(StaticValue::new(vec![]));

        let test_client = callback(TestClient::builder(), callback_cipher_engine)
            .await
            .build();

        let data = NativeClientData {
            cipher_engine,
            access_token,
            naming_key,
        };

        Stash::new(test_client, data)
    }

    #[tokio::test]
    async fn test_get_record() {
        let record = record! {
            "id" => uuid!("00000000-0000-0000-0000-000000000000"),
            "name" => "John",
            "age" => 20
        };

        let mut client = create_test_client(|c, cipher| {
            let record = &record;

            async move {
                c.add_success_response(
                    "get",
                    GetReply {
                        source: Some(
                            record
                                .encrypt(&cipher)
                                .await
                                .expect("Failed to encrypt record"),
                        ),
                    },
                )
            }
        })
        .await;

        let record_response = client
            .get_record(
                &Collection::new_test_collection(),
                uuid!("00000000-0000-0000-0000-000000000000"),
            )
            .await
            .expect("Failed to get record");

        assert_eq!(record_response, Some(record));
    }

    #[tokio::test]
    async fn test_put_invalid_error_code() {
        let mut client = create_test_client(|c, _| async move {
            c.add_success_response("put", PutReply { error: 100 })
        })
        .await;

        let error_response = client
            .put_record(
                &mut Collection::new_test_collection(),
                &record! {
                    "id" => uuid!("00000000-0000-0000-0000-000000000000"),
                },
            )
            .await
            .expect_err("Expected put record to fail");

        assert_eq!(
            error_response.to_string(),
            "Failed to parse error response enum. Got: 100"
        );
    }

    #[tokio::test]
    async fn test_put_no_error() {
        let mut client = create_test_client(|c, _| async move {
            c.add_success_response("put", PutReply { error: 0 })
        })
        .await;

        let response = client
            .put_record(
                &mut Collection::new_test_collection(),
                &record! {
                    "id" => uuid!("00000000-0000-0000-0000-000000000000"),
                },
            )
            .await
            .expect("Failed to put record");

        assert_eq!(response, ());
    }
}