// Copyright (c) Microsoft Corporation.
// Licensed under the MIT License.

use std::collections::HashMap;
use std::convert::TryFrom;
use std::sync::{Arc, Mutex};

use async_std::future;
use azure_iot_operations_mqtt::azure_mqtt;
use azure_iot_operations_mqtt::session::{SessionManagedClient, SessionMonitor};
use azure_iot_operations_protocol::application::ApplicationContextBuilder;
use azure_iot_operations_protocol::common::aio_protocol_error::{
    AIOProtocolError, AIOProtocolErrorKind,
};
use azure_iot_operations_protocol::rpc_command;
use bytes::Bytes;
use serde_json;
use tokio::time;
use uuid::Uuid;

use crate::metl::aio_protocol_error_checker;
use crate::metl::countdown_event_map::CountdownEventMap;
use crate::metl::defaults::ExecutorDefaults;
use crate::metl::mqtt_hub::{MqttHub, to_is_utf8};
use crate::metl::qos::{self, new_packet_identifier_dup_qos};
use crate::metl::test_case::TestCase;
use crate::metl::test_case_action::TestCaseAction;
use crate::metl::test_case_catch::TestCaseCatch;
use crate::metl::test_case_executor::TestCaseExecutor;
use crate::metl::test_case_published_message::TestCasePublishedMessage;
use crate::metl::test_case_serializer::TestCaseSerializer;
use crate::metl::test_payload::TestPayload;

const TEST_TIMEOUT: time::Duration = time::Duration::from_secs(30);

// TODO: this struct could probably be removed
pub struct CommandExecutorTester {}

impl CommandExecutorTester {
    pub async fn test_command_executor(
        test_case: TestCase<ExecutorDefaults>,
        managed_client: SessionManagedClient,
        session_monitor: SessionMonitor,
        mut mqtt_hub: MqttHub,
    ) {
        if let Some(push_acks) = test_case.prologue.push_acks.as_ref() {
            for ack_kind in &push_acks.publish {
                mqtt_hub.enqueue_puback(ack_kind.clone());
            }

            for ack_kind in &push_acks.subscribe {
                mqtt_hub.enqueue_suback(ack_kind.clone());
            }

            for ack_kind in &push_acks.unsubscribe {
                mqtt_hub.enqueue_unsuback(ack_kind.clone());
            }
        }

        // force connack to happen before other events are injected
        mqtt_hub.await_operation().await;
        session_monitor.connected().await;

        let mut countdown_events = CountdownEventMap::new();
        for (event_name, &init_count) in &test_case.prologue.countdown_events {
            countdown_events.insert(event_name.clone(), init_count);
        }

        let executor_count = test_case.prologue.executors.len();
        let mut execution_counts = Vec::with_capacity(executor_count);

        let mut ix = 0;
        for test_case_executor in &test_case.prologue.executors {
            ix += 1;
            let catch = if ix == executor_count {
                test_case.prologue.catch.as_ref()
            } else {
                None
            };

            let execution_count = Arc::new(Mutex::new(0));
            execution_counts.push(execution_count.clone());

            if let Some(executor) = Self::get_command_executor(
                managed_client.clone(),
                test_case_executor,
                catch,
                &mut mqtt_hub,
            )
            .await
            {
                tokio::task::spawn(Self::executor_loop(
                    executor,
                    (*test_case_executor).clone(),
                    countdown_events.clone(),
                    execution_count,
                ));
            }
        }

        let test_case_serializer = &test_case.prologue.executors[0].serializer;

        let mut source_ids: HashMap<i32, Uuid> = HashMap::new();
        let mut correlation_ids: HashMap<i32, Option<azure_mqtt::mqtt_proto::BinaryData<Bytes>>> =
            HashMap::new();
        let mut packet_ids: HashMap<i32, u16> = HashMap::new();

        for test_case_action in &test_case.actions {
            match test_case_action {
                action_receive_request @ TestCaseAction::ReceiveRequest { .. } => {
                    Self::receive_request(
                        action_receive_request,
                        &mut mqtt_hub,
                        &mut source_ids,
                        &mut correlation_ids,
                        &mut packet_ids,
                        test_case_serializer,
                    );
                }
                action_await_ack @ TestCaseAction::AwaitAck { .. } => {
                    Self::await_acknowledgement(action_await_ack, &mut mqtt_hub, &packet_ids).await;
                }
                action_await_publish @ TestCaseAction::AwaitPublish { .. } => {
                    Self::await_publish(action_await_publish, &mut mqtt_hub, &correlation_ids)
                        .await;
                }
                action_sync @ TestCaseAction::Sync { .. } => {
                    Self::sync(action_sync, &countdown_events).await;
                }
                action_sleep @ TestCaseAction::Sleep { .. } => {
                    Self::sleep(action_sleep).await;
                }
                _action_disconnect @ TestCaseAction::Disconnect { .. } => {
                    Self::disconnect(&mut mqtt_hub);
                }
                _action_freeze_time @ TestCaseAction::FreezeTime { .. } => {
                    Self::freeze_time();
                }
                _action_unfreeze_time @ TestCaseAction::UnfreezeTime { .. } => {
                    Self::unfreeze_time();
                }
                _ => {
                    panic!("unexpected action kind");
                }
            }
        }

        if let Some(test_case_epilogue) = test_case.epilogue.as_ref() {
            for topic in &test_case_epilogue.subscribed_topics {
                assert!(
                    mqtt_hub.has_subscribed(topic),
                    "topic {topic} has not been subscribed"
                );
            }

            if let Some(publication_count) = test_case_epilogue.publication_count {
                assert_eq!(
                    publication_count,
                    mqtt_hub.get_publication_count(),
                    "publication count"
                );
            }

            for published_message in &test_case_epilogue.published_messages {
                Self::check_published_message(published_message, &mqtt_hub, &correlation_ids);
            }

            if let Some(acknowledgement_count) = test_case_epilogue.acknowledgement_count {
                assert_eq!(
                    acknowledgement_count,
                    mqtt_hub.get_acknowledgement_count(),
                    "acknowledgement count"
                );
            }

            if let Some(execution_count) = test_case_epilogue.execution_count {
                assert_eq!(execution_count, *execution_counts[0].lock().unwrap());
            }

            for (executor_index, execution_count) in &test_case_epilogue.execution_counts {
                assert_eq!(
                    *execution_count,
                    *execution_counts[*executor_index].lock().unwrap()
                );
            }
        }
    }

    async fn executor_loop(
        mut executor: rpc_command::Executor<TestPayload, TestPayload>,
        test_case_executor: TestCaseExecutor<ExecutorDefaults>,
        countdown_events: CountdownEventMap,
        execution_count: Arc<Mutex<i32>>,
    ) {
        let mut sequencers = HashMap::new();
        for req in test_case_executor.request_responses_map.keys() {
            sequencers.insert(req.clone(), 0);
        }
        let arc_sequencers = Arc::new(Mutex::new(sequencers));

        loop {
            if let Some(Ok(request)) = executor.recv().await {
                *execution_count.lock().unwrap() += 1;

                if test_case_executor.raise_error {
                    drop(request);
                    continue;
                }
                tokio::task::spawn({
                    let test_case_executor = test_case_executor.clone();
                    let countdown_events = countdown_events.clone();
                    let arc_sequencers = arc_sequencers.clone();
                    async move {
                        for test_case_sync in &test_case_executor.sync {
                            if let Some(signal_event) = &test_case_sync.signal_event {
                                countdown_events.signal(signal_event.as_str());
                            }

                            if let Some(wait_event) = &test_case_sync.wait_event {
                                log::info!("waiting for event {}", wait_event.as_str());
                                countdown_events
                                    .wait_timeout(wait_event.as_str(), TEST_TIMEOUT)
                                    .await
                                    .expect("test timeout in countdown_event.wait");
                                log::info!("done waiting for event {}", wait_event.as_str());
                            }
                        }

                        let response_value = match request.payload.payload.as_ref() {
                            Some(request_value) => {
                                if let Some(response_sequence) =
                                    test_case_executor.request_responses_map.get(request_value)
                                {
                                    let mut guard = arc_sequencers.lock().unwrap();
                                    let sequencer = guard.get_mut(request_value).unwrap();
                                    let index = *sequencer % response_sequence.len();
                                    *sequencer += 1;
                                    Some(response_sequence[index].clone())
                                } else {
                                    None
                                }
                            }
                            _ => None,
                        };

                        let response_payload = TestPayload {
                            payload: response_value,
                            out_content_type: test_case_executor
                                .serializer
                                .out_content_type
                                .clone(),
                            accept_content_types: test_case_executor
                                .serializer
                                .accept_content_types
                                .clone(),
                            indicate_character_data: test_case_executor
                                .serializer
                                .indicate_character_data,
                            allow_character_data: test_case_executor
                                .serializer
                                .allow_character_data,
                            fail_deserialization: test_case_executor
                                .serializer
                                .fail_deserialization,
                        };

                        let mut metadata =
                            Vec::with_capacity(test_case_executor.response_metadata.len());
                        for (key, value) in &test_case_executor.response_metadata {
                            if let Some(val) = value {
                                metadata.push((key.clone(), val.clone()));
                            } else if let Some(kvp) =
                                request.custom_user_data.iter().find(|&m| m.0 == *key)
                            {
                                metadata.push((key.clone(), kvp.1.clone()));
                            }
                        }

                        if let Some(token_prefix) = &test_case_executor.token_metadata_prefix {
                            for (key, value) in &request.topic_tokens {
                                metadata.push((format!("{token_prefix}{key}"), value.clone()));
                            }
                        }

                        let response = rpc_command::executor::ResponseBuilder::default()
                            .payload(response_payload)
                            .unwrap()
                            .custom_user_data(metadata)
                            .build()
                            .unwrap();

                        _ = request.complete(response).await;
                    }
                });
            }
        }
    }

    async fn get_command_executor(
        managed_client: SessionManagedClient,
        tce: &TestCaseExecutor<ExecutorDefaults>,
        catch: Option<&TestCaseCatch>,
        mqtt_hub: &mut MqttHub,
    ) -> Option<rpc_command::Executor<TestPayload, TestPayload>> {
        let mut executor_options_builder = rpc_command::executor::OptionsBuilder::default();

        if let Some(request_topic) = tce.request_topic.as_ref() {
            executor_options_builder.request_topic_pattern(request_topic);
        }

        if let Some(topic_namespace) = tce.topic_namespace.as_ref() {
            executor_options_builder.topic_namespace(topic_namespace);
        }

        if let Some(topic_token_map) = tce.topic_token_map.as_ref() {
            executor_options_builder.topic_token_map(topic_token_map.clone());
        }

        if let Some(command_name) = tce.command_name.as_ref() {
            executor_options_builder.command_name(command_name);
        }

        executor_options_builder.is_idempotent(tce.idempotent);

        let options_result = executor_options_builder.build();
        if let Err(error) = options_result {
            if let Some(catch) = catch {
                aio_protocol_error_checker::check_error(
                    catch,
                    &Self::from_executor_options_builder_error(error),
                );
            } else {
                panic!("Unexpected error when building CommandExecutor options: {error}");
            }

            return None;
        }

        let executor_options = options_result.unwrap();

        match rpc_command::Executor::new(
            ApplicationContextBuilder::default().build().unwrap(),
            managed_client,
            executor_options,
        ) {
            Ok(mut executor) => {
                if let Some(catch) = catch {
                    // CommandExecutor has no start method, so if an exception is expected, recv may be needed to trigger it.
                    let (recv_result, _) = tokio::join!(
                        time::timeout(TEST_TIMEOUT, executor.recv()),
                        time::timeout(TEST_TIMEOUT, mqtt_hub.await_operation())
                    );
                    match recv_result {
                        Ok(Some(Ok(_))) => {
                            panic!(
                                "Expected {} error when constructing CommandExecutor but no error returned",
                                catch.error_kind
                            );
                        }
                        Ok(Some(Err(error))) => {
                            aio_protocol_error_checker::check_error(catch, &error);
                        }
                        _ => {
                            panic!(
                                "Expected {} error when calling recv() on CommandExecutor but got timeout instead",
                                catch.error_kind
                            );
                        }
                    }
                    None
                } else {
                    Some(executor)
                }
            }
            Err(error) => {
                if let Some(catch) = catch {
                    aio_protocol_error_checker::check_error(catch, &error);
                    None
                } else {
                    panic!("Unexpected error when constructing CommandExecutor: {error}");
                }
            }
        }
    }

    fn receive_request(
        action: &TestCaseAction<ExecutorDefaults>,
        mqtt_hub: &mut MqttHub,
        source_ids: &mut HashMap<i32, Uuid>,
        correlation_ids: &mut HashMap<i32, Option<azure_mqtt::mqtt_proto::BinaryData<Bytes>>>,
        packet_ids: &mut HashMap<i32, u16>,
        tcs: &TestCaseSerializer<ExecutorDefaults>,
    ) {
        if let TestCaseAction::ReceiveRequest {
            defaults_type: _,
            topic,
            payload,
            content_type,
            format_indicator,
            metadata,
            correlation_index,
            correlation_id,
            qos,
            message_expiry,
            response_topic,
            source_index,
            packet_index,
        } = action
        {
            let mut user_properties: Vec<(
                azure_mqtt::packet::ByteStr<Bytes>,
                azure_mqtt::packet::ByteStr<Bytes>,
            )> = metadata
                .iter()
                .map(|(k, v)| (k.as_str().into(), v.as_str().into()))
                .collect();

            if let Some(source_index) = source_index {
                if let Some(source_id) = source_ids.get(source_index) {
                    user_properties.push((
                        "__srcId".into(),
                        source_id.hyphenated().to_string().as_str().into(),
                    ));
                } else {
                    let source_id = Uuid::new_v4();
                    user_properties.push((
                        "__srcId".into(),
                        source_id.hyphenated().to_string().as_str().into(),
                    ));
                    source_ids.insert(*source_index, source_id);
                }
            }

            let correlation_data = if let Some(correlation_index) = correlation_index {
                if let Some(correlation_data) = correlation_ids.get(correlation_index) {
                    correlation_data.clone()
                } else {
                    let correlation_data = if let Some(correlation_id) = correlation_id {
                        Some(correlation_id.as_bytes().into())
                    } else {
                        Some(Uuid::new_v4().as_bytes().into())
                    };
                    correlation_ids.insert(*correlation_index, correlation_data.clone());
                    correlation_data
                }
            } else {
                None
            }
            .map(|cd| cd.as_ref().into());

            let message_expiry_interval = message_expiry.as_ref().map(|message_expiry| {
                u32::try_from(message_expiry.to_duration().as_secs()).unwrap()
            });

            let packet_id = if let Some(packet_index) = packet_index {
                packet_ids.get(packet_index)
            } else {
                None
            };
            let packet_id = if let Some(packet_id) = packet_id {
                *packet_id
            } else {
                mqtt_hub.get_new_packet_id()
            };
            if let Some(packet_index) = packet_index {
                packet_ids.insert(*packet_index, packet_id);
            }

            let topic = if let Some(topic) = topic {
                topic.clone()
            } else {
                String::new()
            };

            let payload = serde_json::to_vec(&TestPayload {
                payload: payload.clone(),
                out_content_type: tcs.out_content_type.clone(),
                accept_content_types: tcs.accept_content_types.clone(),
                indicate_character_data: tcs.indicate_character_data,
                allow_character_data: tcs.allow_character_data,
                fail_deserialization: tcs.fail_deserialization,
            })
            .unwrap();

            let properties = azure_mqtt::mqtt_proto::PublishOtherProperties {
                payload_is_utf8: to_is_utf8(format_indicator.as_ref()),
                message_expiry_interval,
                response_topic: response_topic
                    .clone()
                    .map(|rt| azure_mqtt::mqtt_proto::Topic::new(rt).unwrap().into()),
                correlation_data,
                user_properties,
                content_type: content_type.as_ref().map(|ct| ct.as_str().into()),
                ..Default::default()
            };

            let publish = azure_mqtt::mqtt_proto::Publish {
                packet_identifier_dup_qos: new_packet_identifier_dup_qos(
                    qos::to_enum(*qos),
                    false,
                    packet_id,
                ),
                topic_name: azure_mqtt::mqtt_proto::Topic::new(topic).unwrap().into(),
                payload: payload.into(),
                other_properties: properties,
                retain: false,
            };

            mqtt_hub.receive_message(publish);
        } else {
            panic!("internal logic error");
        }
    }

    async fn await_acknowledgement(
        action: &TestCaseAction<ExecutorDefaults>,
        mqtt_hub: &mut MqttHub,
        packet_ids: &HashMap<i32, u16>,
    ) {
        if let TestCaseAction::AwaitAck {
            defaults_type: _,
            packet_index,
        } = action
        {
            let packet_id = future::timeout(TEST_TIMEOUT, mqtt_hub.await_acknowledgement())
                .await
                .expect("test timeout in await_acknowledgement");
            if let Some(packet_index) = packet_index {
                assert_eq!(
                    *packet_ids
                        .get(packet_index)
                        .expect("packet index {packet_index} not found in packet id map"),
                    packet_id,
                    "packet ID"
                );
            }
        } else {
            panic!("internal logic error");
        }
    }

    async fn await_publish(
        action: &TestCaseAction<ExecutorDefaults>,
        mqtt_hub: &mut MqttHub,
        correlation_ids: &HashMap<i32, Option<azure_mqtt::mqtt_proto::BinaryData<Bytes>>>,
    ) {
        if let TestCaseAction::AwaitPublish {
            defaults_type: _,
            correlation_index,
        } = action
        {
            let correlation_id = future::timeout(TEST_TIMEOUT, mqtt_hub.await_publish())
                .await
                .expect("test timeout in await_publish");
            if let Some(correlation_index) = correlation_index {
                assert_eq!(
                    correlation_ids.get(correlation_index).expect(
                        "correlation index {correlation_index} not found in correlation map"
                    ),
                    &correlation_id,
                    "correlation ID"
                );
            }
        } else {
            panic!("internal logic error");
        }
    }

    async fn sync(action: &TestCaseAction<ExecutorDefaults>, countdown_events: &CountdownEventMap) {
        if let TestCaseAction::Sync {
            defaults_type: _,
            signal_event,
            wait_event,
        } = action
        {
            if let Some(wait_event) = wait_event {
                countdown_events
                    .wait_timeout(wait_event.as_str(), TEST_TIMEOUT)
                    .await
                    .expect("test timeout in countdown_event.wait");
            }

            if let Some(signal_event) = signal_event {
                countdown_events.signal(signal_event.as_str());
            }
        } else {
            panic!("internal logic error");
        }
    }

    async fn sleep(action: &TestCaseAction<ExecutorDefaults>) {
        if let TestCaseAction::Sleep {
            defaults_type: _,
            duration,
        } = action
        {
            time::sleep(duration.to_duration()).await;
        } else {
            panic!("internal logic error");
        }
    }

    fn disconnect(mqtt_hub: &mut MqttHub) {
        mqtt_hub.disconnect();
    }

    fn freeze_time() {}

    fn unfreeze_time() {}

    fn check_published_message(
        expected_message: &TestCasePublishedMessage,
        mqtt_hub: &MqttHub,
        correlation_ids: &HashMap<i32, Option<azure_mqtt::mqtt_proto::BinaryData<Bytes>>>,
    ) {
        let published_message =
            if let Some(correlation_index) = expected_message.correlation_index {
                if let Some(correlation_id) = correlation_ids.get(&correlation_index) {
                    let publish = mqtt_hub.get_published_message(correlation_id);
                    if publish.is_some() {
                        publish
                    } else {
                        panic!("no message published with correlation data corresponding to index {correlation_index}");
                    }
                } else {
                    panic!("no correlation data recorded for correlation index {correlation_index}");
                }
            } else {
                let publish = mqtt_hub.get_published_message(&None);
                if publish.is_some() {
                    publish
                } else {
                    panic!("no message published with empty correlation data");
                }
            }.unwrap();

        if let Some(topic) = expected_message.topic.as_ref() {
            assert_eq!(topic, published_message.topic_name.as_str(), "topic");
        }

        if let Some(payload) = expected_message.payload.as_ref() {
            if let Some(payload) = payload {
                assert_eq!(published_message.payload, *payload.as_bytes(), "payload");
            } else {
                assert!(&published_message.payload.is_empty());
            }
        }

        if expected_message.content_type.is_some() {
            assert_eq!(
                published_message
                    .other_properties
                    .content_type
                    .as_ref()
                    .map(std::convert::AsRef::as_ref),
                expected_message.content_type.as_deref()
            );
        }

        if expected_message.format_indicator.is_some() {
            assert_eq!(
                to_is_utf8(expected_message.format_indicator.as_ref()),
                published_message.other_properties.payload_is_utf8
            );
        }

        if !expected_message.metadata.is_empty() {
            for (key, value) in &expected_message.metadata {
                let found = published_message
                    .other_properties
                    .user_properties
                    .iter()
                    .find(|&k| k.0 == **key);
                if let Some(value) = value {
                    assert_eq!(
                        found.unwrap().1,
                        **value,
                        "metadata key {key} expected {value}"
                    );
                } else {
                    assert_eq!(None, found, "metadata key {key} not expected");
                }
            }
        }

        if let Some(command_status) = expected_message.command_status {
            let found = published_message
                .other_properties
                .user_properties
                .iter()
                .find(|&k| k.0 == "__stat");
            if let Some(command_status) = command_status {
                assert_eq!(
                    found.unwrap().1,
                    *command_status.to_string(),
                    "status property expected {command_status}"
                );
            } else {
                assert_eq!(None, found, "status property not expected");
            }
        }

        if let Some(is_application_error) = expected_message.is_application_error {
            let found: Option<_> = published_message
                .other_properties
                .user_properties
                .iter()
                .find(|&k| &k.0 == "__apErr");
            if is_application_error {
                assert!(
                    found.unwrap().1.as_ref().to_lowercase() == "true",
                    "is application error"
                );
            } else {
                assert!(
                    found.is_none() || found.unwrap().1.as_ref().to_lowercase() == "false",
                    "is application error"
                );
            }
        }

        if expected_message.expiry.is_some() {
            assert_eq!(
                expected_message.expiry,
                published_message.other_properties.message_expiry_interval
            );
        }
    }

    fn from_executor_options_builder_error(
        builder_error: rpc_command::executor::OptionsBuilderError,
    ) -> AIOProtocolError {
        let property_name = match builder_error {
            rpc_command::executor::OptionsBuilderError::UninitializedField(field_name) => {
                Some(field_name.to_string())
            }
            _ => None,
        };

        let mut protocol_error = AIOProtocolError {
            message: None,
            kind: AIOProtocolErrorKind::ConfigurationInvalid,
            is_shallow: true,
            is_remote: false,
            nested_error: Some(Box::new(builder_error)),
            header_name: None,
            header_value: None,
            timeout_name: None,
            timeout_value: None,
            property_name,
            property_value: None,
            command_name: None,
            protocol_version: None,
            supported_protocol_major_versions: None,
        };

        protocol_error.ensure_error_message();
        protocol_error
    }
}
