// Copyright © 2019 sqsquatch developers
//
// Licensed under the Apache License, Version 2.0
// <LICENSE-APACHE or http://www.apache.org/licenses/LICENSE-2.0> or the MIT
// license <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. All files in the project carrying such notice may not be copied,
// modified, or distributed except according to those terms.

//! Subscribe to an SQS queue for message retrieval

use crate::error::{Error, Result};
use futures01::future::Future;
use lazy_static::lazy_static;
use rusoto_core::{
    credential::ProvideAwsCredentials, request::HttpClient, DispatchSignedRequest, Region,
    RusotoError,
};
use rusoto_credential::DefaultCredentialsProvider;
use rusoto_mock::{MockResponseReader, ReadMockResponse};
use rusoto_sqs::{
    DeleteMessageError, DeleteMessageRequest, GetQueueUrlError, GetQueueUrlRequest,
    GetQueueUrlResult, ReceiveMessageError, ReceiveMessageRequest, ReceiveMessageResult, Sqs,
    SqsClient,
};
use std::collections::HashMap;
use std::fmt;

/// Subscribe to as SQS queue
#[derive(Clone)]
pub struct Sub {
    client: SqsClient,
    queue_name: String,
}

impl fmt::Debug for Sub {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "Sub -> Queue({})", self.queue_name)
    }
}

impl Sub {
    /// Initialize [`Sub`]
    ///
    /// The following sources are checked in order for AWS credentials when calling `initialize`:
    ///
    /// 1. Environment variables: `AWS_ACCESS_KEY_ID` and `AWS_SECRET_ACCESS_KEY`
    /// 2. `credential_process` command in the AWS config file, usually located at `~/.aws/config`.
    /// 3. AWS credentials file. Usually located at `~/.aws/credentials`.
    /// 4. IAM instance profile. Will only work if running on an EC2 instance with an instance profile/role.
    ///
    /// If the sources are exhausted without finding credentials, an error is returned.
    ///
    /// See the documentation for [`DefaultCredentialsProvider`](https://docs.rs/rusoto_credential/latest/rusoto_credential/struct.DefaultCredentialsProvider.html) and [`ChainProvider`](https://docs.rs/rusoto_credential/latest/rusoto_credential/struct.ChainProvider.html) for more information.
    ///
    /// # Example
    /// ```
    /// # use sqsquatch::{Sub, Result};
    /// # use rusoto_core::Region;
    /// #
    /// # fn main() -> Result<()> {
    /// let subscriber = Sub::initialize(Region::UsEast2, "Orders")?;
    /// # Ok(())
    /// # }
    /// ```
    ///
    pub fn initialize<T>(region: Region, queue_name: T) -> Result<Self>
    where
        T: Into<String>,
    {
        let provider = DefaultCredentialsProvider::new()?;
        let rusoto_client = HttpClient::from_connector(super::new_connector()?);
        Self::initialize_internal(region, queue_name, provider, rusoto_client)
    }

    #[doc(hidden)]
    pub fn initialize_internal<P, R, T>(
        region: Region,
        queue_name: T,
        provider: P,
        request_dispatcher: R,
    ) -> Result<Self>
    where
        P: ProvideAwsCredentials + Send + Sync + 'static,
        P::Future: Send,
        R: DispatchSignedRequest + Send + Sync + 'static,
        R::Future: Send,
        T: Into<String>,
    {
        let client = SqsClient::new_with(request_dispatcher, provider, region);

        Ok(Self {
            client,
            queue_name: queue_name.into(),
        })
    }

    /// Receive a message from the queue.
    ///
    /// * `visibility_timeout` - the duration (in seconds) that the received messages are hidden from subsequent retrieve requests after being retrieved by a `receive_message` request.
    /// * `wait_time_seconds` - The duration (in seconds) for which the call waits for a message to arrive in the queue before returning. If a message is available, the call returns sooner than `wait_time_seconds`. If no messages are available and the wait time expires, the call returns successfully with an empty list of messages.
    ///
    /// # Example
    /// ```
    /// # use sqsquatch::{MockType, Sub, Result, utils::MockRequestDispatcher};
    /// # use rusoto_core::Region;
    /// # use rusoto_mock::MockCredentialsProvider;
    /// # use uuid::Uuid;
    /// #
    /// # fn main() -> Result<()> {
    /// #   let mut subscriber = Sub::initialize_internal(
    /// #     Region::UsEast2,
    /// #     "Orders",
    /// #     MockCredentialsProvider,
    /// #     MockRequestDispatcher::default().with_body_map(Sub::mock_responses(MockType::Receive)),
    /// #   )?;
    /// // subscriber initialized as above..
    /// let result = subscriber.receive_message(None, Some(30), Some(20))?;
    /// assert!(result.messages.is_some());
    ///
    /// if let Some(messages) = result.messages {
    ///     assert_eq!(messages.len(), 1);
    ///     assert!(messages[0].receipt_handle.is_some());
    ///
    ///     if let Some(receipt_handle) = &messages[0].receipt_handle {
    ///         assert_eq!(receipt_handle, "MbZj6wDWli+JvwwJaBV+3dcjk2YW2vA3+STFFljTM8tJJg6HRG6PYSasuWXPJB+CwLj1FjgXUv1uSj1gUPAWV66FU/WeR4mq2OKpEGYWbnLmpRCJVAyeMjeU5ZBdtcQ+QEauMZc8ZRv37sIW2iJKq3M9MFx1YvV11A2x/KSbkJ0=");
    ///     }
    /// }
    /// # Ok(())
    /// # }
    /// ```
    pub fn receive_message(
        &mut self,
        message_attribute_names: Option<Vec<String>>,
        visibility_timeout: Option<i64>,
        wait_time_seconds: Option<i64>,
    ) -> Result<ReceiveMessageResult> {
        let queue_url_result = self
            .client
            .get_queue_url(super::get_queue_url_request(self.queue_name.clone()))
            .sync()?;

        if let Some(queue_url) = queue_url_result.queue_url {
            let receive_result = self
                .client
                .receive_message(receive_message_request(
                    queue_url,
                    message_attribute_names,
                    visibility_timeout,
                    wait_time_seconds,
                ))
                .sync()?;

            Ok(receive_result)
        } else {
            Err(Error::publish_error())
        }
    }

    /// Delete a message from the queue.
    ///
    /// * `receipt_handle` - The receipt handle associated with the message to delete. **NOTE** this is not the message id.  It's a unique identifier on a per-receive request.
    ///
    /// # Example
    /// ```
    /// # use sqsquatch::{MockType, Sub, Result, utils::MockRequestDispatcher};
    /// # use rusoto_core::Region;
    /// # use rusoto_mock::MockCredentialsProvider;
    /// # use uuid::Uuid;
    /// #
    /// # fn main() -> Result<()> {
    /// #   let mut subscriber = Sub::initialize_internal(
    /// #     Region::UsEast2,
    /// #     "Orders",
    /// #     MockCredentialsProvider,
    /// #     MockRequestDispatcher::default().with_body_map(Sub::mock_responses(MockType::Delete)),
    /// #   )?;
    /// // subscriber initialized as above..
    /// let result = subscriber.delete_message("a receipt_handle".to_string());
    /// assert!(result.is_ok());
    /// # Ok(())
    /// # }
    /// ```
    pub fn delete_message(&mut self, receipt_handle: String) -> Result<()> {
        let queue_url_result = self
            .client
            .get_queue_url(super::get_queue_url_request(self.queue_name.clone()))
            .sync()?;

        if let Some(queue_url) = queue_url_result.queue_url {
            let _ = self
                .client
                .delete_message(DeleteMessageRequest {
                    receipt_handle,
                    queue_url: queue_url,
                })
                .sync()?;

            Ok(())
        } else {
            Err(Error::publish_error())
        }
    }

    #[doc(hidden)]
    #[must_use]
    pub fn mock_responses(kind: MockType) -> HashMap<String, Vec<u8>> {
        lazy_static! {
            static ref BODY_MAP: HashMap<String, Vec<u8>> = {
                let mut body_map = HashMap::new();
                let _ = body_map.insert(
                    "GetQueueUrl",
                    MockResponseReader::read_response("test-data", "get_queue_url_response.xml"),
                );
                let _ = body_map.insert(
                    "ReceiveMessage",
                    MockResponseReader::read_response("test-data", "receive_message_response.xml"),
                );
                body_map
                    .iter()
                    .map(|(k, v)| ((*k).to_string(), v.as_bytes().to_vec()))
                    .collect()
            };
            static ref DELETE_BODY_MAP: HashMap<String, Vec<u8>> = {
                let mut body_map = HashMap::new();
                let _ = body_map.insert(
                    "GetQueueUrl",
                    MockResponseReader::read_response("test-data", "get_queue_url_response.xml"),
                );
                let _ = body_map.insert(
                    "DeleteMessage",
                    MockResponseReader::read_response("test-data", "delete_message_response.xml"),
                );
                body_map
                    .iter()
                    .map(|(k, v)| ((*k).to_string(), v.as_bytes().to_vec()))
                    .collect()
            };
        }

        match kind {
            MockType::Receive => BODY_MAP.clone(),
            MockType::Delete => DELETE_BODY_MAP.clone(),
        }
    }
}

/// Create futures that will intract with an SQS queue in a subscribe manner
#[derive(Clone)]
pub struct SubFut {
    client: SqsClient,
    queue_name: String,
}

impl fmt::Debug for SubFut {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "Sub -> Queue({})", self.queue_name)
    }
}

impl SubFut {
    /// Initialize [`SubFut`]
    ///
    /// The following sources are checked in order for AWS credentials when calling `initialize`:
    ///
    /// 1. Environment variables: `AWS_ACCESS_KEY_ID` and `AWS_SECRET_ACCESS_KEY`
    /// 2. `credential_process` command in the AWS config file, usually located at `~/.aws/config`.
    /// 3. AWS credentials file. Usually located at `~/.aws/credentials`.
    /// 4. IAM instance profile. Will only work if running on an EC2 instance with an instance profile/role.
    ///
    /// If the sources are exhausted without finding credentials, an error is returned.
    ///
    /// See the documentation for [`DefaultCredentialsProvider`](https://docs.rs/rusoto_credential/latest/rusoto_credential/struct.DefaultCredentialsProvider.html) and [`ChainProvider`](https://docs.rs/rusoto_credential/latest/rusoto_credential/struct.ChainProvider.html) for more information.
    ///
    /// # Example
    /// ```
    /// # use sqsquatch::{SubFut, Result};
    /// # use rusoto_core::Region;
    /// #
    /// # fn main() -> Result<()> {
    /// let subscriber = SubFut::initialize(Region::UsEast2, "Orders")?;
    /// # Ok(())
    /// # }
    /// ```
    ///
    pub fn initialize<T>(region: Region, queue_name: T) -> Result<Self>
    where
        T: Into<String>,
    {
        let provider = DefaultCredentialsProvider::new()?;
        let rusoto_client = HttpClient::from_connector(super::new_connector()?);
        Self::initialize_internal(region, queue_name, provider, rusoto_client)
    }

    #[doc(hidden)]
    pub fn initialize_internal<P, R, T>(
        region: Region,
        queue_name: T,
        provider: P,
        request_dispatcher: R,
    ) -> Result<Self>
    where
        P: ProvideAwsCredentials + Send + Sync + 'static,
        P::Future: Send,
        R: DispatchSignedRequest + Send + Sync + 'static,
        R::Future: Send,
        T: Into<String>,
    {
        let client = SqsClient::new_with(request_dispatcher, provider, region);

        Ok(Self {
            client,
            queue_name: queue_name.into(),
        })
    }

    /// Create a future that will receive a message from the queue.
    ///
    /// * `visibility_timeout` - the duration (in seconds) that the received messages are hidden from subsequent retrieve requests after being retrieved by a `receive_message` request.
    /// * `wait_time_seconds` - The duration (in seconds) for which the call waits for a message to arrive in the queue before returning. If a message is available, the call returns sooner than `wait_time_seconds`. If no messages are available and the wait time expires, the call returns successfully with an empty list of messages.
    ///
    /// # Example
    /// ```
    /// # use futures::compat::Future01CompatExt;
    /// # use sqsquatch::{MockType, SubFut, Result, utils::MockRequestDispatcher};
    /// # use rusoto_core::Region;
    /// # use rusoto_mock::MockCredentialsProvider;
    /// # use uuid::Uuid;
    /// #
    /// # #[tokio::main]
    /// # async fn main() -> Result<()> {
    /// #   let mut subscriber = SubFut::initialize_internal(
    /// #     Region::UsEast2,
    /// #     "Orders",
    /// #     MockCredentialsProvider,
    /// #     MockRequestDispatcher::default().with_body_map(SubFut::mock_responses(MockType::Receive)),
    /// #   )?;
    /// // subscriber initialized as above..
    /// let result = subscriber.receive_message_fut(None, Some(30), Some(20)).compat().await?;
    /// assert!(result.messages.is_some());
    ///
    /// if let Some(messages) = result.messages {
    ///     assert_eq!(messages.len(), 1);
    ///     assert!(messages[0].receipt_handle.is_some());
    ///
    ///     if let Some(receipt_handle) = &messages[0].receipt_handle {
    ///         assert_eq!(receipt_handle, "MbZj6wDWli+JvwwJaBV+3dcjk2YW2vA3+STFFljTM8tJJg6HRG6PYSasuWXPJB+CwLj1FjgXUv1uSj1gUPAWV66FU/WeR4mq2OKpEGYWbnLmpRCJVAyeMjeU5ZBdtcQ+QEauMZc8ZRv37sIW2iJKq3M9MFx1YvV11A2x/KSbkJ0=");
    ///     }
    /// }
    /// # Ok(())
    /// # }
    /// ```
    pub fn receive_message_fut(
        &mut self,
        message_attribute_names: Option<Vec<String>>,
        visibility_timeout: Option<i64>,
        wait_time_seconds: Option<i64>,
    ) -> impl Future<Item = ReceiveMessageResult, Error = Error> {
        let get_queue_url = SubFut::make_get_queue_url_future(
            &self.client,
            super::get_queue_url_request(self.queue_name.clone()),
        );
        let client_clone = self.client.clone();
        let receive_message = move |result: GetQueueUrlResult| {
            client_clone.receive_message(receive_message_request(
                result.queue_url.unwrap(),
                message_attribute_names,
                visibility_timeout,
                wait_time_seconds,
            ))
        };

        get_queue_url
            .map_err(map_re)
            .and_then(receive_message)
            .map_err(map_e)
    }

    /// Delete a message from the queue.
    ///
    /// * `receipt_handle` - The receipt handle associated with the message to delete. **NOTE** this is not the message id.  It's a unique identifier on a per-receive request.
    ///
    /// # Example
    /// ```
    /// # use futures::compat::Future01CompatExt;
    /// # use sqsquatch::{MockType, SubFut, Result, utils::MockRequestDispatcher};
    /// # use rusoto_core::Region;
    /// # use rusoto_mock::MockCredentialsProvider;
    /// # use uuid::Uuid;
    /// #
    /// # #[tokio::main]
    /// # async fn main() -> Result<()> {
    /// #   let mut subscriber = SubFut::initialize_internal(
    /// #     Region::UsEast2,
    /// #     "Orders",
    /// #     MockCredentialsProvider,
    /// #     MockRequestDispatcher::default().with_body_map(SubFut::mock_responses(MockType::Delete)),
    /// #   )?;
    /// // subscriber initialized as above..
    /// let _ = subscriber.delete_message_fut("a receipt_handle".to_string()).compat().await?;
    /// # Ok(())
    /// # }
    /// ```
    pub fn delete_message_fut(
        &mut self,
        receipt_handle: String,
    ) -> impl Future<Item = (), Error = Error> {
        let get_queue_url = SubFut::make_get_queue_url_future(
            &self.client,
            super::get_queue_url_request(self.queue_name.clone()),
        );
        let client_clone = self.client.clone();
        let delete_message = move |result: GetQueueUrlResult| {
            client_clone.delete_message(DeleteMessageRequest {
                receipt_handle,
                queue_url: result.queue_url.unwrap(),
            })
        };

        get_queue_url
            .map_err(map_del)
            .and_then(delete_message)
            .map_err(map_f)
    }

    fn make_get_queue_url_future(
        client: &SqsClient,
        request: GetQueueUrlRequest,
    ) -> impl Future<Item = GetQueueUrlResult, Error = RusotoError<GetQueueUrlError>> {
        client.get_queue_url(request)
    }

    #[doc(hidden)]
    #[must_use]
    pub fn mock_responses(kind: MockType) -> HashMap<String, Vec<u8>> {
        lazy_static! {
            static ref BODY_MAP: HashMap<String, Vec<u8>> = {
                let mut body_map = HashMap::new();
                let _ = body_map.insert(
                    "GetQueueUrl",
                    MockResponseReader::read_response("test-data", "get_queue_url_response.xml"),
                );
                let _ = body_map.insert(
                    "ReceiveMessage",
                    MockResponseReader::read_response("test-data", "receive_message_response.xml"),
                );
                body_map
                    .iter()
                    .map(|(k, v)| ((*k).to_string(), v.as_bytes().to_vec()))
                    .collect()
            };
            static ref DELETE_BODY_MAP: HashMap<String, Vec<u8>> = {
                let mut body_map = HashMap::new();
                let _ = body_map.insert(
                    "GetQueueUrl",
                    MockResponseReader::read_response("test-data", "get_queue_url_response.xml"),
                );
                let _ = body_map.insert(
                    "DeleteMessage",
                    MockResponseReader::read_response("test-data", "delete_message_response.xml"),
                );
                body_map
                    .iter()
                    .map(|(k, v)| ((*k).to_string(), v.as_bytes().to_vec()))
                    .collect()
            };
        }

        match kind {
            MockType::Receive => BODY_MAP.clone(),
            MockType::Delete => DELETE_BODY_MAP.clone(),
        }
    }
}

#[doc(hidden)]
#[derive(Clone, Copy, Debug)]
pub enum MockType {
    Receive,
    Delete,
}

#[allow(clippy::needless_pass_by_value)]
fn map_re(e: RusotoError<GetQueueUrlError>) -> RusotoError<ReceiveMessageError> {
    RusotoError::Service(ReceiveMessageError::OverLimit(e.to_string()))
}

#[allow(clippy::needless_pass_by_value)]
fn map_del(e: RusotoError<GetQueueUrlError>) -> RusotoError<DeleteMessageError> {
    RusotoError::Service(DeleteMessageError::ReceiptHandleIsInvalid(e.to_string()))
}

fn map_e(e: RusotoError<ReceiveMessageError>) -> Error {
    e.into()
}

fn map_f(e: RusotoError<DeleteMessageError>) -> Error {
    e.into()
}

#[allow(clippy::needless_pass_by_value)]
fn receive_message_request(
    queue_url: String,
    _message_attribute_names: Option<Vec<String>>,
    visibility_timeout: Option<i64>,
    wait_time_seconds: Option<i64>,
) -> ReceiveMessageRequest {
    ReceiveMessageRequest {
        attribute_names: Some(vec!["All".to_string()]),
        max_number_of_messages: Some(1),
        message_attribute_names: None,
        queue_url,
        receive_request_attempt_id: None,
        visibility_timeout,
        wait_time_seconds,
    }
}

#[cfg(test)]
mod test {
    use super::{MockType, Sub};
    use crate::{error::Result, utils::MockRequestDispatcher};
    use rusoto_core::Region;
    use rusoto_mock::MockCredentialsProvider;

    #[test]
    fn receive_message() -> Result<()> {
        let _ = pretty_env_logger::try_init_timed();
        let mut subscriber = Sub::initialize_internal(
            Region::UsEast2,
            "Orders",
            MockCredentialsProvider,
            MockRequestDispatcher::default().with_body_map(Sub::mock_responses(MockType::Receive)),
        )?;
        assert!(subscriber.receive_message(None, None, None).is_ok());
        Ok(())
    }

    #[test]
    fn delete() -> Result<()> {
        let _ = pretty_env_logger::try_init_timed();
        let mut subscriber = Sub::initialize_internal(
            Region::UsEast2,
            "Orders",
            MockCredentialsProvider,
            MockRequestDispatcher::default().with_body_map(Sub::mock_responses(MockType::Delete)),
        )?;
        assert!(subscriber.delete_message("".to_string()).is_ok());
        Ok(())
    }
}