// 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::telemetry;
use bytes::Bytes;
use chrono::SecondsFormat;
use serde_json;
use tokio::sync::mpsc;
use tokio::time;
use uuid::Uuid;

use crate::metl::aio_protocol_error_checker;
use crate::metl::defaults::ReceiverDefaults;
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_cloud_event::TestCaseCloudEvent;
use crate::metl::test_case_received_telemetry::TestCaseReceivedTelemetry;
use crate::metl::test_case_receiver::TestCaseReceiver;
use crate::metl::test_case_serializer::TestCaseSerializer;
use crate::metl::test_payload::TestPayload;

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

struct ReceivedTelemetry {
    telemetry_value: Option<String>,
    metadata: HashMap<String, String>,
    topic_tokens: HashMap<String, String>,
    cloud_event: Option<TestCaseCloudEvent>,
    source_id: Option<String>,
}

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

impl TelemetryReceiverTester {
    pub async fn test_telemetry_receiver(
        test_case: TestCase<ReceiverDefaults>,
        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 receiver_count = test_case.prologue.receivers.len();
        let mut telemetry_counts = Vec::with_capacity(receiver_count);
        let (telemetry_tx, mut telemetry_rx) = mpsc::unbounded_channel();

        let mut ix = 0;
        for test_case_receiver in &test_case.prologue.receivers {
            ix += 1;
            let catch = if ix == receiver_count {
                test_case.prologue.catch.as_ref()
            } else {
                None
            };

            let telemetry_count = Arc::new(Mutex::new(0));
            telemetry_counts.push(telemetry_count.clone());

            if let Some(receiver) = Self::get_telemetry_receiver(
                managed_client.clone(),
                test_case_receiver,
                catch,
                &mut mqtt_hub,
            )
            .await
                && !test_case.actions.is_empty()
            {
                tokio::task::spawn(Self::receiver_loop(
                    receiver,
                    telemetry_count,
                    telemetry_tx.clone(),
                ));
            }
        }

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

        let mut source_ids: HashMap<i32, Uuid> = 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_telemetry @ TestCaseAction::ReceiveTelemetry { .. } => {
                    Self::receive_telemetry(
                        action_receive_telemetry,
                        &mut mqtt_hub,
                        &mut source_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_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(acknowledgement_count) = test_case_epilogue.acknowledgement_count {
                assert_eq!(
                    acknowledgement_count,
                    mqtt_hub.get_acknowledgement_count(),
                    "acknowledgement count"
                );
            }

            if let Some(telemetry_count) = test_case_epilogue.telemetry_count {
                assert_eq!(
                    telemetry_count,
                    *telemetry_counts[0].lock().unwrap(),
                    "telemetry count"
                );
            }

            for (receiver_index, telemetry_count) in &test_case_epilogue.telemetry_counts {
                assert_eq!(
                    *telemetry_count,
                    *telemetry_counts[*receiver_index].lock().unwrap(),
                    "telemetry count"
                );
            }

            for received_telemetry in &test_case_epilogue.received_telemetries {
                Self::check_received_telemetry(received_telemetry, &mut telemetry_rx, &source_ids)
                    .await;
            }
        }
    }

    async fn receiver_loop(
        mut receiver: telemetry::Receiver<TestPayload>,
        telemetry_count: Arc<Mutex<i32>>,
        telemetry_tx: mpsc::UnboundedSender<ReceivedTelemetry>,
    ) {
        while let Some(message) = receiver.recv().await {
            match message {
                Ok((telemetry, ack_token)) => {
                    *telemetry_count.lock().unwrap() += 1;

                    let cloud_event =
                        match telemetry::receiver::cloud_event_from_telemetry(&telemetry) {
                            Ok(cloud_event) => Some(TestCaseCloudEvent {
                                source: Some(cloud_event.source),
                                event_type: Some(cloud_event.event_type),
                                spec_version: Some(cloud_event.spec_version),
                                id: Some(cloud_event.id),
                                time: Some(
                                    cloud_event
                                        .time
                                        .map(|t| t.to_rfc3339_opts(SecondsFormat::Secs, true)),
                                ),
                                data_content_type: cloud_event.data_content_type,
                                subject: Some(cloud_event.subject),
                                data_schema: Some(cloud_event.data_schema),
                            }),
                            Err(_) => None,
                        };

                    let mut metadata = HashMap::new();
                    for (key, value) in telemetry.custom_user_data {
                        metadata.insert(key, value);
                    }

                    telemetry_tx
                        .send(ReceivedTelemetry {
                            telemetry_value: telemetry.payload.payload,
                            metadata,
                            topic_tokens: telemetry.topic_tokens.clone(),
                            cloud_event,
                            source_id: telemetry.sender_id,
                        })
                        .unwrap();

                    if let Some(ack_token) = ack_token {
                        ack_token.ack().await.unwrap();
                    }
                }
                Err(e) => {
                    panic!("Error receiving telemetry message: {e:?}");
                }
            }
        }
    }

    async fn get_telemetry_receiver(
        managed_client: SessionManagedClient,
        tcr: &TestCaseReceiver<ReceiverDefaults>,
        catch: Option<&TestCaseCatch>,
        mqtt_hub: &mut MqttHub,
    ) -> Option<telemetry::Receiver<TestPayload>> {
        let mut receiver_options_builder = telemetry::receiver::OptionsBuilder::default();

        if let Some(telemetry_topic) = tcr.telemetry_topic.as_ref() {
            receiver_options_builder.topic_pattern(telemetry_topic);
        }

        if let Some(topic_namespace) = tcr.topic_namespace.as_ref() {
            receiver_options_builder.topic_namespace(topic_namespace);
        }

        if let Some(topic_token_map) = tcr.topic_token_map.as_ref() {
            receiver_options_builder.topic_token_map(topic_token_map.clone());
        }

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

            return None;
        }

        let receiver_options = options_result.unwrap();

        match telemetry::Receiver::new(
            ApplicationContextBuilder::default().build().unwrap(),
            managed_client,
            receiver_options,
        ) {
            Ok(mut receiver) => {
                if let Some(catch) = catch {
                    // TelemetryReceiver 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, receiver.recv()),
                        time::timeout(TEST_TIMEOUT, mqtt_hub.await_operation())
                    );
                    match recv_result {
                        Ok(Some(Ok(_))) => {
                            panic!(
                                "Expected {} error when constructing TelemetryReceiver 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 TelemetryReceiver but got timeout instead",
                                catch.error_kind
                            );
                        }
                    }
                    None
                } else {
                    Some(receiver)
                }
            }
            Err(error) => {
                if let Some(catch) = catch {
                    aio_protocol_error_checker::check_error(catch, &error);
                    None
                } else {
                    panic!("Unexpected error when constructing TelemetryReceiver: {error}");
                }
            }
        }
    }

    fn receive_telemetry(
        action: &TestCaseAction<ReceiverDefaults>,
        mqtt_hub: &mut MqttHub,
        source_ids: &mut HashMap<i32, Uuid>,
        packet_ids: &mut HashMap<i32, u16>,
        tcs: &TestCaseSerializer<ReceiverDefaults>,
    ) {
        if let TestCaseAction::ReceiveTelemetry {
            defaults_type: _,
            topic,
            payload,
            content_type,
            format_indicator,
            metadata,
            qos,
            message_expiry,
            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 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,
                user_properties,
                content_type: content_type.as_ref().map(|ct| ct.as_str().into()),
                ..Default::default()
            };

            let publish = azure_mqtt::mqtt_proto::Publish::<Bytes> {
                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<ReceiverDefaults>,
        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 sleep(action: &TestCaseAction<ReceiverDefaults>) {
        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() {}

    async fn check_received_telemetry(
        expected_telemetry: &TestCaseReceivedTelemetry,
        telemetry_rx: &mut mpsc::UnboundedReceiver<ReceivedTelemetry>,
        source_ids: &HashMap<i32, Uuid>,
    ) {
        let received_telemetry = telemetry_rx
            .recv()
            .await
            .expect("missing expected telemetry");

        if let Some(expected_value) = expected_telemetry.telemetry_value.as_ref() {
            assert_eq!(expected_value, &received_telemetry.telemetry_value);
        }

        for (key, value) in &expected_telemetry.metadata {
            assert_eq!(
                value.as_ref(),
                received_telemetry.metadata.get(key),
                "metadata key {key} expected {value:?}"
            );
        }

        for (key, value) in &expected_telemetry.topic_tokens {
            assert_eq!(
                &value,
                &received_telemetry
                    .topic_tokens
                    .get(key)
                    .expect("expected telemetry topic token {key}"),
                "topic token {key} expected replacement {value:?}"
            );
        }

        if let Some(expected_cloud_event) = &expected_telemetry.cloud_event {
            match expected_cloud_event {
                Some(expected_cloud_event) => {
                    if let Some(received_cloud_event) = &received_telemetry.cloud_event {
                        if let Some(expected_source) = expected_cloud_event.source.as_ref() {
                            assert_eq!(
                                expected_source,
                                received_cloud_event
                                    .source
                                    .as_ref()
                                    .expect("missing cloud event source")
                            );
                        }
                        if let Some(expected_event_type) = expected_cloud_event.event_type.as_ref()
                        {
                            assert_eq!(
                                expected_event_type,
                                received_cloud_event
                                    .event_type
                                    .as_ref()
                                    .expect("missing cloud event type")
                            );
                        }
                        if let Some(expected_spec_version) =
                            expected_cloud_event.spec_version.as_ref()
                        {
                            assert_eq!(
                                expected_spec_version,
                                received_cloud_event
                                    .spec_version
                                    .as_ref()
                                    .expect("missing cloud event spec version")
                            );
                        }
                        if let Some(expected_id) = expected_cloud_event.id.as_ref() {
                            assert_eq!(
                                expected_id,
                                received_cloud_event
                                    .id
                                    .as_ref()
                                    .expect("missing cloud event type")
                            );
                        }
                        if let Some(expected_time) = expected_cloud_event.time.as_ref() {
                            if let Some(expected_time) = expected_time {
                                assert_eq!(
                                    expected_time,
                                    received_cloud_event
                                        .time
                                        .as_ref()
                                        .unwrap()
                                        .as_ref()
                                        .expect("missing cloud event time")
                                );
                            } else {
                                assert!(received_cloud_event.time.as_ref().unwrap().is_none());
                            }
                        }
                        if let Some(expected_data_content_type) =
                            expected_cloud_event.data_content_type.as_ref()
                        {
                            assert_eq!(
                                expected_data_content_type,
                                received_cloud_event
                                    .data_content_type
                                    .as_ref()
                                    .expect("missing cloud event data content type")
                            );
                        }
                        if let Some(expected_subject) = expected_cloud_event.subject.as_ref() {
                            if let Some(expected_subject) = expected_subject {
                                assert_eq!(
                                    expected_subject,
                                    received_cloud_event
                                        .subject
                                        .as_ref()
                                        .unwrap()
                                        .as_ref()
                                        .expect("missing cloud event subject")
                                );
                            } else {
                                assert!(received_cloud_event.subject.as_ref().unwrap().is_none());
                            }
                        }
                        if let Some(expected_data_schema) =
                            expected_cloud_event.data_schema.as_ref()
                        {
                            if let Some(expected_data_schema) = expected_data_schema {
                                assert_eq!(
                                    expected_data_schema,
                                    received_cloud_event
                                        .data_schema
                                        .as_ref()
                                        .unwrap()
                                        .as_ref()
                                        .expect("missing cloud event data schema")
                                );
                            } else {
                                assert!(
                                    received_cloud_event.data_schema.as_ref().unwrap().is_none()
                                );
                            }
                        }
                    } else {
                        panic!("expected cloud event but not found in received telemetry");
                    }
                }
                None => {
                    assert!(
                        received_telemetry.cloud_event.is_none(),
                        "unexpected cloud event found in received telemetry"
                    );
                }
            }
        }

        if let Some(source_index) = &expected_telemetry.source_index {
            let source_id = source_ids
                .get(source_index)
                .expect("source index {source_index} not found in source id map");
            assert_eq!(
                source_id.hyphenated().to_string(),
                received_telemetry.source_id.expect("missing source id")
            );
        }
    }

    fn from_receiver_options_builder_error(
        builder_error: telemetry::receiver::OptionsBuilderError,
    ) -> AIOProtocolError {
        let property_name = match builder_error {
            telemetry::receiver::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
    }
}
