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

//! Types for Azure IoT Operations Connectors.

use std::{borrow::Cow, collections::HashMap, hash::Hash, path::PathBuf, sync::Arc};

use azure_iot_operations_services::{
    azure_device_registry::{
        self, HealthStatus, RuntimeHealth,
        health_reporter::HealthReporterSender,
        models::{self as adr_models, Asset, DeviceRef},
    },
    schema_registry,
};
use chrono::{DateTime, Utc};
use thiserror::Error;
use tokio::sync::watch;
use tokio::sync::{
    Notify,
    mpsc::{self, UnboundedReceiver, UnboundedSender},
};
use tokio_retry2::{Retry, RetryError};
use tokio_util::sync::CancellationToken;

use crate::{
    AdrConfigError, Data, DataOperationKind, DataOperationName, DataOperationRef,
    ManagementActionRef, MessageSchema, MessageSchemaReference,
    base_connector::ConnectorContext,
    deployment_artifacts::{
        self,
        azure_device_registry::{AssetRef, DeviceEndpointRef},
    },
    destination_endpoint::{self, DataOperationForwarder},
    management_action_executor::{self, ManagementActionExecutor},
};

/// Used as the strategy when using [`tokio_retry2::Retry`]
const RETRY_STRATEGY: tokio_retry2::strategy::ExponentialFactorBackoff =
    tokio_retry2::strategy::ExponentialFactorBackoff::from_millis(500, 2.0);

/// Represents the runtime health of a resource.
#[derive(Debug, Clone)]
pub enum RuntimeHealthEvent {
    /// Resource is currently available.
    Available,
    /// Resource is currently unavailable.
    Unavailable {
        /// A human-readable message describing the last transition.
        message: Option<String>,
        /// Unique, CamelCase reason code describing the cause of the last health state transition.
        reason_code: Option<String>,
    },
}

/// Notifications that can be received for a Client
pub enum ClientNotification<T> {
    /// Indicates that the Client's specification has been updated in place
    Updated,
    /// Indicates that the Client has been deleted.
    Deleted,
    /// Indicates that there is a new `T` for this Client, which is included in the notification
    Created(T),
}

/// Represents the result of a network modification
pub enum ModifyResult {
    /// Indicates that the modification was reported
    Reported,
    /// Indicates that the modification did not occur
    NotModified,
}

/// A cloneable status reporter for Device and Endpoint status reporting.
///
/// This provides a way to report Device and Endpoint status changes from outside the [`DeviceEndpointClient`].
/// Each clone maintains its own independent version snapshot, allowing different tasks to work with
/// different versions of the device specification.
#[derive(Clone, Debug)]
pub struct DeviceEndpointStatusReporter {
    connector_context: Arc<ConnectorContext>,
    device_endpoint_status: Arc<tokio::sync::RwLock<DeviceEndpointStatus>>,
    device_endpoint_specification: Arc<std::sync::RwLock<DeviceSpecification>>,
    health_sender: HealthReporterSender,
    device_endpoint_ref: DeviceEndpointRef,
    snapshotted_version: u64,
}

impl DeviceEndpointStatusReporter {
    /// This function is used to report the current runtime health event for the Endpoint.
    ///
    /// NOTE: This will only report the health event to the service if needed, so it is
    /// okay to call this function frequently. It is required to call this function any time
    /// the health status changes to ensure that the service has the latest information.
    /// If the health status does not change, the underlying client will report the last known
    /// health event at the base connector `health_report_interval`'s frequency unless `pause_health_reporting` is called.
    ///
    /// The version used for reporting is the snapshotted version from the last call to
    /// `refresh_health_version()` or `pause_and_refresh_health_version()`, not necessarily the current specification version.
    pub fn report_health_event(&self, health_event: RuntimeHealthEvent) {
        let (status, message, reason_code) = match health_event {
            RuntimeHealthEvent::Available => (HealthStatus::Available, None, None),
            RuntimeHealthEvent::Unavailable {
                message,
                reason_code,
            } => (HealthStatus::Unavailable, message, reason_code),
        };
        let runtime_health = RuntimeHealth {
            last_update_time: Utc::now(),
            message,
            reason_code,
            status,
            version: self.snapshotted_version,
        };
        if let Err(e) = self.health_sender.report(runtime_health) {
            log::warn!(
                "Failed to send health event for device endpoint {:?}: {e}",
                self.device_endpoint_ref
            );
        }
    }

    /// Pauses background health event reporting until a new event is reported.
    /// This should be called when the component is updated to indicate that the previous health event may no longer be applicable.
    pub fn pause_health_reporting(&self) {
        if let Err(e) = self.health_sender.pause() {
            log::warn!(
                "Failed to pause health reporting for device endpoint {:?}: {e}",
                self.device_endpoint_ref
            );
        }
    }

    /// Snapshots the current specification version for use in future health events.
    /// Call this when starting to use a new definition to "lock in" the version you're working with.
    ///
    /// # Panics
    /// if the specification mutex has been poisoned, which should not be possible
    pub fn refresh_health_version(&mut self) {
        self.snapshotted_version = self
            .device_endpoint_specification
            .read()
            .unwrap()
            .version
            .unwrap_or(0);
    }

    /// Combines `pause_health_reporting()` and `refresh_health_version()`.
    /// Pauses background reporting AND snapshots the current specification version.
    /// Useful when receiving an update in the same task that handles notifications.
    pub fn pause_and_refresh_health_version(&mut self) {
        self.pause_health_reporting();
        self.refresh_health_version();
    }

    /// Used to conditionally report the device status and then updates the device with the new status returned.
    ///
    /// The `modify` function is called with the current device status (if any) and should return:
    /// - `Some(new_status)` if the status should be updated and reported
    /// - `None` if no update is needed
    ///
    /// # Returns
    /// - [`ModifyResult::Reported`] if the status was updated and successfully reported
    /// - [`ModifyResult::NotModified`] if no modification was needed or the version changed during processing
    ///
    /// # Errors
    /// [`azure_device_registry::Error`] of kind [`AIOProtocolError`](azure_device_registry::ErrorKind::AIOProtocolError) if
    /// there are any underlying errors from the AIO RPC protocol. This error will be retried
    /// 10 times with exponential backoff and jitter and only returned if it still is failing.
    ///
    /// [`azure_device_registry::Error`] of kind [`ServiceError`](azure_device_registry::ErrorKind::ServiceError) if an error is returned
    /// by the Azure Device Registry service.
    ///
    /// # Panics
    /// if the specification mutex has been poisoned, which should not be possible
    pub async fn report_device_status_if_modified<F>(
        &self,
        modify: F,
    ) -> Result<ModifyResult, azure_device_registry::Error>
    where
        F: Fn(Option<Result<(), &AdrConfigError>>) -> Option<Result<(), AdrConfigError>>,
    {
        // Get the current version of the device endpoint specification
        let cached_version = self.device_endpoint_specification.read().unwrap().version;

        {
            // Get the current device status
            let status_read_guard = self.device_endpoint_status.read().await;
            let current_device_endpoint_status =
                status_read_guard.get_current_device_endpoint_status(cached_version);

            let modify_input = Self::get_device_modify_input(&current_device_endpoint_status);

            // We do not use the result since the status once we acquire the read lock might change
            // and we will have to re-evaluate
            let Some(_modify_result) = modify(modify_input) else {
                // If no modification is needed, return Ok
                return Ok(ModifyResult::NotModified);
            };
        }

        // To modify, we need to acquire the write lock
        let mut status_write_guard = self.device_endpoint_status.write().await;

        if cached_version != self.device_endpoint_specification.read().unwrap().version {
            // Our modify is no longer valid
            log::debug!(
                "Device endpoint specification is out-of-date for {:?}; skipping modification",
                self.device_endpoint_ref
            );
            return Ok(ModifyResult::NotModified);
        }

        // We can continue here because the device endpoint status will not change. The specification might
        // update but we will be reporting for an old version.

        // Get the current device endpoint status in case it has changed
        let current_device_endpoint_status =
            status_write_guard.get_current_device_endpoint_status(cached_version);

        let modify_result = {
            let modify_input = Self::get_device_modify_input(&current_device_endpoint_status);

            let Some(modify_result) = modify(modify_input) else {
                // If no modification is needed, return Ok
                return Ok(ModifyResult::NotModified);
            };

            modify_result
        };

        let device_endpoint_status_to_report = current_device_endpoint_status.into_owned();

        let endpoints = {
            if let Some(inbound_endpoint_status) =
                device_endpoint_status_to_report.inbound_endpoint_status
            {
                // If the inbound endpoint status is present, include it in the report
                HashMap::from([(
                    self.device_endpoint_ref.inbound_endpoint_name.clone(),
                    inbound_endpoint_status.err(),
                )])
            } else {
                // If the inbound endpoint status is not present, we still have to include it in the report
                HashMap::from([(self.device_endpoint_ref.inbound_endpoint_name.clone(), None)])
            }
        };

        // Create a device status
        let device_status_to_report = adr_models::DeviceStatus {
            config: Some(azure_device_registry::ConfigStatus {
                version: cached_version,
                error: modify_result.err(),
                last_transition_time: Some(Utc::now()),
            }),
            endpoints,
        };

        log::debug!(
            "Reporting device status from app for {:?}",
            self.device_endpoint_ref
        );

        Self::internal_report_status(
            &self.connector_context,
            device_status_to_report,
            &mut status_write_guard,
            &self.device_endpoint_ref,
        )
        .await?;

        Ok(ModifyResult::Reported)
    }

    fn get_device_modify_input(
        current_status: &DeviceEndpointStatus,
    ) -> Option<Result<(), &AdrConfigError>> {
        current_status
            .config
            .as_ref()
            .map(|config| match &config.error {
                Some(e) => Err(e),
                None => Ok(()),
            })
    }

    /// Used to conditionally report the endpoint status and then updates the device with the new status returned.
    ///
    /// The `modify` function is called with the current endpoint status (if any) and should return:
    /// - `Some(new_status)` if the status should be updated and reported
    /// - `None` if no update is needed
    ///
    /// # Returns
    /// - [`ModifyResult::Reported`] if the status was updated and successfully reported
    /// - [`ModifyResult::NotModified`] if no modification was needed or the version changed during processing
    ///
    /// # Errors
    /// [`azure_device_registry::Error`] of kind [`AIOProtocolError`](azure_device_registry::ErrorKind::AIOProtocolError) if
    /// there are any underlying errors from the AIO RPC protocol. This error will be retried
    /// 10 times with exponential backoff and jitter and only returned if it still is failing.
    ///
    /// [`azure_device_registry::Error`] of kind [`ServiceError`](azure_device_registry::ErrorKind::ServiceError) if an error is returned
    /// by the Azure Device Registry service.
    ///
    /// # Panics
    /// if the specification mutex has been poisoned, which should not be possible
    pub async fn report_endpoint_status_if_modified<F>(
        &self,
        modify: F,
    ) -> Result<ModifyResult, azure_device_registry::Error>
    where
        F: Fn(Option<Result<(), &AdrConfigError>>) -> Option<Result<(), AdrConfigError>>,
    {
        // Get the current version of the device endpoint specification
        let cached_version = self.device_endpoint_specification.read().unwrap().version;

        {
            // Get the current endpoint status
            let status_read_guard = self.device_endpoint_status.read().await;
            let current_device_endpoint_status =
                status_read_guard.get_current_device_endpoint_status(cached_version);

            let modify_input = Self::get_endpoint_modify_input(&current_device_endpoint_status);

            // We do not use the result since the status once we acquire the read lock might change
            // and we will have to re-evaluate
            let Some(_modify_result) = modify(modify_input) else {
                // If no modification is needed, return Ok
                return Ok(ModifyResult::NotModified);
            };
        }

        // To modify, we need to acquire the write lock
        let mut status_write_guard = self.device_endpoint_status.write().await;

        if cached_version != self.device_endpoint_specification.read().unwrap().version {
            // Our modify is no longer valid
            log::debug!(
                "Device endpoint specification is out-of-date for {:?}; skipping modification",
                self.device_endpoint_ref
            );
            return Ok(ModifyResult::NotModified);
        }

        // We can continue here because the device endpoint status will not change. The specification might
        // update but we will be reporting for an old version.

        // Get the current device endpoint status in case it has changed
        let current_device_endpoint_status =
            status_write_guard.get_current_device_endpoint_status(cached_version);

        let modify_result = {
            let modify_input = Self::get_endpoint_modify_input(&current_device_endpoint_status);

            let Some(modify_result) = modify(modify_input) else {
                // If no modification is needed, return Ok
                return Ok(ModifyResult::NotModified);
            };

            modify_result
        };

        let device_endpoint_status_to_report = current_device_endpoint_status.into_owned();

        // If the config is None, we need to create a new one to report along with the endpoint status
        let mut device_config_status = device_endpoint_status_to_report.config.unwrap_or_default();

        device_config_status.version = cached_version;
        device_config_status.last_transition_time = Some(Utc::now());

        // Create a device status
        let device_status_to_report = adr_models::DeviceStatus {
            config: Some(device_config_status),
            endpoints: HashMap::from([(
                self.device_endpoint_ref.inbound_endpoint_name.clone(),
                modify_result.clone().err(),
            )]),
        };

        log::debug!(
            "Reporting endpoint status from app for {:?}",
            self.device_endpoint_ref
        );

        // Report and update status
        Self::internal_report_status(
            &self.connector_context,
            device_status_to_report,
            &mut status_write_guard,
            &self.device_endpoint_ref,
        )
        .await?;

        Ok(ModifyResult::Reported)
    }

    fn get_endpoint_modify_input(
        current_status: &DeviceEndpointStatus,
    ) -> Option<Result<(), &AdrConfigError>> {
        current_status
            .inbound_endpoint_status
            .as_ref()
            .map(|status| match status {
                Ok(()) => Ok(()),
                Err(e) => Err(e),
            })
    }

    /// Reports an already built status to the service, with retries, and then updates the device with the new status returned
    async fn internal_report_status(
        connector_context: &Arc<ConnectorContext>,
        adr_device_status: adr_models::DeviceStatus,
        adr_device_status_ref: &mut DeviceEndpointStatus,
        device_endpoint_ref: &DeviceEndpointRef,
    ) -> Result<(), azure_device_registry::Error> {
        // send status update to the service
        let updated_device_status = Retry::spawn(
            RETRY_STRATEGY.map(tokio_retry2::strategy::jitter).take(10),
            async || -> Result<adr_models::DeviceStatus, RetryError<azure_device_registry::Error>> {
                connector_context
                    .azure_device_registry_client
                    .update_device_plus_endpoint_status(
                        device_endpoint_ref.device_name.clone(),
                        device_endpoint_ref.inbound_endpoint_name.clone(),
                        adr_device_status.clone(),
                        connector_context.azure_device_registry_timeout,
                    )
                    .await
                    .map_err(|e| adr_error_into_retry_error(e, "Update Device Status"))
            },
        )
        .await?;

        // update self with new returned status
        *adr_device_status_ref = DeviceEndpointStatus::new(
            updated_device_status,
            &device_endpoint_ref.inbound_endpoint_name,
        );
        Ok(())
    }
}

/// An Observation for device endpoint creation events that uses
/// multiple underlying clients to get full device endpoint information.
pub struct DeviceEndpointClientCreationObservation {
    connector_context: Arc<ConnectorContext>,
    device_endpoint_create_observation:
        deployment_artifacts::azure_device_registry::DeviceEndpointCreateObservation,
    /// Flag to track if device creation is in progress
    pending_device_creation: bool,
    /// Channels for sending and receiving completed device endpoint clients
    /// This is used to ensure that we only process one device creation at a time
    device_completion_rx: mpsc::Receiver<Option<DeviceEndpointClient>>,
    device_completion_tx: mpsc::Sender<Option<DeviceEndpointClient>>,
}
impl DeviceEndpointClientCreationObservation {
    /// Creates a new [`DeviceEndpointClientCreationObservation`] that uses the given [`ConnectorContext`]
    pub(crate) fn new(connector_context: Arc<ConnectorContext>) -> Self {
        // TODO: handle unwrap in a better way
        let device_endpoint_create_observation =
            deployment_artifacts::azure_device_registry::DeviceEndpointCreateObservation::new(
                connector_context.debounce_duration,
            )
            .unwrap();

        let (device_completion_tx, device_completion_rx) = mpsc::channel(1);

        Self {
            connector_context,
            device_endpoint_create_observation,
            pending_device_creation: false,
            device_completion_rx,
            device_completion_tx,
        }
    }

    /// Receives a notification for a newly created device endpoint. This
    /// notification includes the [`DeviceEndpointClient`], which can be used
    /// to receive Assets related to this Device Endpoint
    ///
    /// # Panics
    /// If the `device_endpoint_create_observation` channel is closed, which should not be possible
    pub async fn recv_notification(&mut self) -> DeviceEndpointClient {
        loop {
            tokio::select! {
                // Check for completed device creation
                Some(device_client_option) = self.device_completion_rx.recv() => {
                    self.pending_device_creation = false;
                    if let Some(device_client) = device_client_option {
                        return device_client;
                    }
                    // If device_client_option is None, creation failed, continue loop
                },
                // Get new device creation notifications only if not already processing one
                create_notification = self.device_endpoint_create_observation.recv_notification(), if !self.pending_device_creation => {
                    let (device_endpoint_ref, asset_create_observation) =
                        create_notification.expect("Device Endpoint Create Observation should never return None because the device_endpoint_create_observation struct holds the sending side of the channel");

                    // Start device creation task
                    self.pending_device_creation = true;
                    let connector_context = self.connector_context.clone();
                    let device_completion_tx = self.device_completion_tx.clone();

                    tokio::task::spawn(async move {
                        let device_client = Self::create_device_endpoint_client(
                            connector_context,
                            device_endpoint_ref,
                            asset_create_observation,
                        ).await;

                        // Always send the result (Some or None) to unblock the receiver
                        let _ = device_completion_tx.send(device_client).await;
                    });
                }
            }
        }
    }

    /// Internal helper to create a [`DeviceEndpointClient`]
    async fn create_device_endpoint_client(
        connector_context: Arc<ConnectorContext>,
        device_endpoint_ref: DeviceEndpointRef,
        asset_create_observation: deployment_artifacts::azure_device_registry::AssetCreateObservation,
    ) -> Option<DeviceEndpointClient> {
        // Obtain the device update observation
        let device_endpoint_update_observation =  match Retry::spawn(
            RETRY_STRATEGY.map(tokio_retry2::strategy::jitter),
            async || -> Result<azure_device_registry::DeviceUpdateObservation, RetryError<azure_device_registry::Error>> {
                connector_context
                    .azure_device_registry_client
                    .observe_device_update_notifications(
                        device_endpoint_ref.device_name.clone(),
                        device_endpoint_ref.inbound_endpoint_name.clone(),
                        connector_context.azure_device_registry_timeout,
                    )
                    // retry on network errors, otherwise don't retry on config/dev errors
                    .await
                    .map_err(|e| adr_error_into_retry_error(e, "Observe for Device Updates"))
        }).await {
            Ok(device_update_observation) => device_update_observation,
            Err(e) => {
              log::error!("Dropping device endpoint create notification: {device_endpoint_ref:?}. Failed to observe for device update notifications after retries: {e}");
              return None;
            },
        };

        // Get the device definition
        let device = match Retry::spawn(
            RETRY_STRATEGY.map(tokio_retry2::strategy::jitter),
            async || -> Result<adr_models::Device, RetryError<azure_device_registry::Error>> {
                connector_context
                    .azure_device_registry_client
                    .get_device(
                        device_endpoint_ref.device_name.clone(),
                        device_endpoint_ref.inbound_endpoint_name.clone(),
                        connector_context.azure_device_registry_timeout,
                    )
                    .await
                    .map_err(|e| adr_error_into_retry_error(e, "Get Device Definition"))
            },
        )
        .await
        {
            Ok(device) => device,
            Err(e) => {
                log::error!(
                    "Dropping device endpoint create notification: {device_endpoint_ref:?}. Failed to get Device definition after retries: {e}"
                );
                // unobserve as cleanup
                DeviceEndpointClient::unobserve_device(
                    &connector_context,
                    &device_endpoint_ref,
                    false,
                )
                .await;
                return None;
            }
        };

        // Get the device status
        let device_status = match Retry::spawn(
            RETRY_STRATEGY.map(tokio_retry2::strategy::jitter),
            async || -> Result<adr_models::DeviceStatus, RetryError<azure_device_registry::Error>> {
                connector_context
                    .azure_device_registry_client
                    .get_device_status(
                        device_endpoint_ref.device_name.clone(),
                        device_endpoint_ref.inbound_endpoint_name.clone(),
                        connector_context.azure_device_registry_timeout,
                    )
                    .await
                    .map_err(|e| adr_error_into_retry_error(e, "Get Device Status"))
            },
        )
        .await
        {
            Ok(device_status) => device_status,
            Err(e) => {
                log::error!("Dropping device endpoint create notification: {device_endpoint_ref:?}. Failed to get Device Status after retries: {e}");
                // unobserve as cleanup
                DeviceEndpointClient::unobserve_device(
                    &connector_context,
                    &device_endpoint_ref,
                    false)
                .await;
                return None;
            }
        };

        // turn the device definition into a DeviceEndpointClient
        match DeviceEndpointClient::new(
            device,
            device_status,
            device_endpoint_ref.clone(),
            device_endpoint_update_observation,
            asset_create_observation,
            connector_context.clone(),
        ) {
            Ok(managed_device) => Some(managed_device),
            Err(e) => {
                // the device definition didn't include the inbound_endpoint, so it likely no longer exists
                // TODO: This won't be a possible failure point in the future once the service returns errors
                log::error!(
                    "Dropping device endpoint create notification: {device_endpoint_ref:?}. {e}"
                );
                // unobserve as cleanup
                DeviceEndpointClient::unobserve_device(
                    &connector_context,
                    &device_endpoint_ref,
                    false,
                )
                .await;
                None
            }
        }
    }
}

/// Azure Device Registry Device Endpoint that includes additional functionality to report status and receive updates
#[derive(Debug, Getters)]
pub struct DeviceEndpointClient {
    /// The names of the Device and Inbound Endpoint
    device_endpoint_ref: DeviceEndpointRef,
    /// The 'specification' Field.
    #[getter(skip)]
    specification: Arc<std::sync::RwLock<DeviceSpecification>>,
    /// The 'status' Field.
    #[getter(skip)]
    status: Arc<tokio::sync::RwLock<DeviceEndpointStatus>>,
    // Internally used fields
    /// The internal observation for updates
    #[getter(skip)]
    device_update_observation: azure_device_registry::DeviceUpdateObservation,
    #[getter(skip)]
    asset_create_observation: deployment_artifacts::azure_device_registry::AssetCreateObservation,
    /// Flag to track if asset creation is in progress
    #[getter(skip)]
    pending_asset_creation: bool,
    /// Health reporter sender for sending health events
    #[getter(skip)]
    health_sender: HealthReporterSender,
    /// Cancellation token for health reporting task - cancelled on deletion
    #[getter(skip)]
    health_cancellation_token: CancellationToken,
    /// Channels for sending and receiving completed asset clients.
    /// This is used to ensure that we only process one asset creation at a time
    #[getter(skip)]
    asset_completion_rx: mpsc::UnboundedReceiver<Option<AssetClient>>,
    #[getter(skip)]
    asset_completion_tx: mpsc::UnboundedSender<Option<AssetClient>>,
    #[getter(skip)]
    connector_context: Arc<ConnectorContext>,
}
impl DeviceEndpointClient {
    pub(crate) fn new(
        device: adr_models::Device,
        device_status: adr_models::DeviceStatus,
        device_endpoint_ref: DeviceEndpointRef,
        device_update_observation: azure_device_registry::DeviceUpdateObservation,
        asset_create_observation: deployment_artifacts::azure_device_registry::AssetCreateObservation,
        connector_context: Arc<ConnectorContext>,
        // TODO: This won't need to return an error once the service properly sends errors if the endpoint doesn't exist
    ) -> Result<Self, String> {
        let (asset_completion_tx, asset_completion_rx) = mpsc::unbounded_channel();

        let health_cancellation_token = CancellationToken::new();
        let device_ref = DeviceRef {
            device_name: device_endpoint_ref.device_name.clone(),
            endpoint_name: device_endpoint_ref.inbound_endpoint_name.clone(),
        };
        let health_sender = connector_context
            .azure_device_registry_client
            .new_device_endpoint_health_reporter(
                device_ref,
                connector_context.azure_device_registry_timeout,
                connector_context.health_report_interval,
                health_cancellation_token.clone(),
            );
        Ok(DeviceEndpointClient {
            specification: Arc::new(std::sync::RwLock::new(DeviceSpecification::new(
                device,
                connector_context
                    .connector_artifacts
                    .device_endpoint_credentials_mount
                    .as_ref(),
                &device_endpoint_ref.inbound_endpoint_name,
            )?)),
            status: Arc::new(tokio::sync::RwLock::new(DeviceEndpointStatus::new(
                device_status,
                &device_endpoint_ref.inbound_endpoint_name,
            ))),
            device_endpoint_ref,
            device_update_observation,
            asset_create_observation,
            pending_asset_creation: false,
            health_sender,
            health_cancellation_token,
            asset_completion_rx,
            asset_completion_tx,
            connector_context,
        })
    }

    /// Used to receive notifications related to the Device/Inbound Endpoint
    /// from the Azure Device Registry Service.
    ///
    /// Returns [`ClientNotification::Updated`] if the Device Endpoint
    /// Specification has been updated in place.
    ///
    /// Returns [`ClientNotification::Deleted`] if the Device Endpoint has been
    /// deleted. The [`DeviceEndpointClient`] should not be used after this
    /// point, and no more notifications will be received.
    ///
    /// Returns [`ClientNotification::Created`] with a new [`AssetClient`] if a new
    /// Asset has been created.
    ///
    /// # Panics
    /// If the Azure Device Registry Service provides a notification that isn't for this Device Endpoint. This should not be possible.
    ///
    /// If the specification mutex has been poisoned, which should not be possible
    pub async fn recv_notification(&mut self) -> ClientNotification<AssetClient> {
        loop {
            tokio::select! {
                biased;
                update = self.device_update_observation.recv_notification() => {
                    // handle the notification
                    // We set auto ack to true, so there's never an ack here to deal with. If we restart, then we'll implicitly
                    // get the update again because we'll pull the latest definition on the restart, so we don't need to get
                    // the notification again.
                    let Some((updated_device, _)) = update else {
                        // if the update notification is None, then the device endpoint has been deleted
                        // Cancel health reporting task
                        self.health_cancellation_token.cancel();
                        // unobserve as cleanup
                        // Spawn a new task to prevent a possible cancellation and ensure the deleted
                        // notification reaches the application.
                        tokio::task::spawn(
                            {
                                let connector_context_clone = self.connector_context.clone();
                                let device_endpoint_ref_clone = self.device_endpoint_ref.clone();
                                async move {
                                    Self::unobserve_device(&connector_context_clone, &device_endpoint_ref_clone, true).await;
                                }
                            }
                        );
                        return ClientNotification::Deleted;
                    };
                    // update self with updated specification
                    let mut unlocked_specification = self.specification.write().unwrap(); // unwrap can't fail unless lock is poisoned
                    *unlocked_specification = DeviceSpecification::new(
                        updated_device,
                        self.connector_context
                            .connector_artifacts
                            .device_endpoint_credentials_mount
                            .as_ref(),
                        &self.device_endpoint_ref.inbound_endpoint_name,
                    ).expect("Device Update Notification should never provide a device that doesn't have the inbound endpoint");

                    return ClientNotification::Updated;
                },
                // Check for completed asset creation
                Some(asset_client_option) = self.asset_completion_rx.recv() => {
                    self.pending_asset_creation = false;
                    if let Some(asset_client) = asset_client_option {
                        return ClientNotification::Created(asset_client);
                    }
                    // If asset_client_option is None, creation failed, continue loop
                },
                create_notification = self.asset_create_observation.recv_notification(), if !self.pending_asset_creation => {
                    let Some((asset_ref, asset_deletion_token)) = create_notification else {
                        // if the create notification is None, then the device endpoint has been deleted
                        log::info!("Device Endpoint Deletion detected, stopping device update observation for {:?}", self.device_endpoint_ref);
                        // Cancel health reporting task
                        self.health_cancellation_token.cancel();
                        // unobserve as cleanup
                        // Spawn a new task to prevent a possible cancellation and ensure the deleted
                        // notification reaches the application.
                        tokio::task::spawn(
                            {
                                let connector_context_clone = self.connector_context.clone();
                                let device_endpoint_ref_clone = self.device_endpoint_ref.clone();
                                async move {
                                    Self::unobserve_device(&connector_context_clone, &device_endpoint_ref_clone, true).await;
                                }
                            }
                        );
                        return ClientNotification::Deleted;
                    };

                    // Start asset creation task
                    self.pending_asset_creation = true;
                    let connector_context = self.connector_context.clone();
                    let specification = self.specification.clone();
                    let status = self.status.clone();
                    let asset_completion_tx = self.asset_completion_tx.clone();

                    tokio::task::spawn(async move {
                        let asset_client = Self::create_asset_client(
                            connector_context,
                            asset_ref,
                            asset_deletion_token,
                            specification,
                            status,
                        ).await;

                        // Always send the result (Some or None) to unblock the receiver
                        let _ = asset_completion_tx.send(asset_client);
                    });
                }
            }
        }
    }

    /// Creates a new status reporter for this [`DeviceEndpointClient`].
    /// The reporter's version snapshot is initialized to the current specification version.
    ///
    /// # Panics
    /// if the specification mutex has been poisoned, which should not be possible
    #[must_use]
    pub fn get_status_reporter(&self) -> DeviceEndpointStatusReporter {
        let snapshotted_version = self.specification.read().unwrap().version.unwrap_or(0);
        DeviceEndpointStatusReporter {
            connector_context: self.connector_context.clone(),
            device_endpoint_status: self.status.clone(),
            device_endpoint_specification: self.specification.clone(),
            health_sender: self.health_sender.clone(),
            device_endpoint_ref: self.device_endpoint_ref.clone(),
            snapshotted_version,
        }
    }

    /// Internal helper to create an [`AssetClient`]
    async fn create_asset_client(
        connector_context: Arc<ConnectorContext>,
        asset_ref: AssetRef,
        asset_deletion_token: CancellationToken,
        specification: Arc<std::sync::RwLock<DeviceSpecification>>,
        status: Arc<tokio::sync::RwLock<DeviceEndpointStatus>>,
    ) -> Option<AssetClient> {
        // Get asset update observation
        let asset_update_observation = match Retry::spawn(
            RETRY_STRATEGY.map(tokio_retry2::strategy::jitter),
            async || -> Result<azure_device_registry::AssetUpdateObservation, RetryError<azure_device_registry::Error>> {
                connector_context
                    .azure_device_registry_client
                    .observe_asset_update_notifications(
                        asset_ref.device_name.clone(),
                        asset_ref.inbound_endpoint_name.clone(),
                        asset_ref.name.clone(),
                        connector_context.azure_device_registry_timeout,
                    )
                    // retry on network errors, otherwise don't retry on config/dev errors
                    .await
                    .map_err(|e| adr_error_into_retry_error(e, "Observe for Asset Updates"))
        }).await {
            Ok(asset_update_observation) => asset_update_observation,
            Err(e) => {
                log::error!("Dropping asset create notification: {asset_ref:?}. Failed to observe for asset update notifications after retries: {e}");
                return None;
            },
        };

        // Get the asset definition
        let asset = match Retry::spawn(
            RETRY_STRATEGY.map(tokio_retry2::strategy::jitter),
            async || -> Result<adr_models::Asset, RetryError<azure_device_registry::Error>> {
                connector_context
                    .azure_device_registry_client
                    .get_asset(
                        asset_ref.device_name.clone(),
                        asset_ref.inbound_endpoint_name.clone(),
                        asset_ref.name.clone(),
                        connector_context.azure_device_registry_timeout,
                    )
                    .await
                    .map_err(|e| adr_error_into_retry_error(e, "Get Asset Definition"))
            },
        )
        .await
        {
            Ok(asset) => asset,
            Err(e) => {
                log::error!(
                    "Dropping asset create notification: {asset_ref:?}. Failed to get Asset definition after retries: {e}"
                );
                // unobserve as cleanup
                AssetClient::unobserve_asset(&connector_context, &asset_ref, false).await;
                return None;
            }
        };

        // Get the asset status
        let asset_status = match Retry::spawn(
            RETRY_STRATEGY.map(tokio_retry2::strategy::jitter),
            async || -> Result<adr_models::AssetStatus, RetryError<azure_device_registry::Error>> {
                connector_context
                    .azure_device_registry_client
                    .get_asset_status(
                        asset_ref.device_name.clone(),
                        asset_ref.inbound_endpoint_name.clone(),
                        asset_ref.name.clone(),
                        connector_context.azure_device_registry_timeout,
                    )
                    .await
                    .map_err(|e| adr_error_into_retry_error(e, "Get Asset Status"))
            },
        )
        .await {
            Ok(asset_status) => asset_status,
            Err(e) => {
                log::error!("Dropping asset create notification: {asset_ref:?}. Failed to get Asset status after retries: {e}");
                // unobserve as cleanup
                AssetClient::unobserve_asset(&connector_context, &asset_ref, false).await;
                return None;
            },
        };

        // Create the AssetClient
        Some(
            AssetClient::new(
                asset,
                asset_status,
                asset_ref,
                specification,
                status,
                asset_update_observation,
                asset_deletion_token,
                connector_context,
            )
            .await,
        )
    }

    // Returns a clone of the current device status
    /// # Panics
    /// if the status mutex has been poisoned, which should not be possible
    #[must_use]
    pub async fn status(&self) -> DeviceEndpointStatus {
        (*self.status.read().await).clone()
    }

    /// Returns a clone of the current device specification
    /// # Panics
    /// if the specification mutex has been poisoned, which should not be possible
    #[must_use]
    pub fn specification(&self) -> DeviceSpecification {
        (*self.specification.read().unwrap()).clone()
    }

    /// Internal convenience function to unobserve from a device's update notifications for cleanup
    async fn unobserve_device(
        connector_context: &Arc<ConnectorContext>,
        device_endpoint_ref: &DeviceEndpointRef,
        is_device_endpoint_deleted: bool,
    ) {
        let _ = Retry::spawn(
            RETRY_STRATEGY.map(tokio_retry2::strategy::jitter),
            async || -> Result<(), RetryError<azure_device_registry::Error>> {
                connector_context
                    .azure_device_registry_client
                    .unobserve_device_update_notifications(
                        device_endpoint_ref.device_name.clone(),
                        device_endpoint_ref.inbound_endpoint_name.clone(),
                        connector_context.azure_device_registry_timeout,
                    )
                    // retry on network errors, otherwise don't retry on config/dev errors
                    .await
                    .map_err(|e| adr_error_into_retry_error(e, "Unobserve for Device Updates"))
            },
        )
        .await
        .inspect_err(|e| {
            // If this unobserve is happening because the device_endpoint has been deleted, the unobserve call isn't necessary and will likely return
            // an error indicating that the device_endpoint can't be found. In the future, we may have a service change to allow us to not need to
            // call unobserve in these cases to clean up our local state, but for now this call is needed and the error should be ignored.
            if is_device_endpoint_deleted {
                match e.kind() {
                    azure_device_registry::ErrorKind::ServiceError(_) => {
                        log::debug!("Expected failure unobserving device update notifications for {device_endpoint_ref:?}. Since the device/endpoint has been deleted, this is expected and not an error if the error indicates that the device/endpoint is not found: {e}");
                    },
                    _ => {
                        log::warn!("Failed to unobserve device update notifications for deleted {device_endpoint_ref:?} after retries. If error indicates that the device/endpoint is not found, this is expected and not a true error: {e}");
                    }
                }
            } else {
                log::error!("Failed to unobserve device update notifications for {device_endpoint_ref:?} after retries. If error indicates that the device/endpoint is not found, this may be expected and not a true error: {e}");
            }
        });
    }
}

impl Drop for DeviceEndpointClient {
    fn drop(&mut self) {
        self.health_cancellation_token.cancel();
    }
}

/// A cloneable status reporter for Asset status reporting.
///
/// This provides a way to report Asset status changes from outside the [`AssetClient`].
#[derive(Clone, Debug)]
pub struct AssetStatusReporter {
    connector_context: Arc<ConnectorContext>,
    asset_status: Arc<tokio::sync::RwLock<adr_models::AssetStatus>>,
    asset_specification: Arc<std::sync::RwLock<AssetSpecification>>,
    asset_ref: AssetRef,
}

impl AssetStatusReporter {
    /// Used to conditionally report the asset status and then updates the asset with the new status returned.
    ///
    /// The `modify` function is called with the current asset status (if any) and should return:
    /// - `Some(new_status)` if the status should be updated and reported
    /// - `None` if no update is needed
    ///
    /// # Returns
    /// - [`ModifyResult::Reported`] if the status was updated and successfully reported
    /// - [`ModifyResult::NotModified`] if no modification was needed or the version changed during processing
    ///
    /// # Errors
    /// [`azure_device_registry::Error`] of kind [`AIOProtocolError`](azure_device_registry::ErrorKind::AIOProtocolError) if
    /// there are any underlying errors from the AIO RPC protocol. This error will be retried
    /// 10 times with exponential backoff and jitter and only returned if it still is failing.
    ///
    /// [`azure_device_registry::Error`] of kind [`ServiceError`](azure_device_registry::ErrorKind::ServiceError) if an error is returned
    /// by the Azure Device Registry service.
    ///
    /// # Panics
    /// if the specification mutex has been poisoned, which should not be possible
    pub async fn report_status_if_modified<F>(
        &self,
        modify: F,
    ) -> Result<ModifyResult, azure_device_registry::Error>
    where
        F: Fn(Option<Result<(), &AdrConfigError>>) -> Option<Result<(), AdrConfigError>>,
    {
        // Get the current version of the asset specification
        let cached_version = self.asset_specification.read().unwrap().version;

        {
            // Get the current asset status
            let status_read_guard = self.asset_status.read().await;
            let current_asset_status =
                AssetClient::get_current_asset_status(&status_read_guard, cached_version);

            let modify_input = Self::get_modify_input(&current_asset_status);

            // We do not use the result since the status once we acquire the read lock might change
            // and we will have to re-evaluate
            let Some(_modify_result) = modify(modify_input) else {
                // If no modification is needed, return Ok
                return Ok(ModifyResult::NotModified);
            };
        }

        // To modify, we need to acquire the write lock
        let mut status_write_guard = self.asset_status.write().await;

        if cached_version != self.asset_specification.read().unwrap().version {
            // Our modify is no longer valid
            log::debug!(
                "Asset specification is out-of-date for {:?}; skipping modification",
                self.asset_ref
            );
            return Ok(ModifyResult::NotModified);
        }

        // We can continue here because the asset status will not change. The specification might
        // update but we will be reporting for an old version.

        // Get the current asset status in case it has changed
        let current_asset_status =
            AssetClient::get_current_asset_status(&status_write_guard, cached_version);

        let modify_result = {
            let modify_input = Self::get_modify_input(&current_asset_status);

            let Some(modify_result) = modify(modify_input) else {
                // If no modification is needed, return Ok
                return Ok(ModifyResult::NotModified);
            };

            modify_result
        };

        let mut asset_status_to_report = current_asset_status.into_owned();

        // Update the config status
        asset_status_to_report.config = Some(azure_device_registry::ConfigStatus {
            version: cached_version,
            error: modify_result.clone().err(),
            last_transition_time: Some(chrono::Utc::now()),
        });

        log::debug!("Reporting asset status from app for {:?}", self.asset_ref);

        Self::internal_report_status(
            asset_status_to_report,
            &self.connector_context,
            &self.asset_ref,
            &mut status_write_guard,
            "AssetClient::report_status_if_modified",
        )
        .await?;

        Ok(ModifyResult::Reported)
    }

    fn get_modify_input(
        current_status: &adr_models::AssetStatus,
    ) -> Option<Result<(), &AdrConfigError>> {
        current_status
            .config
            .as_ref()
            .map(|config| match &config.error {
                Some(err) => Err(err),
                None => Ok(()),
            })
    }

    pub(crate) async fn internal_report_status(
        adr_asset_status: adr_models::AssetStatus,
        connector_context: &ConnectorContext,
        asset_ref: &AssetRef,
        asset_status_ref: &mut adr_models::AssetStatus,
        log_identifier: &str,
    ) -> Result<(), azure_device_registry::Error> {
        // send status update to the service
        let updated_asset_status = Retry::spawn(
            RETRY_STRATEGY.map(tokio_retry2::strategy::jitter).take(10),
            async || -> Result<adr_models::AssetStatus, RetryError<azure_device_registry::Error>> {
                connector_context
                    .azure_device_registry_client
                    .update_asset_status(
                        asset_ref.device_name.clone(),
                        asset_ref.inbound_endpoint_name.clone(),
                        asset_ref.name.clone(),
                        adr_asset_status.clone(),
                        connector_context.azure_device_registry_timeout,
                    )
                    .await
                    .map_err(|e| adr_error_into_retry_error(e, &format!("Update Asset Status for {log_identifier}")))
            },
        )
        .await?;
        // update self with new returned status
        *asset_status_ref = updated_asset_status;
        Ok(())
    }
}

/// Struct used to hold the updates for an Asset's components
/// until all components have been processed and the function
/// using this struct is beyond the point where it needs to worry about
/// cancel safety.
struct AssetComponentUpdates {
    new_status: adr_models::AssetStatus,
    status_updated: bool,
    data_operation_updates: Vec<(
        watch::Sender<DataOperationUpdateNotification>,
        DataOperationUpdateNotification,
    )>,
    management_action_updates: Vec<(
        watch::Sender<ManagementActionUpdateNotification>,
        ManagementActionUpdateNotification,
    )>,
    new_asset_component_clients: Vec<AssetComponentClient>,
}

/// Azure Device Registry Asset that includes additional functionality
/// to report status and receive Asset updates
#[derive(Debug, Getters)]
pub struct AssetClient {
    /// Asset, device, and inbound endpoint names
    asset_ref: AssetRef,
    /// Specification for the Asset
    #[getter(skip)]
    specification: Arc<std::sync::RwLock<AssetSpecification>>,
    /// Status for the Asset
    #[getter(skip)]
    status: Arc<tokio::sync::RwLock<adr_models::AssetStatus>>,
    /// Specification of the device that this Asset is tied to
    #[getter(skip)]
    device_specification: Arc<std::sync::RwLock<DeviceSpecification>>,
    /// Status of the device that this Asset is tied to
    #[getter(skip)]
    device_status: Arc<tokio::sync::RwLock<DeviceEndpointStatus>>,
    // Internally used fields
    /// Internal `CancellationToken` for when the Asset is deleted. Surfaced to the user through the receive update flow
    #[getter(skip)]
    asset_deletion_token: CancellationToken,
    /// The internal observation for updates
    #[getter(skip)]
    asset_update_observation: azure_device_registry::AssetUpdateObservation,
    /// Internal watcher receiver that holds a snapshot of the latest update and whether it has
    /// been fully processed or not
    #[getter(skip)]
    asset_update_watcher_rx: watch::Receiver<Asset>,
    /// Internal watcher sender that sends the latest update
    #[getter(skip)]
    asset_update_watcher_tx: watch::Sender<Asset>,
    /// Internal sender for when new asset components are created
    #[getter(skip)]
    asset_component_creation_tx: UnboundedSender<AssetComponentClient>,
    /// Internal channel for receiving notifications about asset component creation events.
    #[getter(skip)]
    asset_component_creation_rx: UnboundedReceiver<AssetComponentClient>,
    /// Internal watch sender for releasing asset component create/update notifications
    #[getter(skip)]
    release_asset_component_notifications_tx: watch::Sender<()>,
    /// hashmap of current dataset names to their sender to send dataset updates
    #[getter(skip)]
    dataset_hashmap: HashMap<String, watch::Sender<DataOperationUpdateNotification>>,
    /// hashmap of current event/event group names to their sender to send event updates
    #[getter(skip)]
    event_hashmap: HashMap<
        (String, String), // (EventGroup name, Event name)
        watch::Sender<DataOperationUpdateNotification>,
    >,
    /// hashmap of current stream names to their sender to send stream updates
    #[getter(skip)]
    stream_hashmap: HashMap<String, watch::Sender<DataOperationUpdateNotification>>,
    /// hashmap of current management action/management group names to their current definitions and a sender to send action updates
    #[getter(skip)]
    management_action_hashmap: HashMap<
        (String, String), // (ManagementGroup name, ManagementAction name)
        watch::Sender<ManagementActionUpdateNotification>,
    >,
    #[getter(skip)]
    connector_context: Arc<ConnectorContext>,
}

impl AssetClient {
    #[allow(clippy::too_many_arguments)]
    pub(crate) async fn new(
        asset: adr_models::Asset,
        asset_status: adr_models::AssetStatus,
        asset_ref: AssetRef,
        device_specification: Arc<std::sync::RwLock<DeviceSpecification>>,
        device_status: Arc<tokio::sync::RwLock<DeviceEndpointStatus>>,
        asset_update_observation: azure_device_registry::AssetUpdateObservation,
        asset_deletion_token: CancellationToken,
        connector_context: Arc<ConnectorContext>,
    ) -> Self {
        let (asset_update_watcher_tx, asset_update_watcher_rx) = watch::channel(asset.clone());
        let specification = AssetSpecification::from(asset.clone());
        let specification_version = specification.version;
        let (asset_component_creation_tx, asset_component_creation_rx) = mpsc::unbounded_channel();

        // Create the AssetClient so that we can use the same helper functions for processing the asset components as we do during the update flow
        let mut asset_client = AssetClient {
            asset_ref,
            specification: Arc::new(std::sync::RwLock::new(specification)),
            status: Arc::new(tokio::sync::RwLock::new(asset_status)),
            device_specification,
            device_status,
            asset_update_observation,
            asset_update_watcher_rx,
            asset_update_watcher_tx,
            asset_component_creation_tx,
            asset_component_creation_rx,
            dataset_hashmap: HashMap::new(),
            event_hashmap: HashMap::new(),
            stream_hashmap: HashMap::new(),
            management_action_hashmap: HashMap::new(),
            connector_context,
            release_asset_component_notifications_tx: watch::Sender::new(()),
            asset_deletion_token,
        };

        {
            // lock the status write guard so that no other threads can modify the status while we update it
            // (not possible in new, but allows use of Self:: helper fns)
            let mut status_write_guard = asset_client.status.write().await;
            // if there are any config errors when parsing the asset, collect them all so we can report them at once
            let mut updates = AssetComponentUpdates {
                new_status: Self::get_current_asset_status(
                    &status_write_guard,
                    specification_version,
                )
                .into_owned(),
                status_updated: false,
                data_operation_updates: Vec::new(),
                new_asset_component_clients: Vec::new(),
                management_action_updates: Vec::new(),
            };
            // Handle "updates" for each type of asset component. Since we don't currently have any asset
            // components tracked yet, everything in the definition will be treated as a new asset component.
            let mut temp_dataset_hashmap = asset_client.dataset_hashmap.clone();
            // asset_client.dataset_hashmap will be empty, so all datasets will be treated as new (as it should be).
            // Note that I could use vec::new() for temp_dataset_hashmap, but for extra safety, I'll clone the asset's dataset hashmap instead
            asset_client.handle_data_operation_kind_updates(
                &mut temp_dataset_hashmap,
                &asset,
                &asset.datasets,
                &mut updates,
            );
            asset_client.dataset_hashmap = temp_dataset_hashmap;

            let mut temp_event_hashmap = asset_client.event_hashmap.clone();
            asset_client.handle_event_group_updates(
                &mut temp_event_hashmap,
                &asset,
                &asset.event_groups,
                &mut updates,
            );
            asset_client.event_hashmap = temp_event_hashmap;

            let mut temp_stream_hashmap = asset_client.stream_hashmap.clone();
            asset_client.handle_data_operation_kind_updates(
                &mut temp_stream_hashmap,
                &asset,
                &asset.streams,
                &mut updates,
            );
            asset_client.stream_hashmap = temp_stream_hashmap;

            let mut temp_management_action_hashmap = asset_client.management_action_hashmap.clone();
            asset_client.handle_management_action_updates(
                &mut temp_management_action_hashmap,
                &asset.management_groups,
                &mut updates,
            );
            asset_client.management_action_hashmap = temp_management_action_hashmap;

            // if there were any config errors, report them to the ADR service together
            if updates.status_updated {
                log::debug!(
                    "Reporting error asset status on new for {:?}",
                    asset_client.asset_ref
                );
                if let Err(e) = AssetStatusReporter::internal_report_status(
                    updates.new_status,
                    &asset_client.connector_context,
                    &asset_client.asset_ref,
                    &mut status_write_guard,
                    "AssetClient::new",
                )
                .await
                {
                    log::error!(
                        "Failed to report error Asset status for new Asset {:?}: {e}",
                        asset_client.asset_ref
                    );
                }
            }

            // Send all new asset component create notifications
            for new_asset_component_client in updates.new_asset_component_clients {
                // error is not possible since the receiving side of the channel is owned by the AssetClient
                let _ = asset_client
                    .asset_component_creation_tx
                    .send(new_asset_component_client);
            }
            // Note, updates.data_operation_updates/management_action_updates is not used because there will be no updates on new
        }

        asset_client
    }

    /// Helper function to handle updates for all management actions on an Asset
    ///
    /// Parses and validates all Asset updates pertaining to management actions: detects any deleted, updated, and new management actions
    /// Modifies `updates` and `management_action_hashmap` in place:
    ///     Modifies `updates.new_status` with any validation errors found
    ///     Adds any management action updates to `updates.management_action_updates` that can be sent after the update task can't be cancelled
    ///     Adds any new Management Action Clients to `updates.new_asset_component_clients` that can be sent after the update task can't be cancelled
    ///     Removes any deleted Management Actions from the `management_action_hashmap` that can be applied after the update task can't be cancelled
    fn handle_management_action_updates(
        &self,
        management_action_hashmap: &mut HashMap<
            (String, String), // (ManagementGroup name, ManagementAction name)
            watch::Sender<ManagementActionUpdateNotification>,
        >,
        updated_management_groups: &[adr_models::ManagementGroup],
        updates: &mut AssetComponentUpdates,
    ) {
        // remove the management actions that are no longer present in the new asset definition.
        // This triggers deletion notification since this drops the update sender.
        management_action_hashmap.retain(|(management_group_name, action_name), _| {
            if let Some(updated_management_group) = updated_management_groups
                .iter()
                .find(|management_group| management_group.name == *management_group_name)
            {
                updated_management_group
                    .actions
                    .iter()
                    .any(|updated_action| updated_action.name == *action_name)
            } else {
                false
            }
        });

        // For all received actions, check if the existing action needs to be updated or if a new one needs to be created
        for received_management_group in updated_management_groups {
            for received_management_action in &received_management_group.actions {
                let update_notification = ManagementActionUpdateNotification {
                    definition: (
                        received_management_group.clone(),
                        received_management_action.clone(),
                    )
                        .into(),
                    release_asset_component_notifications_rx: self
                        .release_asset_component_notifications_tx
                        .subscribe(),
                };
                // if it already exists
                if let Some(management_action_update_tx) =
                    management_action_hashmap.get(update_notification.definition.hash_name())
                {
                    // save update to send to the action after the task can't get cancelled
                    // Send an update whether the action changed or not, because the client needs to know the asset was updated at minimum
                    updates
                        .management_action_updates
                        .push((management_action_update_tx.clone(), update_notification));
                }
                // it needs to be created
                else {
                    let received_management_specification = update_notification.definition.clone();
                    let (management_action_update_watcher_tx, management_action_update_watcher_rx) =
                        watch::channel(update_notification);
                    let (new_management_action_client, new_executor_res) =
                        ManagementActionClient::new(
                            received_management_specification,
                            management_action_update_watcher_rx,
                            self.asset_ref.clone(),
                            self.status.clone(),
                            self.specification.clone(),
                            self.device_specification.clone(),
                            self.device_status.clone(),
                            self.connector_context.clone(),
                        );
                    // if there were errors, report them to the status to be send to ADR
                    if let Err(e) = &new_executor_res {
                        // add the error to the status to be reported to ADR, and then continue to process
                        // other actions
                        ManagementActionClient::update_action_status(
                            &mut updates.new_status,
                            &received_management_group.name,
                            &received_management_action.name,
                            Err(e.clone()),
                        );
                        updates.status_updated = true;
                    }
                    // insert the management action into the hashmap so we can handle updates
                    management_action_hashmap.insert(
                        new_management_action_client.definition.hash_name().clone(),
                        management_action_update_watcher_tx,
                    );
                    // save the new management action client to be sent on self.asset_component_creation_tx after the task can't get cancelled
                    updates.new_asset_component_clients.push(
                        AssetComponentClient::ManagementAction((
                            new_management_action_client,
                            new_executor_res,
                        )),
                    );
                }
            }
        }
    }

    /// Helper function to handle updates for event groups on an Asset
    /// Transforms a list of `adr_models::EventGroup`s into a list of
    /// [`EventSpecification`]s per Event and calls `handle_data_operation_kind_updates`
    fn handle_event_group_updates(
        &self,
        event_hashmap: &mut HashMap<
            (String, String), // (eg name, event name)
            watch::Sender<DataOperationUpdateNotification>,
        >,
        updated_asset: &Asset,
        updated_asset_event_groups: &[adr_models::EventGroup],
        updates: &mut AssetComponentUpdates,
    ) {
        let mut updated_asset_events: Vec<EventSpecification> = Vec::new();
        for event_group in updated_asset_event_groups {
            // Creates an [`EventSpecification`] for each event in the event group
            for event in &event_group.events {
                updated_asset_events.push((event_group.clone(), event.clone()).into());
            }
        }

        // Do the actual updates now that we have the correct data format
        self.handle_data_operation_kind_updates(
            event_hashmap,
            updated_asset,
            &updated_asset_events,
            updates,
        );
    }

    /// Helper function to handle updates for all of a type of data operations on an Asset
    /// This reduces duplicate code for each data operation kind - instead this function is called once for each
    ///
    /// Parses and validates all Asset updates pertaining to this data operation kind:
    ///     Detects any deleted, updated, and new data operations
    ///     Parses the default destination for that data operation
    /// Modifies `updates` and `data_operation_hashmap` in place:
    ///     Modifies `updates.new_status` with any validation errors found
    ///     Adds any Data Operation updates to `updates.data_operation_updates` that can be sent after the update task can't be cancelled
    ///     Adds any new Data Operation Clients to `updates.new_asset_component_clients` that can be sent after the update task can't be cancelled
    ///     Removes any deleted Data Operations from the `data_operation_hashmap` that can be applied after the update task can't be cancelled
    fn handle_data_operation_kind_updates<T: Clone + DataOperation + PartialEq>(
        &self,
        data_operation_hashmap: &mut HashMap<
            T::HashName,
            watch::Sender<DataOperationUpdateNotification>,
        >,
        updated_asset: &Asset,
        updated_asset_data_operations: &[T],
        updates: &mut AssetComponentUpdates,
    ) {
        // remove the data operations that are no longer present in the new asset definition.
        // This triggers deletion notification since this drops the update sender.
        data_operation_hashmap.retain(|data_operation_name, _| {
            updated_asset_data_operations
                .iter()
                .any(|data_operation| data_operation.hash_name() == data_operation_name)
        });
        // Get the new default data operation destinations and track whether they're different or not from the current one
        let default_data_operation_destination_updated = match T::kind() {
            DataOperationKind::Dataset => {
                updated_asset.default_datasets_destinations
                    != self
                        .specification
                        .read()
                        .unwrap()
                        .default_datasets_destinations
            }
            DataOperationKind::Event => {
                updated_asset.default_events_destinations
                    != self
                        .specification
                        .read()
                        .unwrap()
                        .default_events_destinations
            }
            DataOperationKind::Stream => {
                updated_asset.default_streams_destinations
                    != self
                        .specification
                        .read()
                        .unwrap()
                        .default_streams_destinations
            }
        };

        let default_destinations_result = match T::kind() {
            DataOperationKind::Dataset => {
                destination_endpoint::Destination::new_dataset_destinations(
                    &updated_asset.default_datasets_destinations,
                    &self.asset_ref,
                    updated_asset.uuid.as_ref(),
                    updated_asset.external_asset_id.as_ref(),
                    &self.connector_context,
                )
            }
            DataOperationKind::Event => {
                // TODO: To support default destinations on the event group in the future,
                // we can create a hashmap of eventGroups to their default destinations here.
                // However, any failures need to be reported only on an event that would actually
                // use that bad configuration
                destination_endpoint::Destination::new_event_stream_destinations(
                    &updated_asset.default_events_destinations,
                    &self.asset_ref,
                    updated_asset.uuid.as_ref(),
                    updated_asset.external_asset_id.as_ref(),
                    &self.connector_context,
                )
            }
            DataOperationKind::Stream => {
                destination_endpoint::Destination::new_event_stream_destinations(
                    &updated_asset.default_streams_destinations,
                    &self.asset_ref,
                    updated_asset.uuid.as_ref(),
                    updated_asset.external_asset_id.as_ref(),
                    &self.connector_context,
                )
            }
        };
        let default_data_operation_destinations = match default_destinations_result {
            Ok(res) => res.into_iter().map(Arc::new).collect(),
            Err(e) => {
                log::error!(
                    "Invalid default data operation destination for Asset {:?}: {e:?}",
                    self.asset_ref
                );
                // Add this to the status to be reported to ADR
                updates.new_status.config = Some(azure_device_registry::ConfigStatus {
                    version: updated_asset.version,
                    error: Some(e),
                    last_transition_time: Some(chrono::Utc::now()),
                });
                updates.status_updated = true;
                // set this to an empty vec instead of skipping parsing the rest of the asset because if all
                // data operations have a destination specified, this might not cause the asset to be unusable
                vec![]
            }
        };

        // For all received data operations, check if the existing data operation needs an update or if a new one needs to be created
        for received_data_operation in updated_asset_data_operations {
            // it already exists
            if let Some(data_operation_update_tx) =
                data_operation_hashmap.get(received_data_operation.hash_name())
            {
                // TODO: To support default destinations on the event group in the future,
                // we can check the previous event group default destinations hashmap here
                // and overwrite the asset event default destination if it's present before
                // sending the update.
                // To support splitting updates for the eventGroup as a separate notification
                // from updates on an event, that logic would need to be added here as well.

                // save update to send to the data operation after the task can't get cancelled
                // Send an update whether the data operation changed or not, because the client needs to know the asset was updated at minimum
                updates.data_operation_updates.push((
                    data_operation_update_tx.clone(),
                    DataOperationUpdateNotification {
                        definition: received_data_operation.clone().into(),
                        default_destinations: default_data_operation_destinations.clone(),
                        default_destination_has_changed: default_data_operation_destination_updated,
                        release_asset_component_notifications_rx: self
                            .release_asset_component_notifications_tx
                            .subscribe(),
                    },
                ));
            }
            // it needs to be created
            else {
                // TODO: To support default destinations on the event group in the future,
                // here we also need to determine whether to pass in the asset or eventGroup
                // default event destination

                let data_operation_definition = received_data_operation.clone().into();
                let (data_operation_update_watcher_tx, data_operation_update_watcher_rx) =
                    watch::channel(DataOperationUpdateNotification {
                        definition: data_operation_definition.clone(),
                        default_destinations: default_data_operation_destinations.clone(),
                        default_destination_has_changed: true,
                        release_asset_component_notifications_rx: self
                            .release_asset_component_notifications_tx
                            .subscribe(),
                    });
                let (new_data_operation_client, data_operation_client_result) =
                    DataOperationClient::new(
                        data_operation_definition,
                        data_operation_update_watcher_rx,
                        &default_data_operation_destinations,
                        self.asset_ref.clone(),
                        self.status.clone(),
                        self.specification.clone(),
                        self.device_specification.clone(),
                        self.device_status.clone(),
                        self.connector_context.clone(),
                    );
                // insert the data operation client into the hashmap so we can handle updates
                data_operation_hashmap.insert(
                    received_data_operation.hash_name().clone(),
                    data_operation_update_watcher_tx,
                );
                // save new data operation client to be sent on self.asset_component_creation_tx after the task can't get cancelled
                updates
                    .new_asset_component_clients
                    .push(AssetComponentClient::DataOperation((
                        new_data_operation_client,
                        data_operation_client_result.clone(),
                    )));

                // If there were errors, report them to the status to be sent to ADR
                if let Err(e) = data_operation_client_result {
                    // Add the error to the status to be reported to ADR, and then continue to process
                    // other data operations even if one isn't valid.
                    match received_data_operation.data_operation_name() {
                        DataOperationName::Dataset {
                            name: ref dataset_name,
                        } => {
                            DataOperationClient::update_dataset_status(
                                &mut updates.new_status,
                                dataset_name,
                                Err(e),
                            );
                        }
                        DataOperationName::Event {
                            name: ref event_name,
                            ref event_group_name,
                        } => {
                            DataOperationClient::update_event_status(
                                &mut updates.new_status,
                                event_group_name,
                                event_name,
                                Err(e),
                            );
                        }
                        DataOperationName::Stream {
                            name: ref stream_name,
                        } => {
                            DataOperationClient::update_stream_status(
                                &mut updates.new_status,
                                stream_name,
                                Err(e),
                            );
                        }
                    }
                    updates.status_updated = true;
                }
            }
        }
    }

    /// Helper function to handle an asset update
    async fn handle_update(
        &mut self,
        updated_asset: Asset,
    ) -> ClientNotification<AssetComponentClient> {
        // lock the status write guard so that no other threads can modify the status while we update it
        let mut status_write_guard = self.status.write().await;

        let mut adr_asset_status =
            AssetClient::get_current_asset_status(&status_write_guard, updated_asset.version)
                .into_owned();

        // Update the config status
        adr_asset_status.config = match adr_asset_status.config {
            Some(mut config) => {
                config.last_transition_time = Some(Utc::now());
                Some(config)
            }
            None => {
                // If the config is None, we need to create a new one to report along with the asset component statuses
                Some(azure_device_registry::ConfigStatus {
                    version: updated_asset.version,
                    last_transition_time: Some(Utc::now()),
                    ..Default::default()
                })
            }
        };

        // if there are any config errors when parsing the asset, collect them all so we can report them at once
        // track all asset components to update and save notifications for once the task can't be cancelled
        let mut updates = AssetComponentUpdates {
            new_status: adr_asset_status,
            status_updated: false,
            data_operation_updates: Vec::new(),
            management_action_updates: Vec::new(),
            new_asset_component_clients: Vec::new(),
        };

        // Handle updates for each type of asset component
        // make changes to copies of the asset component hashmaps so that this function is cancel safe
        let mut temp_dataset_hashmap = self.dataset_hashmap.clone();
        self.handle_data_operation_kind_updates(
            &mut temp_dataset_hashmap,
            &updated_asset,
            &updated_asset.datasets,
            &mut updates,
        );
        let mut temp_event_hashmap = self.event_hashmap.clone();
        self.handle_event_group_updates(
            &mut temp_event_hashmap,
            &updated_asset,
            &updated_asset.event_groups,
            &mut updates,
        );
        let mut temp_stream_hashmap = self.stream_hashmap.clone();
        self.handle_data_operation_kind_updates(
            &mut temp_stream_hashmap,
            &updated_asset,
            &updated_asset.streams,
            &mut updates,
        );
        let mut temp_management_action_hashmap = self.management_action_hashmap.clone();
        self.handle_management_action_updates(
            &mut temp_management_action_hashmap,
            &updated_asset.management_groups,
            &mut updates,
        );

        // if there were any config errors, report them to the ADR service together
        if updates.status_updated {
            log::debug!(
                "Reporting error asset status on recv_notification for {:?}",
                self.asset_ref
            );
            if let Err(e) = AssetStatusReporter::internal_report_status(
                updates.new_status,
                &self.connector_context,
                &self.asset_ref,
                &mut status_write_guard,
                "AssetClient::recv_notification",
            )
            .await
            {
                log::error!(
                    "Failed to report error Asset status for updated Asset {:?}: {e}",
                    self.asset_ref
                );
            }
        }

        // update specification
        let mut unlocked_specification = self.specification.write().unwrap(); // unwrap can't fail unless lock is poisoned
        *unlocked_specification = AssetSpecification::from(updated_asset);

        // update asset component hashmaps now that this task can't be cancelled
        self.dataset_hashmap = temp_dataset_hashmap;
        self.event_hashmap = temp_event_hashmap;
        self.stream_hashmap = temp_stream_hashmap;
        self.management_action_hashmap = temp_management_action_hashmap;

        // send all notifications associated with this asset update
        for (data_operation_update_tx, data_operation_update_notification) in
            updates.data_operation_updates
        {
            // send update to the data operation
            let _ = data_operation_update_tx.send(data_operation_update_notification).inspect_err(|tokio::sync::watch::error::SendError(DataOperationUpdateNotification {definition, ..})| {
                // TODO: should this trigger the DataOperationClient create flow, or is this just indicative of an application bug?
                log::warn!(
                    "Update received for {} on asset {:?}, but DataOperationClient has been dropped",
                    definition.name(),
                    self.asset_ref
                );
            });
        }
        for (management_action_update_tx, management_action_update_notification) in
            updates.management_action_updates
        {
            // send update to the management action
            let _ = management_action_update_tx.send(management_action_update_notification).inspect_err(|tokio::sync::watch::error::SendError(ManagementActionUpdateNotification {definition, ..})| {
                log::warn!(
                    "Update received for management action {} on asset {:?}, but ManagementActionClient has been dropped",
                    definition.command_name(),
                    self.asset_ref
                );
            });
        }
        for new_asset_component_client in updates.new_asset_component_clients {
            // error is not possible since the receiving side of the channel is owned by the AssetClient
            let _ = self
                .asset_component_creation_tx
                .send(new_asset_component_client);
        }

        // Asset update has been fully processed, mark as seen.
        self.asset_update_watcher_rx.mark_unchanged();
        ClientNotification::Updated
    }

    /// Used to receive notifications related to the Asset from the Azure Device
    /// Registry Service.
    ///
    /// Returns [`ClientNotification::Updated`] if the Asset Specification has
    /// been updated in place.
    ///
    /// Returns [`ClientNotification::Deleted`] if the Asset has been deleted.
    /// The [`AssetClient`] should not be used after this point, and no more
    /// notifications will be received.
    ///
    /// Returns [`ClientNotification::Created`] with a new [`AssetComponentClient`] if a new dataset, event, stream, or
    /// management action has been created. Within the [`AssetComponentClient`] enum, the Client and any configuration
    /// errors detected during creation are included. If there was an error detected, the [`AssetComponentClient`] should
    /// not be used until an Ok update is received. The [`AssetComponentClient`] can be used to receive updates for the new component.
    ///
    /// Receiving an update will also trigger update/deletion notifications for asset components that
    /// are linked to this asset. To ensure the asset update is received before asset component notifications,
    /// asset component notifications won't be released until this function is polled again after receiving an
    /// update.
    ///
    /// # Cancel safety
    /// This method is cancel safe. If you use it as the event in a `tokio::select!` statement and some other branch
    /// completes first, then it is guaranteed that no asset notifications will be lost, and the asset will not
    /// be updated without a notification being returned.
    ///
    /// # Panics
    /// If the specification mutex has been poisoned, which should not be possible
    pub async fn recv_notification(&mut self) -> ClientNotification<AssetComponentClient> {
        // release any pending asset component create/update notifications
        self.release_asset_component_notifications_tx
            .send_modify(|()| ());
        tokio::select! {
            biased;
            () = self.asset_deletion_token.cancelled() => {
                log::info!("Asset deletion token received, stopping asset update observation for {:?}", self.asset_ref);
                // unobserve as cleanup
                // Spawn a new task to prevent a possible cancellation and ensure the deleted
                // notification reaches the application.
                tokio::task::spawn(
                    {
                        let connector_context_clone = self.connector_context.clone();
                        let asset_ref_clone = self.asset_ref.clone();
                        async move {
                            Self::unobserve_asset(&connector_context_clone, &asset_ref_clone, true).await;
                        }
                    }
                );
                ClientNotification::Deleted
            },
            notification = self.asset_update_observation.recv_notification() => {
                // We set auto ack to true, so there's never an ack here to deal with. If we restart, then we'll implicitly
                // get the update again because we'll pull the latest definition on the restart, so we don't need to get
                // the notification again as an MQTT message.
                if let Some((updated_asset, _)) = notification {
                    // The asset client holds the receiving end so `send` will never return an error,
                    // we are saving the updated asset in the watcher in case `handle_update` is
                    // cancelled
                    let _ = self.asset_update_watcher_tx.send(updated_asset.clone());
                    self.handle_update(updated_asset).await
                } else {
                    // unobserve as cleanup
                    // Spawn a new task to prevent a possible cancellation and ensure the deleted
                    // notification reaches the application.
                    tokio::task::spawn(
                        {
                            let connector_context_clone = self.connector_context.clone();
                            let asset_ref_clone = self.asset_ref.clone();
                            async move {
                                Self::unobserve_asset(&connector_context_clone, &asset_ref_clone, true).await;
                            }
                        }
                    );
                    // Asset update has been fully processed, mark as seen.
                    self.asset_update_watcher_rx.mark_unchanged();
                    ClientNotification::Deleted
                }
            },
            // An error here is not possible since the asset client holds the sender end
            Ok(()) = self.asset_update_watcher_rx.changed() => {
                // We need to mark as changed in case `handle_update` is cancelled so we can enter
                // this branch again
                self.asset_update_watcher_rx.mark_changed();
                let updated_asset = self.asset_update_watcher_rx.borrow().clone();

                self.handle_update(updated_asset).await
            }
            create_notification = self.asset_component_creation_rx.recv() => {
                let Some(asset_component_client) = create_notification else {
                    // unobserve as cleanup
                    // Spawn a new task to prevent a possible cancellation and ensure the deleted
                    // notification reaches the application.
                    tokio::task::spawn(
                        {
                            let connector_context_clone = self.connector_context.clone();
                            let asset_ref_clone = self.asset_ref.clone();
                            async move {
                                Self::unobserve_asset(&connector_context_clone, &asset_ref_clone, true).await;
                            }
                        }
                    );
                    return ClientNotification::Deleted;
                };
                ClientNotification::Created(asset_component_client)
            },
        }
    }

    // Returns a clone of the current asset specification
    /// # Panics
    /// if the specification mutex has been poisoned, which should not be possible
    #[must_use]
    pub fn specification(&self) -> AssetSpecification {
        (*self.specification.read().unwrap()).clone()
    }

    /// Returns a clone of the current asset status
    /// Note that this is the value of the last reported status (or the status
    /// from when the asset was first received if one hasn't been reported yet),
    /// so it is possible that this is not the latest value that the ADR service has
    #[must_use]
    pub async fn status(&self) -> adr_models::AssetStatus {
        (*self.status.read().await).clone()
    }

    /// Returns a clone of the current device specification
    /// # Panics
    /// if the device specification mutex has been poisoned, which should not be possible
    #[must_use]
    pub fn device_specification(&self) -> DeviceSpecification {
        (*self.device_specification.read().unwrap()).clone()
    }

    /// Returns a clone of the current device status
    #[must_use]
    pub async fn device_status(&self) -> DeviceEndpointStatus {
        (*self.device_status.read().await).clone()
    }

    /// Creates a new status reporter for this [`AssetClient`]
    #[must_use]
    pub fn get_status_reporter(&self) -> AssetStatusReporter {
        AssetStatusReporter {
            connector_context: self.connector_context.clone(),
            asset_status: self.status.clone(),
            asset_specification: self.specification.clone(),
            asset_ref: self.asset_ref.clone(),
        }
    }

    /// Internal helper to get an [`adr_models::AssetStatus`] that can be used as a starting place
    /// to modify the current status with whatever new things we want to report.
    ///
    /// Note that it returns a `Cow`. The reason is that most of the times that we are reporting
    /// a status we will not end up modifying it. `Cow` allows us to only clone when we are going to
    /// modify.
    pub(crate) fn get_current_asset_status(
        current_status: &adr_models::AssetStatus,
        adr_version: Option<u64>,
    ) -> Cow<'_, adr_models::AssetStatus> {
        match &current_status.config {
            Some(config) => {
                if config.version == adr_version {
                    // If the version in our config matches the one in ADR we return our current
                    // asset status
                    Cow::Borrowed(current_status)
                } else {
                    // If the version doesn't match, the config version we are holding is stale so
                    // we pass a cleared version
                    Cow::Owned(adr_models::AssetStatus::default())
                }
            }
            None => {
                // If there is no config, then the status is not set, we can pass our current asset
                // status
                Cow::Borrowed(current_status)
            }
        }
    }

    /// Internal convenience function to unobserve from an asset's update notifications for cleanup
    async fn unobserve_asset(
        connector_context: &Arc<ConnectorContext>,
        asset_ref: &AssetRef,
        is_asset_deleted: bool,
    ) {
        // unobserve as cleanup
        let _ = Retry::spawn(
            RETRY_STRATEGY.map(tokio_retry2::strategy::jitter),
            async || -> Result<(), RetryError<azure_device_registry::Error>> {
                connector_context
                    .azure_device_registry_client
                    .unobserve_asset_update_notifications(
                        asset_ref.device_name.clone(),
                        asset_ref.inbound_endpoint_name.clone(),
                        asset_ref.name.clone(),
                        connector_context.azure_device_registry_timeout,
                    )
                    // retry on network errors, otherwise don't retry on config/dev errors
                    .await
                    .map_err(|e| adr_error_into_retry_error(e, "Unobserve for Asset Updates"))
            },
        )
        .await
        .inspect_err(|e| {
            // If this unobserve is happening because the asset has been deleted, the unobserve call isn't necessary and will likely return
            // an error indicating that the asset can't be found. In the future, we may have a service change to allow us to not need to
            // call unobserve in these cases to clean up our local state, but for now this call is needed and the error should be ignored.
            if is_asset_deleted {
                match e.kind() {
                    azure_device_registry::ErrorKind::ServiceError(_) => {
                        log::debug!("Expected failure unobserving asset update notifications for {asset_ref:?}. Since the asset has been deleted, this is expected and not an error if the error indicates that the asset is not found: {e}");
                    },
                    _ => {
                        log::warn!("Failed to unobserve asset update notifications for deleted {asset_ref:?} after retries. If error indicates that the asset is not found, this is expected and not a true error: {e}");
                    }
                }
            } else {
                log::error!("Failed to unobserve asset update notifications for {asset_ref:?} after retries. If error indicates that the asset is not found, this may be expected and not a true error: {e}");
            }
        });
    }
}

/// Azure Device Registry Asset Component Client represents either a Dataset, Event,
/// Stream, or Management Action Client
#[allow(clippy::large_enum_variant)] // Large variants are mostly composed of already heap allocated data
#[non_exhaustive]
pub enum AssetComponentClient {
    /// Azure Device Registry Data Operation Client represents either a Dataset, Event,
    /// or Stream Client along with any configuration errors detected during creation
    DataOperation((DataOperationClient, Result<(), AdrConfigError>)),
    /// Azure Device Registry Management Action Client represents a Management Action Client
    /// along with either any configuration error detected during creation or the initial
    /// Management Action Executor
    ManagementAction(
        (
            ManagementActionClient,
            Result<ManagementActionExecutor, AdrConfigError>,
        ),
    ),
}

/// Errors that can be returned when reporting a message schema for an asset component
#[derive(Error, Debug)]
pub enum MessageSchemaError {
    /// An error occurred while putting the Schema in the Schema Registry
    #[error(transparent)]
    PutSchemaError(#[from] schema_registry::Error),
    /// An error occurred while reporting the Schema to the Azure Device Registry Service.
    #[error(transparent)]
    AzureDeviceRegistryError(#[from] azure_device_registry::Error),
}

/// Result of a schema modification attempt
pub enum SchemaModifyResult {
    /// Indicates that the schema was reported successfully and status was modified
    Reported(MessageSchemaReference),
    /// Indicates that the schema or status were not modified
    NotModified,
}

trait AssetComponentRef: std::fmt::Debug {
    /// Gets the input needed to modify the status for this asset component
    fn get_modify_input<'a>(
        &self,
        current_status: &'a adr_models::AssetStatus,
    ) -> Option<Result<(), &'a AdrConfigError>>;

    /// Internal function to report the status for this asset component
    async fn internal_report_status(
        &self,
        connector_context: &Arc<ConnectorContext>,
        asset_ref: &AssetRef,
        adr_asset_status: adr_models::AssetStatus,
        asset_status_write_guard: &mut adr_models::AssetStatus,
        desired_asset_component_status: Result<(), AdrConfigError>,
        log_identifier: &str,
    ) -> Result<(), azure_device_registry::Error>;
}

/// A cloneable status reporter for Asset Component status reporting.
///
/// This provides a way to report Asset Component status changes from outside the [`AssetComponentClient`].
/// Each clone maintains its own independent version snapshot, allowing different tasks to work with
/// different versions of the asset specification.
#[allow(private_bounds)]
#[derive(Debug, Clone)]
pub struct AssetComponentStatusReporter<T: AssetComponentRef> {
    connector_context: Arc<ConnectorContext>,
    asset_status: Arc<tokio::sync::RwLock<adr_models::AssetStatus>>,
    asset_specification: Arc<std::sync::RwLock<AssetSpecification>>,
    health_sender: HealthReporterSender,
    asset_component_ref: T,
    asset_ref: AssetRef,
    snapshotted_version: u64,
}

// Allowing private bounds because the customer doesn't actually need to interface with T, and it should be private
#[allow(private_bounds)]
impl<T: AssetComponentRef> AssetComponentStatusReporter<T> {
    /// This function is used to report the current runtime health event for the asset component.
    ///
    /// NOTE: This will only report the health event to the service if needed, so it is
    /// okay to call this function frequently. It is required to call this function any time
    /// the health status changes to ensure that the service has the latest information.
    /// If the health status does not change, the underlying client will report the last known
    /// health event at the base connector `health_report_interval`'s frequency at minimum unless `pause_health_reporting` is called.
    ///
    /// The version used for reporting is the snapshotted version from the last call to
    /// `refresh_health_version()` or `pause_and_refresh_health_version()`, not necessarily the current specification version.
    pub fn report_health_event(&self, health_event: RuntimeHealthEvent) {
        let (status, message, reason_code) = match health_event {
            RuntimeHealthEvent::Available => (HealthStatus::Available, None, None),
            RuntimeHealthEvent::Unavailable {
                message,
                reason_code,
            } => (HealthStatus::Unavailable, message, reason_code),
        };
        let runtime_health = RuntimeHealth {
            last_update_time: Utc::now(),
            message,
            reason_code,
            status,
            version: self.snapshotted_version,
        };
        if let Err(e) = self.health_sender.report(runtime_health) {
            log::warn!(
                "Failed to send health event for {:?}: {e}",
                self.asset_component_ref
            );
        }
    }

    /// Pauses background health event reporting until a new event is reported.
    /// This should be called when the component is updated to indicate that the previous health event may no longer be applicable.
    pub fn pause_health_reporting(&self) {
        if let Err(e) = self.health_sender.pause() {
            log::warn!(
                "Failed to pause health reporting for {:?}: {e}",
                self.asset_component_ref
            );
        }
    }

    /// Snapshots the current asset specification version for use in future health events.
    /// Call this when starting to use a new definition to "lock in" the version you're working with.
    ///
    /// # Panics
    /// if the asset specification mutex has been poisoned, which should not be possible
    pub fn refresh_health_version(&mut self) {
        self.snapshotted_version = self
            .asset_specification
            .read()
            .unwrap()
            .version
            .unwrap_or(0);
    }

    /// Combines `pause_health_reporting()` and `refresh_health_version()`.
    /// Pauses background reporting AND snapshots the current asset specification version.
    /// Useful when receiving an update in the same task that handles notifications.
    pub fn pause_and_refresh_health_version(&mut self) {
        self.pause_health_reporting();
        self.refresh_health_version();
    }

    /// Used to conditionally report the asset component status and then updates the asset with the new status returned.
    ///
    /// The `modify` function is called with the current asset component status (if any) and should return:
    /// - `Some(new_status)` if the status should be updated and reported
    /// - `None` if no update is needed
    ///
    /// # Returns
    /// - [`ModifyResult::Reported`] if the status was updated and successfully reported
    /// - [`ModifyResult::NotModified`] if no modification was needed or the version changed during processing
    ///
    /// # Errors
    /// [`azure_device_registry::Error`] of kind [`AIOProtocolError`](azure_device_registry::ErrorKind::AIOProtocolError) if
    /// there are any underlying errors from the AIO RPC protocol. This error will be retried
    /// 10 times with exponential backoff and jitter and only returned if it still is failing.
    ///
    /// [`azure_device_registry::Error`] of kind [`ServiceError`](azure_device_registry::ErrorKind::ServiceError) if an error is returned
    /// by the Azure Device Registry service.
    ///
    /// # Panics
    /// if the specification mutex has been poisoned, which should not be possible
    pub async fn report_status_if_modified<F>(
        &self,
        modify: F,
    ) -> Result<ModifyResult, azure_device_registry::Error>
    where
        F: Fn(Option<Result<(), &AdrConfigError>>) -> Option<Result<(), AdrConfigError>>,
    {
        // Get the current version of the asset specification
        let cached_version = self.asset_specification.read().unwrap().version;

        {
            // Get the current asset status
            let status_read_guard = self.asset_status.read().await;
            let current_asset_status =
                AssetClient::get_current_asset_status(&status_read_guard, cached_version);

            let modify_input = self
                .asset_component_ref
                .get_modify_input(&current_asset_status);

            // We do not use the result since the status once we acquire the read lock might change
            // and we will have to re-evaluate
            let Some(_modify_result) = modify(modify_input) else {
                // If no modification is needed, return Ok
                return Ok(ModifyResult::NotModified);
            };
        }

        // To modify, we need to acquire the write lock
        let mut status_write_guard = self.asset_status.write().await;

        if cached_version != self.asset_specification.read().unwrap().version {
            // Our modify is no longer valid
            log::debug!(
                "Asset specification is out-of-date for {:?}; skipping modification",
                self.asset_component_ref
            );
            return Ok(ModifyResult::NotModified);
        }

        // We can continue here because the asset status will not change. The specification might
        // update but we will be reporting for an old version.

        // Get the current asset status in case it has changed
        let current_asset_status =
            AssetClient::get_current_asset_status(&status_write_guard, cached_version);

        let modify_result = {
            let modify_input = self
                .asset_component_ref
                .get_modify_input(&current_asset_status);

            let Some(modify_result) = modify(modify_input) else {
                // If no modification is needed, return Ok
                return Ok(ModifyResult::NotModified);
            };

            modify_result
        };

        let mut asset_status_to_report = current_asset_status.into_owned();

        // Update the config status
        asset_status_to_report.config = match asset_status_to_report.config {
            Some(mut config) => {
                config.last_transition_time = Some(Utc::now());
                Some(config)
            }
            None => {
                // If the config is None, we need to create a new one to report along with the
                // asset component's status
                Some(azure_device_registry::ConfigStatus {
                    version: cached_version,
                    last_transition_time: Some(Utc::now()),
                    ..Default::default()
                })
            }
        };

        self.asset_component_ref
            .internal_report_status(
                &self.connector_context,
                &self.asset_ref,
                asset_status_to_report,
                &mut status_write_guard,
                modify_result,
                "AssetComponentStatusReporter::report_asset_component_status_if_modified",
            )
            .await?;

        Ok(ModifyResult::Reported)
    }
}

/// Simplified notification for unsupported components since we don't need to differentiate between update types for them
pub enum UnsupportedComponentNotification {
    /// The component was updated
    Updated,
    /// The component was deleted
    Deleted,
}

/// A trait to support handling all unsupported components in a generic way for the Connector Application
pub trait UnsupportedComponentClient {
    /// The type of the status reporter associated with this unsupported component, which will be used by the Connector Application to report status for this component in a generic way
    #[allow(private_bounds)]
    type StatusReporterType: AssetComponentRef;
    /// Gets a generic status reporter for this unsupported component
    fn get_status_reporter(&self) -> AssetComponentStatusReporter<Self::StatusReporterType>;
    /// Returns whether the notification is Deleted or Updated, since for unsupported components we don't need to differentiate between update types
    fn recv_notification(
        &mut self,
    ) -> impl std::future::Future<Output = UnsupportedComponentNotification> + Send;
}

/// A cloneable status reporter for Data Operation status reporting.
///
/// This provides a way to report Data Operation status changes from outside the [`DataOperationClient`].
pub type DataOperationStatusReporter = AssetComponentStatusReporter<DataOperationRef>;

impl AssetComponentRef for DataOperationRef {
    fn get_modify_input<'a>(
        &self,
        current_status: &'a adr_models::AssetStatus,
    ) -> Option<Result<(), &'a AdrConfigError>> {
        match self.data_operation_name {
            DataOperationName::Dataset {
                name: ref dataset_name,
            } => current_status
                .datasets
                .as_ref()
                .and_then(|datasets| {
                    datasets
                        .iter()
                        .find(|ds_status| ds_status.name == *dataset_name)
                })
                .map(|ds_status| ds_status.error.as_ref().map_or(Ok(()), Err)),
            DataOperationName::Event {
                name: ref event_name,
                ref event_group_name,
            } => current_status
                .event_groups
                .as_ref()
                .and_then(|event_groups| {
                    event_groups
                        .iter()
                        .find(|eg_status| eg_status.name == *event_group_name)
                })
                .and_then(|eg_status| {
                    eg_status.events.as_ref().and_then(|events| {
                        events.iter().find(|e_status| e_status.name == *event_name)
                    })
                })
                .map(|e_status| e_status.error.as_ref().map_or(Ok(()), Err)),
            DataOperationName::Stream {
                name: ref stream_name,
            } => current_status
                .streams
                .as_ref()
                .and_then(|streams| {
                    streams
                        .iter()
                        .find(|s_status| s_status.name == *stream_name)
                })
                .map(|s_status| s_status.error.as_ref().map_or(Ok(()), Err)),
        }
    }

    async fn internal_report_status(
        &self,
        connector_context: &Arc<ConnectorContext>,
        asset_ref: &AssetRef,
        mut adr_asset_status: adr_models::AssetStatus,
        asset_status_write_guard: &mut adr_models::AssetStatus,
        desired_asset_component_status: Result<(), AdrConfigError>,
        log_identifier: &str,
    ) -> Result<(), azure_device_registry::Error> {
        match self.data_operation_name {
            DataOperationName::Dataset {
                name: ref dataset_name,
            } => DataOperationClient::update_dataset_status(
                &mut adr_asset_status,
                dataset_name,
                desired_asset_component_status,
            ),
            DataOperationName::Event {
                name: ref event_name,
                ref event_group_name,
            } => DataOperationClient::update_event_status(
                &mut adr_asset_status,
                event_group_name,
                event_name,
                desired_asset_component_status,
            ),
            DataOperationName::Stream {
                name: ref stream_name,
            } => DataOperationClient::update_stream_status(
                &mut adr_asset_status,
                stream_name,
                desired_asset_component_status,
            ),
        }

        log::debug!("Reporting data operation {self:?} status from app");

        AssetStatusReporter::internal_report_status(
            adr_asset_status,
            connector_context,
            asset_ref,
            asset_status_write_guard,
            log_identifier,
        )
        .await
    }
}

impl UnsupportedComponentClient for DataOperationClient {
    type StatusReporterType = DataOperationRef;
    fn get_status_reporter(&self) -> AssetComponentStatusReporter<Self::StatusReporterType> {
        self.get_status_reporter()
    }

    async fn recv_notification(&mut self) -> UnsupportedComponentNotification {
        match self.recv_notification().await {
            DataOperationNotification::Updated(_) | DataOperationNotification::AssetUpdated(_) => {
                UnsupportedComponentNotification::Updated
            }
            DataOperationNotification::Deleted => UnsupportedComponentNotification::Deleted,
        }
    }
}

#[derive(Debug, Clone)]
pub(crate) struct DataOperationUpdateNotification {
    /// data operation definition from update
    definition: DataOperationDefinition,
    /// default data operation destinations from update
    default_destinations: Vec<Arc<destination_endpoint::Destination>>,
    /// whether default destinations have changed or not in this update
    default_destination_has_changed: bool,
    /// watch receiver for when the update notification should be released to the application
    release_asset_component_notifications_rx: watch::Receiver<()>,
}

/// Notifications that can be received for a Data Operation
pub enum DataOperationNotification {
    /// Indicates that the Asset containing the Data Operation has been updated in place. If this is returned, it indicates that the Data Operation definition has not changed
    /// If there was an error detected, it is included in the result, and must be reported by the application.
    /// If an error is returned, the [`DataOperationClient`] should not be used until there is an Ok update
    AssetUpdated(Result<(), AdrConfigError>),
    /// Indicates that the Data Operation's definition has been updated in place.
    /// This will only be returned if the Data Operation definition changed, which also implies that the Asset has been updated.
    /// If the default destination or default config on the asset changed, the `AssetUpdated` variant will be returned.
    /// If there was an error detected, it is included in the result, and must be reported by the application.
    /// If an error is returned, the [`DataOperationClient`] should not be used until there is an Ok update
    Updated(Result<(), AdrConfigError>),
    /// Indicates that the Data Operation has been deleted.
    Deleted,
}

/// Azure Device Registry Data Operation Client represents either a Dataset, Event,
/// or Stream and includes additional functionality
/// to report status, report message schema, receive updates,
/// and send data to the destination
#[derive(Debug, Getters)]
pub struct DataOperationClient {
    /// Data operation kind and data operation, asset, device, and inbound endpoint names
    data_operation_ref: DataOperationRef,
    // Data operation Definition
    definition: DataOperationDefinition,
    /// Current status for the Asset
    #[getter(skip)]
    asset_status: Arc<tokio::sync::RwLock<adr_models::AssetStatus>>,
    /// Current specification for the Asset
    #[getter(skip)]
    asset_specification: Arc<std::sync::RwLock<AssetSpecification>>,
    /// Specification of the device that this data operation is tied to
    #[getter(skip)]
    device_specification: Arc<std::sync::RwLock<DeviceSpecification>>,
    /// Status of the device that this data operation is tied to
    #[getter(skip)]
    device_status: Arc<tokio::sync::RwLock<DeviceEndpointStatus>>,
    // Internally used fields
    /// Internal [`DataOperationForwarder`] that handles forwarding data to the destination defined in the data operation definition if valid
    #[getter(skip)]
    forwarder: DataOperationForwarder,
    #[getter(skip)]
    connector_context: Arc<ConnectorContext>,
    /// Asset reference for internal use
    #[getter(skip)]
    asset_ref: AssetRef,
    /// Internal watcher receiver that holds a snapshot of the latest update and whether it has been
    /// fully processed or not.
    #[getter(skip)]
    data_operation_update_watcher_rx: watch::Receiver<DataOperationUpdateNotification>,
    /// Health reporter sender for sending health events
    #[getter(skip)]
    health_sender: HealthReporterSender,
    /// Cancellation token for health reporting task - cancelled on deletion
    #[getter(skip)]
    health_cancellation_token: CancellationToken,
}

/// Creates a health reporter sender for a data operation.
fn new_data_operation_health_sender(
    connector_context: &Arc<ConnectorContext>,
    data_operation_ref: &DataOperationRef,
    cancellation_token: CancellationToken,
) -> HealthReporterSender {
    match &data_operation_ref.data_operation_name {
        DataOperationName::Dataset { name } => connector_context
            .azure_device_registry_client
            .new_dataset_health_reporter(
                AssetRef::from(data_operation_ref).into(),
                name.clone(),
                connector_context.azure_device_registry_timeout,
                connector_context.health_report_interval,
                cancellation_token,
            ),
        DataOperationName::Event {
            name,
            event_group_name,
        } => connector_context
            .azure_device_registry_client
            .new_event_health_reporter(
                AssetRef::from(data_operation_ref).into(),
                event_group_name.clone(),
                name.clone(),
                connector_context.azure_device_registry_timeout,
                connector_context.health_report_interval,
                cancellation_token,
            ),
        DataOperationName::Stream { name } => connector_context
            .azure_device_registry_client
            .new_stream_health_reporter(
                AssetRef::from(data_operation_ref).into(),
                name.clone(),
                connector_context.azure_device_registry_timeout,
                connector_context.health_report_interval,
                cancellation_token,
            ),
    }
}

impl DataOperationClient {
    /// Data Operation Client will always be created, even if there are configuration errors
    /// so that updates can be received to fix the configuration errors.
    #[allow(clippy::too_many_arguments)]
    pub(crate) fn new(
        definition: DataOperationDefinition,
        data_operation_update_watcher_rx: watch::Receiver<DataOperationUpdateNotification>,
        default_destinations: &[Arc<destination_endpoint::Destination>],
        asset_ref: AssetRef,
        asset_status: Arc<tokio::sync::RwLock<adr_models::AssetStatus>>,
        asset_specification: Arc<std::sync::RwLock<AssetSpecification>>,
        device_specification: Arc<std::sync::RwLock<DeviceSpecification>>,
        device_status: Arc<tokio::sync::RwLock<DeviceEndpointStatus>>,
        connector_context: Arc<ConnectorContext>,
    ) -> (Self, Result<(), AdrConfigError>) {
        // Create a new data_operation
        let data_operation_ref = DataOperationRef {
            data_operation_name: definition.name(),
            asset_name: asset_ref.name.clone(),
            device_name: asset_ref.device_name.clone(),
            inbound_endpoint_name: asset_ref.inbound_endpoint_name.clone(),
        };
        // technically because these fields are under the lock, if they change, we won't update them unless there's a data operation update
        // however, uuids will always be set and can't change and external ids can only change once from None to set, so this is acceptable
        let (device_uuid, device_external_device_id) = {
            let device_spec = device_specification.read().unwrap();
            (
                device_spec.uuid.clone(),
                device_spec.external_device_id.clone(),
            )
        };
        let (asset_uuid, asset_external_asset_id) = {
            let asset_spec = asset_specification.read().unwrap();
            (
                asset_spec.uuid.clone(),
                asset_spec.external_asset_id.clone(),
            )
        };
        let forwarder_res = match definition {
            DataOperationDefinition::Dataset(ref dataset) => {
                destination_endpoint::Forwarder::new_dataset_forwarder(
                    dataset,
                    default_destinations,
                    &asset_ref,
                    device_uuid,
                    device_external_device_id,
                    asset_uuid.as_ref(),
                    asset_external_asset_id.as_ref(),
                    connector_context.clone(),
                )
            }
            DataOperationDefinition::Event(ref event) => {
                destination_endpoint::Forwarder::new_event_stream_forwarder(
                    &event.destinations,
                    default_destinations,
                    &data_operation_ref,
                    event.data_source.clone(),
                    event.type_ref.clone(),
                    device_uuid,
                    device_external_device_id,
                    asset_uuid.as_ref(),
                    asset_external_asset_id.as_ref(),
                    connector_context.clone(),
                )
            }
            DataOperationDefinition::Stream(ref stream) => {
                destination_endpoint::Forwarder::new_event_stream_forwarder(
                    &stream.destinations,
                    default_destinations,
                    &data_operation_ref,
                    None,
                    stream.type_ref.clone(),
                    device_uuid,
                    device_external_device_id,
                    asset_uuid.as_ref(),
                    asset_external_asset_id.as_ref(),
                    connector_context.clone(),
                )
            }
        };
        let health_cancellation_token = CancellationToken::new();
        let health_sender = new_data_operation_health_sender(
            &connector_context,
            &data_operation_ref,
            health_cancellation_token.clone(),
        );
        let (forwarder, res) = match forwarder_res {
            Ok(forwarder) => (
                DataOperationForwarder::Forwarder(Box::new(forwarder)),
                Ok(()),
            ),
            Err(e) => {
                log::warn!("Invalid destination for data_operation: {data_operation_ref:?} {e:?}");
                (DataOperationForwarder::Error(e.clone()), Err(e))
            }
        };
        (
            Self {
                data_operation_ref,
                definition,
                asset_status,
                asset_specification,
                device_specification,
                device_status,
                forwarder,
                connector_context,
                asset_ref,
                data_operation_update_watcher_rx,
                health_sender,
                health_cancellation_token,
            },
            res,
        )
    }

    /// Returns the kind of data operation this client represents
    #[must_use]
    pub fn kind(&self) -> DataOperationKind {
        self.definition.kind()
    }

    /// Used to conditionally report the message schema of a data operation as an existing schema reference
    ///
    /// The `modify` function is called with the current message schema reference (if any) and should return:
    /// - `Some(new_message_schema_reference)` if the schema should be updated and reported
    /// - `None` if no update is needed
    ///
    /// # Returns
    /// - [`SchemaModifyResult::Reported`] if the schema was updated and successfully reported, containing the reported [`MessageSchemaReference`]
    /// - [`SchemaModifyResult::NotModified`] if no modification was needed or the version changed during processing
    ///
    /// # Errors
    /// [`MessageSchemaError`] of kind [`AzureDeviceRegistryError::AIOProtocolError`](azure_device_registry::ErrorKind::AIOProtocolError) if
    /// there are any underlying errors from the AIO RPC protocol. This error will be retried
    /// 10 times with exponential backoff and jitter and only returned if it still is failing.
    ///
    /// [`MessageSchemaError`] of kind [`AzureDeviceRegistryError::ServiceError`](azure_device_registry::ErrorKind::ServiceError) if an error is returned
    /// by the Azure Device Registry service.
    ///
    /// # Panics
    /// If the asset specification mutex has been poisoned, which should not be possible
    pub async fn report_message_schema_reference_if_modified<F>(
        &mut self,
        modify: F,
    ) -> Result<SchemaModifyResult, MessageSchemaError>
    where
        F: Fn(Option<&MessageSchemaReference>) -> Option<MessageSchemaReference>,
    {
        // Get the current version of the asset specification
        let cached_version = self.asset_specification.read().unwrap().version;

        {
            // Get the current asset status
            let status_read_guard = self.asset_status.read().await;
            let current_asset_status =
                AssetClient::get_current_asset_status(&status_read_guard, cached_version);

            let modify_input = self.get_schema_reference_modify_input(&current_asset_status);

            match modify(modify_input) {
                Some(_) => {
                    // A modification was made, we proceed to report schema
                }
                None => {
                    // No modification was made, so no need to report schema
                    return Ok(SchemaModifyResult::NotModified);
                }
            }
        };

        // To modify, we need to acquire the write lock
        let mut status_write_guard = self.asset_status.write().await;

        // We can continue here because the asset status will not change. The specification might
        // update but we will be reporting for an old version.

        if cached_version != self.asset_specification.read().unwrap().version {
            // Our modify is no longer valid
            log::debug!(
                "Reporting for an out-of-date asset specification from a Data Operation client, will not modify"
            );
            return Ok(SchemaModifyResult::NotModified);
        }

        // Get the current asset status in case it has changed
        let current_asset_status =
            AssetClient::get_current_asset_status(&status_write_guard, cached_version);

        let modify_input = self.get_schema_reference_modify_input(&current_asset_status);

        let Some(new_message_schema_reference) = modify(modify_input) else {
            // No modification was made, so no need to report schema
            return Ok(SchemaModifyResult::NotModified);
        };

        let mut asset_status_to_report = current_asset_status.into_owned();

        asset_status_to_report.config = match asset_status_to_report.config {
            Some(mut config) => {
                config.last_transition_time = Some(Utc::now());
                Some(config)
            }
            None => {
                // If the config is None, we need to create a new one to report along with the asset status
                Some(azure_device_registry::ConfigStatus {
                    version: cached_version,
                    last_transition_time: Some(Utc::now()),
                    ..Default::default()
                })
            }
        };

        Self::internal_report_message_schema_reference(
            &self.connector_context,
            &self.asset_ref,
            &self.data_operation_ref,
            &mut self.forwarder,
            asset_status_to_report,
            &mut status_write_guard,
            &new_message_schema_reference,
            "DataOperationClient::report_message_schema_reference_if_modified",
        )
        .await?;

        Ok(SchemaModifyResult::Reported(new_message_schema_reference))
    }

    /// Used to conditionally report the message schema of a data operation
    ///
    /// The `modify` function is called with the current message schema reference (if any) and should return:
    /// - `Some(new_message_schema)` if the schema should be updated and reported
    /// - `None` if no update is needed
    ///
    /// # Returns
    /// - [`SchemaModifyResult::Reported`] if the schema was updated and successfully reported, containing the reported [`MessageSchemaReference`]
    /// - [`SchemaModifyResult::NotModified`] if no modification was needed or the version changed during processing
    ///
    /// # Errors
    /// [`MessageSchemaError`] of kind [`SchemaRegistryError::InvalidRequestArgument`](schema_registry::ErrorKind::InvalidRequestArgument)
    /// if the content of the [`MessageSchema`] is empty or there is an error building the request
    ///
    /// [`MessageSchemaError`] of kind [`SchemaRegistryError::ServiceError`](schema_registry::ErrorKind::ServiceError)
    /// if there is an error returned by the Schema Registry Service. This error will be retried 10
    /// times with exponential backoff and jitter if it is an internal error and only returned if
    /// it still is failing.
    ///
    /// [`MessageSchemaError`] of kind [`AzureDeviceRegistryError::AIOProtocolError`](azure_device_registry::ErrorKind::AIOProtocolError) if
    /// there are any underlying errors from the AIO RPC protocol. This error will be retried
    /// 10 times with exponential backoff and jitter and only returned if it still is failing.
    ///
    /// [`MessageSchemaError`] of kind [`AzureDeviceRegistryError::ServiceError`](azure_device_registry::ErrorKind::ServiceError) if
    /// an error is returned by the Azure Device Registry service.
    ///
    /// # Panics
    /// If the asset specification mutex has been poisoned, which should not be possible
    pub async fn report_message_schema_if_modified<F>(
        &mut self,
        modify: F,
    ) -> Result<SchemaModifyResult, MessageSchemaError>
    where
        F: Fn(Option<&MessageSchemaReference>) -> Option<MessageSchema>,
    {
        // Get the current version of the asset specification
        let cached_version = self.asset_specification.read().unwrap().version;

        {
            // Get the current asset status
            let status_read_guard = self.asset_status.read().await;
            let current_asset_status =
                AssetClient::get_current_asset_status(&status_read_guard, cached_version);

            let modify_input = self.get_schema_reference_modify_input(&current_asset_status);

            if modify(modify_input).is_some() {
                // A modification was made, we proceed to report schema
            } else {
                // No modification was made, so no need to report schema
                return Ok(SchemaModifyResult::NotModified);
            }
        };

        // To modify, we need to acquire the write lock
        let mut status_write_guard = self.asset_status.write().await;

        // We can continue here because the asset status will not change. The specification might
        // update but we will be reporting for an old version.

        if cached_version != self.asset_specification.read().unwrap().version {
            // Our modify is no longer valid
            log::debug!(
                "Reporting for an out-of-date asset specification from a Data Operation client, will not modify"
            );
            return Ok(SchemaModifyResult::NotModified);
        }

        // Get the current asset status in case it has changed
        let current_asset_status =
            AssetClient::get_current_asset_status(&status_write_guard, cached_version);

        let modify_input = self.get_schema_reference_modify_input(&current_asset_status);

        let Some(new_message_schema) = modify(modify_input) else {
            // No modification was made, so no need to report schema
            return Ok(SchemaModifyResult::NotModified);
        };

        let mut asset_status_to_report = current_asset_status.into_owned();

        asset_status_to_report.config = match asset_status_to_report.config {
            Some(mut config) => {
                config.last_transition_time = Some(Utc::now());
                Some(config)
            }
            None => {
                // If the config is None, we need to create a new one to report along with the asset status
                Some(azure_device_registry::ConfigStatus {
                    version: cached_version,
                    last_transition_time: Some(Utc::now()),
                    ..Default::default()
                })
            }
        };

        // First put the schema in the schema registry
        let message_schema_reference = Retry::spawn(
            RETRY_STRATEGY.map(tokio_retry2::strategy::jitter),
            async || -> Result<schema_registry::Schema, RetryError<schema_registry::Error>> {
                self.connector_context
                    .schema_registry_client
                    .put(
                        new_message_schema.clone(),
                        self.connector_context.schema_registry_timeout,
                    )
                    .await
                    .map_err(|e| {
                        match e.kind() {
                            // network/retriable
                            schema_registry::ErrorKind::AIOProtocolError(_) => {
                                log::warn!(
                                    "Reporting message schema failed for {:?}. Retrying: {e}",
                                    self.data_operation_ref
                                );
                                RetryError::transient(e)
                            }
                            schema_registry::ErrorKind::ServiceError(service_error) => {
                                if let schema_registry::ErrorCode::InternalError =
                                    service_error.code
                                {
                                    log::warn!(
                                        "Reporting message schema failed for {:?}. Retrying: {e}",
                                        self.data_operation_ref
                                    );
                                    RetryError::transient(e)
                                } else {
                                    RetryError::permanent(e)
                                }
                            }
                            // indicates an error in the provided message schema, return to caller so they can fix
                            schema_registry::ErrorKind::InvalidRequestArgument(_) => {
                                RetryError::permanent(e)
                            }
                        }
                    })
            },
        )
        .await
        .map(|schema| MessageSchemaReference {
            name: schema.name,
            version: schema.version,
            registry_namespace: schema.namespace,
        })?;

        Self::internal_report_message_schema_reference(
            &self.connector_context,
            &self.asset_ref,
            &self.data_operation_ref,
            &mut self.forwarder,
            asset_status_to_report,
            &mut status_write_guard,
            &message_schema_reference,
            "DataOperationClient::report_message_schema_if_modified",
        )
        .await?;

        Ok(SchemaModifyResult::Reported(message_schema_reference))
    }

    #[allow(clippy::too_many_arguments)]
    async fn internal_report_message_schema_reference(
        connector_context: &Arc<ConnectorContext>,
        asset_ref: &AssetRef,
        data_operation_ref: &DataOperationRef,
        forwarder: &mut DataOperationForwarder,
        asset_status_to_report: adr_models::AssetStatus,
        status_write_guard: &mut adr_models::AssetStatus,
        message_schema_reference: &MessageSchemaReference,
        log_identifier: &str,
    ) -> Result<(), MessageSchemaError> {
        // Use the provided asset_status_to_report instead of creating a new one
        let mut new_status = asset_status_to_report;

        // if data operation is already in the current status, then update the existing data operation with the new message schema
        // Otherwise if the data operation isn't present, or no data operations of that kind have been reported yet, then add it with the new message schema
        match data_operation_ref.data_operation_name {
            DataOperationName::Dataset {
                name: ref dataset_name,
            } => {
                if let Some(dataset_status) = new_status.datasets.as_mut().and_then(|datasets| {
                    datasets
                        .iter_mut()
                        .find(|dataset| dataset.name == *dataset_name)
                }) {
                    // If the dataset already has a status, update the existing dataset with the new message schema
                    dataset_status.message_schema_reference =
                        Some(message_schema_reference.clone());
                } else {
                    // If the dataset doesn't exist in the current status, then add it
                    new_status.datasets.get_or_insert_with(Vec::new).push(
                        adr_models::DatasetEventStreamStatus {
                            name: dataset_name.clone(),
                            message_schema_reference: Some(message_schema_reference.clone()),
                            error: None,
                        },
                    );
                }
            }
            DataOperationName::Event {
                name: ref event_name,
                ref event_group_name,
            } => {
                if let Some(event_group_status) =
                    new_status.event_groups.as_mut().and_then(|event_groups| {
                        event_groups
                            .iter_mut()
                            .find(|event_group| event_group.name == *event_group_name)
                    })
                {
                    if let Some(event_status) =
                        event_group_status.events.as_mut().and_then(|events| {
                            events.iter_mut().find(|event| event.name == *event_name)
                        })
                    {
                        // If the event already has a status, update the existing event with the new message schema
                        event_status.message_schema_reference =
                            Some(message_schema_reference.clone());
                    } else {
                        // If the event doesn't exist in the current status, then add it
                        event_group_status.events.get_or_insert_with(Vec::new).push(
                            adr_models::DatasetEventStreamStatus {
                                name: event_name.clone(),
                                message_schema_reference: Some(message_schema_reference.clone()),
                                error: None,
                            },
                        );
                    }
                } else {
                    // If the event group doesn't exist in the current status, then add it
                    new_status.event_groups.get_or_insert_with(Vec::new).push(
                        adr_models::EventGroupStatus {
                            name: event_group_name.clone(),
                            events: Some(vec![adr_models::DatasetEventStreamStatus {
                                name: event_name.clone(),
                                message_schema_reference: Some(message_schema_reference.clone()),
                                error: None,
                            }]),
                        },
                    );
                }
            }
            DataOperationName::Stream {
                name: ref stream_name,
            } => {
                if let Some(stream_status) = new_status.streams.as_mut().and_then(|streams| {
                    streams
                        .iter_mut()
                        .find(|stream| stream.name == *stream_name)
                }) {
                    // If the stream already has a status, update the existing stream with the new message schema
                    stream_status.message_schema_reference = Some(message_schema_reference.clone());
                } else {
                    // If the stream doesn't exist in the current status, then add it
                    new_status.streams.get_or_insert_with(Vec::new).push(
                        adr_models::DatasetEventStreamStatus {
                            name: stream_name.clone(),
                            message_schema_reference: Some(message_schema_reference.clone()),
                            error: None,
                        },
                    );
                }
            }
        }

        // send status update to the service
        log::debug!("Reporting data operation {data_operation_ref:?} message schema from app");
        AssetStatusReporter::internal_report_status(
            new_status,
            connector_context,
            asset_ref,
            status_write_guard,
            log_identifier,
        )
        .await?;

        // CONSIDERATION: should we do this at the beginning and not allow the message schema reporting if there's no valid destination?
        // For now, don't block reporting the message schema if there's no valid destination for more flexibility
        match forwarder {
            DataOperationForwarder::Forwarder(forwarder) => {
                forwarder.update_message_schema_reference(Some(message_schema_reference.clone()));
            }
            DataOperationForwarder::Error(_) => {}
        }

        Ok(())
    }

    fn get_schema_reference_modify_input<'a>(
        &self,
        asset_status: &'a adr_models::AssetStatus,
    ) -> Option<&'a adr_models::MessageSchemaReference> {
        match self.data_operation_ref.data_operation_name {
            DataOperationName::Dataset {
                name: ref dataset_name,
            } => asset_status
                .datasets
                .as_ref()
                .and_then(|datasets| {
                    datasets
                        .iter()
                        .find(|ds_status| ds_status.name == *dataset_name)
                })
                .and_then(|ds_status| ds_status.message_schema_reference.as_ref()),
            DataOperationName::Event {
                name: ref event_name,
                ref event_group_name,
            } => asset_status
                .event_groups
                .as_ref()
                .and_then(|event_groups| {
                    event_groups
                        .iter()
                        .find(|eg_status| eg_status.name == *event_group_name)
                })
                .and_then(|eg_status| {
                    eg_status.events.as_ref().and_then(|events| {
                        events.iter().find(|e_status| e_status.name == *event_name)
                    })
                })
                .and_then(|e_status| e_status.message_schema_reference.as_ref()),
            DataOperationName::Stream {
                name: ref stream_name,
            } => asset_status
                .streams
                .as_ref()
                .and_then(|streams| {
                    streams
                        .iter()
                        .find(|s_status| s_status.name == *stream_name)
                })
                .and_then(|s_status| s_status.message_schema_reference.as_ref()),
        }
    }

    /// Used to send transformed data to the destination
    /// Returns once the message has been sent successfully
    /// Use `forward_data_provide_protocol_specific_identifier` if it is desired to
    /// provide a Protocol Specific Identifier to be used on the Cloud Event `source`
    /// header used if the destination is `MQTT`. If this fn is used, the Cloud Event Header
    /// will default to using either the device external device id or the device name.
    ///
    /// # Errors
    /// [`destination_endpoint::Error`] of kind [`MissingMessageSchema`](destination_endpoint::ErrorKind::MissingMessageSchema)
    /// if the [`MessageSchema`] has not been reported yet. This is required before forwarding any data
    ///
    /// [`destination_endpoint::Error`] of kind [`ValidationError`](destination_endpoint::ErrorKind::ValidationError)
    /// if the [`Data`] isn't valid or there isn't a valid destination configured for the data operation.
    ///
    /// [`destination_endpoint::Error`] of kind [`BrokerStateStoreError`](destination_endpoint::ErrorKind::BrokerStateStoreError)
    /// if the destination is `BrokerStateStore` and there are any errors setting the data with the service
    ///
    /// [`destination_endpoint::Error`] of kind [`MqttTelemetryError`](destination_endpoint::ErrorKind::MqttTelemetryError)
    /// if the destination is `Mqtt` and there are any errors sending the message to the broker
    pub async fn forward_data(&self, data: Data) -> Result<(), destination_endpoint::Error> {
        self.forwarder.send_data(data, None).await
    }

    /// Used to send transformed data to the destination
    /// Returns once the message has been sent successfully.
    /// `protocol_specific_identifier` will be used on the Cloud Event
    /// `source` header used if the destination is `MQTT`. If `forward_data` is used instead of this fn,
    /// the Cloud Event Header will default to using either the device external device id or the device name.
    /// The inbound endpoint address is a recommended value for this field.
    ///
    /// # Errors
    /// [`destination_endpoint::Error`] of kind [`MissingMessageSchema`](destination_endpoint::ErrorKind::MissingMessageSchema)
    /// if the [`MessageSchema`] has not been reported yet. This is required before forwarding any data
    ///
    /// [`destination_endpoint::Error`] of kind [`ValidationError`](destination_endpoint::ErrorKind::ValidationError)
    /// if the [`Data`] isn't valid or there isn't a valid destination configured for the data operation.
    ///
    /// [`destination_endpoint::Error`] of kind [`BrokerStateStoreError`](destination_endpoint::ErrorKind::BrokerStateStoreError)
    /// if the destination is `BrokerStateStore` and there are any errors setting the data with the service
    ///
    /// [`destination_endpoint::Error`] of kind [`MqttTelemetryError`](destination_endpoint::ErrorKind::MqttTelemetryError)
    /// if the destination is `Mqtt` and there are any errors sending the message to the broker
    pub async fn forward_data_provide_protocol_specific_identifier(
        &self,
        data: Data,
        protocol_specific_identifier: &str,
    ) -> Result<(), destination_endpoint::Error> {
        self.forwarder
            .send_data(data, Some(protocol_specific_identifier))
            .await
    }

    /// Used to receive notifications about the Data Operation from the Azure Device Registry Service.
    ///
    /// Returns [`DataOperationNotification::DataOperationUpdated`] if the Data Operation's definition has been updated in place.
    ///
    /// Returns [`DataOperationNotification::AssetUpdated`] if the Asset's definition has been updated in place, but there were no
    /// changes to the Data Operation's definition.
    ///
    /// Returns either of the above with an `Err(AdrConfigError)` if the Data Operation/Asset received an update, but a config error
    /// for the Data Operation was detected. The definition is still updated in place, but the [`DataOperationClient`] should not be
    /// used until there is a new update because forwarding the data will not be possible.
    ///
    /// Returns [`DataOperationNotification::Deleted`] if the Data Operation has been deleted. The [`DataOperationClient`]
    /// should not be used after this point, and no more notifications will be received.
    ///
    /// # Panics
    /// If the asset specification mutex has been poisoned, which should not be possible
    ///
    /// # Cancel safety
    /// This method is cancel safe. If you use it as the event in a `tokio::select!` statement and some other branch
    /// completes first, then it is guaranteed that no data operation notifications will be lost, and the data operation will not
    /// be updated without a notification being returned.
    pub async fn recv_notification(&mut self) -> DataOperationNotification {
        if self
            .data_operation_update_watcher_rx
            .changed()
            .await
            .is_err()
        {
            // Cancel health reporting task on deletion
            self.health_cancellation_token.cancel();
            return DataOperationNotification::Deleted;
        }
        // In case this function gets cancelled the next time it is called we will process the update again.
        self.data_operation_update_watcher_rx.mark_changed();
        let mut update_notification = self.data_operation_update_watcher_rx.borrow().clone();

        // wait until the update has been released. If the watch sender has been dropped, this means the Asset has been deleted/dropped
        if update_notification
            .release_asset_component_notifications_rx
            .changed()
            .await
            .is_err()
        {
            self.data_operation_update_watcher_rx.mark_unchanged();
            // Cancel health reporting task on deletion
            self.health_cancellation_token.cancel();
            return DataOperationNotification::Deleted;
        }
        let data_operation_changed = update_notification.definition != self.definition;
        if !data_operation_changed && !update_notification.default_destination_has_changed {
            // Asset only update, we don't need to recreate the forwarder
            // Once the data_operation update has been processed we can mark the value in the watcher as seen
            self.data_operation_update_watcher_rx.mark_unchanged();
            // if nothing changed, but there was a destination error before, we need to return that error again since it still applies
            let res = match &self.forwarder {
                DataOperationForwarder::Forwarder(_) => Ok(()),
                DataOperationForwarder::Error(e) => Err(e.clone()),
            };
            return DataOperationNotification::AssetUpdated(res);
        }

        // no await points beyond this point, so this is safe
        self.definition = update_notification.definition;

        // create new forwarder, in case destination has changed
        // technically because these fields are under the lock, if they change, we won't update them unless there's a data operation update
        // however, uuids will always be set and can't change and external ids can only change once from None to set, so this is acceptable
        let (device_uuid, device_external_device_id) = {
            let device_spec = self.device_specification.read().unwrap();
            (
                device_spec.uuid.clone(),
                device_spec.external_device_id.clone(),
            )
        };
        let (asset_uuid, asset_external_asset_id) = {
            let asset_spec = self.asset_specification.read().unwrap();
            (
                asset_spec.uuid.clone(),
                asset_spec.external_asset_id.clone(),
            )
        };
        let forwarder_result = match self.definition {
            DataOperationDefinition::Dataset(ref updated_dataset) => {
                destination_endpoint::Forwarder::new_dataset_forwarder(
                    updated_dataset,
                    &update_notification.default_destinations,
                    &self.asset_ref,
                    device_uuid,
                    device_external_device_id,
                    asset_uuid.as_ref(),
                    asset_external_asset_id.as_ref(),
                    self.connector_context.clone(),
                )
            }
            DataOperationDefinition::Event(ref updated_event) => {
                destination_endpoint::Forwarder::new_event_stream_forwarder(
                    &updated_event.destinations,
                    &update_notification.default_destinations,
                    &self.data_operation_ref,
                    updated_event.data_source.clone(),
                    updated_event.type_ref.clone(),
                    device_uuid,
                    device_external_device_id,
                    asset_uuid.as_ref(),
                    asset_external_asset_id.as_ref(),
                    self.connector_context.clone(),
                )
            }
            DataOperationDefinition::Stream(ref updated_stream) => {
                destination_endpoint::Forwarder::new_event_stream_forwarder(
                    &updated_stream.destinations,
                    &update_notification.default_destinations,
                    &self.data_operation_ref,
                    None,
                    updated_stream.type_ref.clone(),
                    device_uuid,
                    device_external_device_id,
                    asset_uuid.as_ref(),
                    asset_external_asset_id.as_ref(),
                    self.connector_context.clone(),
                )
            }
        };
        self.forwarder = match forwarder_result {
            Ok(forwarder) => DataOperationForwarder::Forwarder(Box::new(forwarder)),
            Err(e) => {
                log::warn!(
                    "Invalid data_operation destination for updated data_operation: {:?} {e:?}",
                    self.data_operation_ref
                );
                DataOperationForwarder::Error(e)
            }
        };
        // Once the data_operation definition has been updated we can mark the value in the watcher as seen
        self.data_operation_update_watcher_rx.mark_unchanged();
        let res = match &self.forwarder {
            DataOperationForwarder::Forwarder(_) => Ok(()),
            DataOperationForwarder::Error(e) => Err(e.clone()),
        };
        if data_operation_changed {
            DataOperationNotification::Updated(res)
        } else {
            DataOperationNotification::AssetUpdated(res)
        }
    }

    /// Creates a new status reporter for this [`DataOperationClient`]
    /// The reporter's version snapshot is initialized to the current asset specification version.
    ///
    /// # Panics
    /// if the asset specification mutex has been poisoned, which should not be possible
    #[must_use]
    pub fn get_status_reporter(&self) -> DataOperationStatusReporter {
        let snapshotted_version = self
            .asset_specification
            .read()
            .unwrap()
            .version
            .unwrap_or(0);
        DataOperationStatusReporter {
            connector_context: self.connector_context.clone(),
            asset_status: self.asset_status.clone(),
            asset_specification: self.asset_specification.clone(),
            health_sender: self.health_sender.clone(),
            asset_component_ref: self.data_operation_ref.clone(),
            asset_ref: self.asset_ref.clone(),
            snapshotted_version,
        }
    }

    /// Returns a clone of this Data Operation's [`MessageSchemaReference`] from
    /// the `AssetStatus`, if it exists
    #[must_use]
    pub async fn message_schema_reference(&self) -> Option<MessageSchemaReference> {
        // unwrap can't fail unless lock is poisoned
        match self.data_operation_ref.data_operation_name {
            DataOperationName::Dataset {
                name: ref dataset_name,
            } => self
                .asset_status
                .read()
                .await
                .datasets
                .as_ref()?
                .iter()
                .find(|dataset| dataset.name == *dataset_name)?
                .message_schema_reference
                .clone(),
            DataOperationName::Event {
                name: ref event_name,
                ref event_group_name,
            } => self
                .asset_status
                .read()
                .await
                .event_groups
                .as_ref()?
                .iter()
                .find(|event_group| event_group.name == *event_group_name)?
                .events
                .as_ref()?
                .iter()
                .find(|event| event.name == *event_name)?
                .message_schema_reference
                .clone(),
            DataOperationName::Stream {
                name: ref stream_name,
            } => self
                .asset_status
                .read()
                .await
                .streams
                .as_ref()?
                .iter()
                .find(|stream| stream.name == *stream_name)?
                .message_schema_reference
                .clone(),
        }
    }

    /// Returns a clone of the current asset specification
    /// # Panics
    /// if the asset specification mutex has been poisoned, which should not be possible
    #[must_use]
    pub fn asset_specification(&self) -> AssetSpecification {
        (*self.asset_specification.read().unwrap()).clone()
    }

    /// Returns a clone of the current asset status, if it exists
    #[must_use]
    pub async fn asset_status(&self) -> adr_models::AssetStatus {
        (*self.asset_status.read().await).clone()
    }

    /// Returns a clone of the current device specification
    /// # Panics
    /// if the device specification mutex has been poisoned, which should not be possible
    #[must_use]
    pub fn device_specification(&self) -> DeviceSpecification {
        (*self.device_specification.read().unwrap()).clone()
    }

    /// Returns a clone of the current device status
    #[must_use]
    pub async fn device_status(&self) -> DeviceEndpointStatus {
        (*self.device_status.read().await).clone()
    }

    /// Helper function to update the specific dataset status within the asset status
    fn update_dataset_status(
        asset_status_to_update: &mut adr_models::AssetStatus,
        dataset_name: &str,
        dataset_status: Result<(), AdrConfigError>,
    ) {
        if let Some(curr_dataset_status) =
            asset_status_to_update
                .datasets
                .as_mut()
                .and_then(|datasets| {
                    datasets
                        .iter_mut()
                        .find(|dataset| dataset.name == dataset_name)
                })
        {
            // If the dataset already has a status, update the existing dataset with the new error
            curr_dataset_status.error = dataset_status.err();
        } else {
            // If the dataset doesn't exist in the current status, then add it
            asset_status_to_update
                .datasets
                .get_or_insert_with(Vec::new)
                .push(adr_models::DatasetEventStreamStatus {
                    name: dataset_name.to_string(),
                    message_schema_reference: None,
                    error: dataset_status.err(),
                });
        }
    }

    /// Helper function to update the specific event status within the asset status
    fn update_event_status(
        asset_status_to_update: &mut adr_models::AssetStatus,
        event_group_name: &str,
        event_name: &str,
        event_status: Result<(), AdrConfigError>,
    ) {
        if let Some(curr_event_group_status) = asset_status_to_update
            .event_groups
            .as_mut()
            .and_then(|event_groups| {
                event_groups
                    .iter_mut()
                    .find(|event_group| event_group.name == event_group_name)
            })
        {
            if let Some(curr_event_status) = curr_event_group_status
                .events
                .as_mut()
                .and_then(|events| events.iter_mut().find(|event| event.name == event_name))
            {
                // If the event already has a status, update the existing event with the new error
                curr_event_status.error = event_status.err();
            } else {
                // If the event doesn't exist in the current event group status, then add it
                curr_event_group_status
                    .events
                    .get_or_insert_with(Vec::new)
                    .push(adr_models::DatasetEventStreamStatus {
                        name: event_name.to_string(),
                        message_schema_reference: None,
                        error: event_status.err(),
                    });
            }
        } else {
            // If the event group doesn't exist in the current status, then add it
            asset_status_to_update
                .event_groups
                .get_or_insert_with(Vec::new)
                .push(adr_models::EventGroupStatus {
                    name: event_group_name.to_string(),
                    events: Some(vec![adr_models::DatasetEventStreamStatus {
                        name: event_name.to_string(),
                        message_schema_reference: None,
                        error: event_status.err(),
                    }]),
                });
        }
    }

    /// Helper function to update the specific stream status within the asset status
    fn update_stream_status(
        asset_status_to_update: &mut adr_models::AssetStatus,
        stream_name: &str,
        stream_status: Result<(), AdrConfigError>,
    ) {
        if let Some(curr_stream_status) = asset_status_to_update
            .streams
            .as_mut()
            .and_then(|streams| streams.iter_mut().find(|stream| stream.name == stream_name))
        {
            // If the stream already has a status, update the existing stream with the new error
            curr_stream_status.error = stream_status.err();
        } else {
            // If the stream doesn't exist in the current status, then add it
            asset_status_to_update
                .streams
                .get_or_insert_with(Vec::new)
                .push(adr_models::DatasetEventStreamStatus {
                    name: stream_name.to_string(),
                    message_schema_reference: None,
                    error: stream_status.err(),
                });
        }
    }
}

impl Drop for DataOperationClient {
    fn drop(&mut self) {
        self.health_cancellation_token.cancel();
    }
}

/// A cloneable status reporter for Management Action status reporting.
///
/// This provides a way to report Management Action status changes from outside the [`ManagementActionClient`].
pub type ManagementActionStatusReporter = AssetComponentStatusReporter<ManagementActionRef>;

impl AssetComponentRef for ManagementActionRef {
    fn get_modify_input<'a>(
        &self,
        current_status: &'a adr_models::AssetStatus,
    ) -> Option<Result<(), &'a AdrConfigError>> {
        current_status
            .management_groups
            .as_ref()
            .and_then(|management_groups| {
                management_groups
                    .iter()
                    .find(|mg_status| mg_status.name == *self.management_group_name)
            })
            .and_then(|mg_status| {
                mg_status.actions.as_ref().and_then(|actions| {
                    actions
                        .iter()
                        .find(|a_status| a_status.name == *self.management_action_name)
                })
            })
            .map(|a_status| a_status.error.as_ref().map_or(Ok(()), Err))
    }

    async fn internal_report_status(
        &self,
        connector_context: &Arc<ConnectorContext>,
        asset_ref: &AssetRef,
        mut adr_asset_status: adr_models::AssetStatus,
        asset_status_write_guard: &mut adr_models::AssetStatus,
        desired_asset_component_status: Result<(), AdrConfigError>,
        log_identifier: &str,
    ) -> Result<(), azure_device_registry::Error> {
        ManagementActionClient::update_action_status(
            &mut adr_asset_status,
            &self.management_group_name,
            &self.management_action_name,
            desired_asset_component_status,
        );

        log::debug!("Reporting management action {self:?} status from app");

        AssetStatusReporter::internal_report_status(
            adr_asset_status,
            connector_context,
            asset_ref,
            asset_status_write_guard,
            log_identifier,
        )
        .await
    }
}

impl UnsupportedComponentClient for ManagementActionClient {
    type StatusReporterType = ManagementActionRef;
    fn get_status_reporter(&self) -> AssetComponentStatusReporter<Self::StatusReporterType> {
        self.get_status_reporter()
    }

    async fn recv_notification(&mut self) -> UnsupportedComponentNotification {
        match self.recv_notification().await {
            ManagementActionNotification::Updated(_)
            | ManagementActionNotification::UpdatedWithNewExecutor(_)
            | ManagementActionNotification::AssetUpdated(_) => {
                UnsupportedComponentNotification::Updated
            }
            ManagementActionNotification::Deleted => UnsupportedComponentNotification::Deleted,
        }
    }
}

#[derive(Debug, Clone)]
pub(crate) struct ManagementActionUpdateNotification {
    /// management action definition from update
    definition: ManagementActionSpecification,
    /// watch receiver for when the update notification should be released to the application
    release_asset_component_notifications_rx: watch::Receiver<()>,
}

/// Notifications that can be received for a Management Action
#[allow(clippy::large_enum_variant)] // Large variants are mostly composed of already heap allocated data
pub enum ManagementActionNotification {
    /// Indicates that the Asset containing the Management Action has been updated in place. If this is returned,
    /// it indicates that the Management Group/Action definition has not changed.
    /// If there was an error detected, it is included in the result, and must be reported by the application.
    /// If an error is returned, the [`ManagementActionClient`] should not be used until there is an Ok update
    AssetUpdated(Result<(), AdrConfigError>),
    /// Indicates that the Management Action's definition has been updated in place.
    /// This will only be returned if the Management Group or Action definition has changed,
    /// which implies that the Asset has also been updated.
    /// If the default config on the asset changed, the `AssetUpdated` variant will be returned.
    /// If there was an error detected, it is included in the result, and must be reported by the application.
    /// If an error is returned, the [`ManagementActionClient`] should not be used until there is an Ok update
    Updated(Result<(), AdrConfigError>),
    /// Indicates that the Management Action's definition has been updated in place and the old executor no longer
    /// can be used with the new definition. This will only be returned if the Management Group or Action
    /// definition has changed, and that change requires a new executor to be created (e.g. topic change).
    /// This changing also implies that the Asset has been updated.
    /// If the topic for the executor has not changed, the `Updated` variant will be returned.
    /// If there was an error detected, it is included in the result, and must be reported by the application.
    /// If an error is returned, the [`ManagementActionClient`] should not be used until there is an Ok update.
    UpdatedWithNewExecutor(Result<ManagementActionExecutor, AdrConfigError>),
    /// Indicates that the Management Action has been deleted.
    Deleted,
}

/// Enum to indicate whether the schema being reported is for a request or response
#[derive(Debug)]
enum ActionSchema {
    Request,
    Response,
}

/// Azure Device Registry Management Action Client represents a Management Action and its Management Group
/// and includes additional functionality to report status, report message schemas, receive updates,
/// and receive execution requests for the Management Action.
pub struct ManagementActionClient {
    /// Management action, management group, asset, device, and inbound endpoint names
    management_action_ref: ManagementActionRef,
    // Management Action and Management Group Definition
    definition: ManagementActionSpecification,
    /// Current status for the Asset
    asset_status: Arc<tokio::sync::RwLock<adr_models::AssetStatus>>,
    /// Current specification for the Asset
    asset_specification: Arc<std::sync::RwLock<AssetSpecification>>,
    /// Specification of the device that this management action is tied to
    device_specification: Arc<std::sync::RwLock<DeviceSpecification>>,
    /// Status of the device that this management action is tied to
    device_status: Arc<tokio::sync::RwLock<DeviceEndpointStatus>>,
    // Internally used fields
    /// Any previously detected config status for the management action
    previous_detected_config_status: Result<(), AdrConfigError>,
    /// used to shutdown the current executor if the definition for it is no longer current
    executor_shutdown_notifier: Option<Arc<Notify>>,
    connector_context: Arc<ConnectorContext>,
    /// Asset reference for internal use
    asset_ref: AssetRef,
    /// Internal watcher receiver that holds a snapshot of the latest update and whether it has been
    /// fully processed or not.
    management_action_update_watcher_rx: watch::Receiver<ManagementActionUpdateNotification>,
    /// Health reporter sender for sending health events
    health_sender: HealthReporterSender,
    /// Cancellation token for health reporting task - cancelled on deletion
    health_cancellation_token: CancellationToken,
}

impl ManagementActionClient {
    #[allow(clippy::too_many_arguments)]
    pub(crate) fn new(
        definition: ManagementActionSpecification,
        management_action_update_watcher_rx: watch::Receiver<ManagementActionUpdateNotification>,
        asset_ref: AssetRef,
        asset_status: Arc<tokio::sync::RwLock<adr_models::AssetStatus>>,
        asset_specification: Arc<std::sync::RwLock<AssetSpecification>>,
        device_specification: Arc<std::sync::RwLock<DeviceSpecification>>,
        device_status: Arc<tokio::sync::RwLock<DeviceEndpointStatus>>,
        connector_context: Arc<ConnectorContext>,
    ) -> (Self, Result<ManagementActionExecutor, AdrConfigError>) {
        // Create a new management_action
        let management_action_ref = ManagementActionRef {
            management_action_name: definition.name.clone(),
            management_group_name: definition.management_group.name.clone(),
            asset_name: asset_ref.name.clone(),
            device_name: asset_ref.device_name.clone(),
            inbound_endpoint_name: asset_ref.inbound_endpoint_name.clone(),
        };
        let health_cancellation_token = CancellationToken::new();
        let health_sender = connector_context
            .azure_device_registry_client
            .new_management_action_health_reporter(
                asset_ref.clone().into(),
                management_action_ref.management_group_name.clone(),
                management_action_ref.management_action_name.clone(),
                connector_context.azure_device_registry_timeout,
                connector_context.health_report_interval,
                health_cancellation_token.clone(),
            );
        // create executor
        let executor_res = management_action_executor::try_executor_topic_from_management_topics(
            definition.topic.as_ref(),
            definition.management_group.default_topic.as_ref(),
        )
        .and_then(|topic| {
            ManagementActionExecutor::new(topic, &management_action_ref, &connector_context)
        });

        (
            Self {
                management_action_ref,
                definition,
                asset_status,
                asset_specification,
                device_specification,
                device_status,
                connector_context,
                asset_ref,
                management_action_update_watcher_rx,
                previous_detected_config_status: executor_res
                    .as_ref()
                    .map(|_| ())
                    .map_err(Clone::clone),
                executor_shutdown_notifier: executor_res
                    .as_ref()
                    .ok()
                    .map(ManagementActionExecutor::get_shutdown_notifier),
                health_sender,
                health_cancellation_token,
            },
            executor_res,
        )
    }

    /// Used to conditionally report the request message schema of a management action as an existing schema reference
    ///
    /// The `modify` function is called with the current message schema reference (if any) and should return:
    /// - `Some(new_message_schema_reference)` if the schema should be updated and reported
    /// - `None` if no update is needed
    ///
    /// # Returns
    /// - [`SchemaModifyResult::Reported`] if the schema was updated and successfully reported, containing the reported [`MessageSchemaReference`]
    /// - [`SchemaModifyResult::NotModified`] if no modification was needed or the version changed during processing
    ///
    /// # Errors
    /// [`MessageSchemaError`] of kind [`AzureDeviceRegistryError::AIOProtocolError`](azure_device_registry::ErrorKind::AIOProtocolError) if
    /// there are any underlying errors from the AIO RPC protocol. This error will be retried
    /// 10 times with exponential backoff and jitter and only returned if it still is failing.
    ///
    /// [`MessageSchemaError`] of kind [`AzureDeviceRegistryError::ServiceError`](azure_device_registry::ErrorKind::ServiceError) if an error is returned
    /// by the Azure Device Registry service.
    ///
    /// # Panics
    /// If the asset specification mutex has been poisoned, which should not be possible
    pub async fn report_request_message_schema_reference_if_modified<F>(
        &mut self,
        modify: F,
    ) -> Result<SchemaModifyResult, MessageSchemaError>
    where
        F: Fn(Option<&MessageSchemaReference>) -> Option<MessageSchemaReference>,
    {
        self.report_message_schema_reference_if_modified(ActionSchema::Request, modify)
            .await
    }

    /// Used to conditionally report the response message schema of a management action as an existing schema reference
    ///
    /// The `modify` function is called with the current message schema reference (if any) and should return:
    /// - `Some(new_message_schema_reference)` if the schema should be updated and reported
    /// - `None` if no update is needed
    ///
    /// # Returns
    /// - [`SchemaModifyResult::Reported`] if the schema was updated and successfully reported, containing the reported [`MessageSchemaReference`]
    /// - [`SchemaModifyResult::NotModified`] if no modification was needed or the version changed during processing
    ///
    /// # Errors
    /// [`MessageSchemaError`] of kind [`AzureDeviceRegistryError::AIOProtocolError`](azure_device_registry::ErrorKind::AIOProtocolError) if
    /// there are any underlying errors from the AIO RPC protocol. This error will be retried
    /// 10 times with exponential backoff and jitter and only returned if it still is failing.
    ///
    /// [`MessageSchemaError`] of kind [`AzureDeviceRegistryError::ServiceError`](azure_device_registry::ErrorKind::ServiceError) if an error is returned
    /// by the Azure Device Registry service.
    ///
    /// # Panics
    /// If the asset specification mutex has been poisoned, which should not be possible
    pub async fn report_response_message_schema_reference_if_modified<F>(
        &mut self,
        modify: F,
    ) -> Result<SchemaModifyResult, MessageSchemaError>
    where
        F: Fn(Option<&MessageSchemaReference>) -> Option<MessageSchemaReference>,
    {
        self.report_message_schema_reference_if_modified(ActionSchema::Response, modify)
            .await
    }

    /// Used to conditionally report the request message schema of a management action
    ///
    /// The `modify` function is called with the current message schema reference (if any) and should return:
    /// - `Some(new_message_schema)` if the schema should be updated and reported
    /// - `None` if no update is needed
    ///
    /// # Returns
    /// - [`SchemaModifyResult::Reported`] if the schema was updated and successfully reported, containing the reported [`MessageSchemaReference`]
    /// - [`SchemaModifyResult::NotModified`] if no modification was needed or the version changed during processing
    ///
    /// # Errors
    /// [`MessageSchemaError`] of kind [`SchemaRegistryError::InvalidRequestArgument`](schema_registry::ErrorKind::InvalidRequestArgument)
    /// if the content of the [`MessageSchema`] is empty or there is an error building the request
    ///
    /// [`MessageSchemaError`] of kind [`SchemaRegistryError::ServiceError`](schema_registry::ErrorKind::ServiceError)
    /// if there is an error returned by the Schema Registry Service. This error will be retried 10
    /// times with exponential backoff and jitter if it is an internal error and only returned if
    /// it still is failing.
    ///
    /// [`MessageSchemaError`] of kind [`AzureDeviceRegistryError::AIOProtocolError`](azure_device_registry::ErrorKind::AIOProtocolError) if
    /// there are any underlying errors from the AIO RPC protocol. This error will be retried
    /// 10 times with exponential backoff and jitter and only returned if it still is failing.
    ///
    /// [`MessageSchemaError`] of kind [`AzureDeviceRegistryError::ServiceError`](azure_device_registry::ErrorKind::ServiceError) if
    /// an error is returned by the Azure Device Registry service.
    ///
    /// # Panics
    /// If the asset specification mutex has been poisoned, which should not be possible
    pub async fn report_request_message_schema_if_modified<F>(
        &mut self,
        modify: F,
    ) -> Result<SchemaModifyResult, MessageSchemaError>
    where
        F: Fn(Option<&MessageSchemaReference>) -> Option<MessageSchema>,
    {
        self.report_message_schema_if_modified(ActionSchema::Request, modify)
            .await
    }

    /// Used to conditionally report the response message schema of a management action
    ///
    /// The `modify` function is called with the current message schema reference (if any) and should return:
    /// - `Some(new_message_schema)` if the schema should be updated and reported
    /// - `None` if no update is needed
    ///
    /// # Returns
    /// - [`SchemaModifyResult::Reported`] if the schema was updated and successfully reported, containing the reported [`MessageSchemaReference`]
    /// - [`SchemaModifyResult::NotModified`] if no modification was needed or the version changed during processing
    ///
    /// # Errors
    /// [`MessageSchemaError`] of kind [`SchemaRegistryError::InvalidRequestArgument`](schema_registry::ErrorKind::InvalidRequestArgument)
    /// if the content of the [`MessageSchema`] is empty or there is an error building the request
    ///
    /// [`MessageSchemaError`] of kind [`SchemaRegistryError::ServiceError`](schema_registry::ErrorKind::ServiceError)
    /// if there is an error returned by the Schema Registry Service. This error will be retried 10
    /// times with exponential backoff and jitter if it is an internal error and only returned if
    /// it still is failing.
    ///
    /// [`MessageSchemaError`] of kind [`AzureDeviceRegistryError::AIOProtocolError`](azure_device_registry::ErrorKind::AIOProtocolError) if
    /// there are any underlying errors from the AIO RPC protocol. This error will be retried
    /// 10 times with exponential backoff and jitter and only returned if it still is failing.
    ///
    /// [`MessageSchemaError`] of kind [`AzureDeviceRegistryError::ServiceError`](azure_device_registry::ErrorKind::ServiceError) if
    /// an error is returned by the Azure Device Registry service.
    ///
    /// # Panics
    /// If the asset specification mutex has been poisoned, which should not be possible
    pub async fn report_response_message_schema_if_modified<F>(
        &mut self,
        modify: F,
    ) -> Result<SchemaModifyResult, MessageSchemaError>
    where
        F: Fn(Option<&MessageSchemaReference>) -> Option<MessageSchema>,
    {
        self.report_message_schema_if_modified(ActionSchema::Response, modify)
            .await
    }

    /// Used to receive notifications about the Management Action from the Azure Device Registry Service.
    ///
    /// Returns [`ManagementActionNotification::ManagementActionNewExecutor`] if the Management Group or Action's definition has been updated in place
    /// and the definition change required a new executor to be created (e.g. topic change). If this is received, the old executor will be shut down,
    /// but it can still be used to drain/complete any already received requests until it is dropped.
    ///
    /// Returns [`ManagementActionNotification::ManagementActionUpdated`] if the Management Group or Action's definition has been updated in place,
    /// but the definition change did not require a new executor to be created.
    ///
    /// Returns [`ManagementActionNotification::AssetUpdated`] if the Asset's definition has been updated in place, but there were no
    /// changes to the Management Group or Action's definition.
    ///
    /// Returns any of the above with an `Err(AdrConfigError)` if the Management Action/Asset received an update, but a config error
    /// for the Management Action was detected. The definition is still updated in place, but the [`ManagementActionClient`] should not be
    /// used until there is a new update because receiving requests will not be possible.
    ///
    /// Returns [`ManagementActionNotification::Deleted`] if the Management Action has been deleted. The [`ManagementActionClient`]
    /// should not be used after this point, and no more notifications will be received. If this is received, the existing executor
    /// will be shut down, but it can still be used to drain/complete any already received requests until it is dropped.
    ///
    /// # Panics
    /// If the asset specification mutex has been poisoned, which should not be possible
    ///
    /// # Cancel safety
    /// This method is cancel safe. If you use it as the event in a `tokio::select!` statement and some other branch
    /// completes first, then it is guaranteed that no management action notifications will be lost, and the management action will not
    /// be updated without a notification being returned.
    pub async fn recv_notification(&mut self) -> ManagementActionNotification {
        if self
            .management_action_update_watcher_rx
            .changed()
            .await
            .is_err()
        {
            // shut down the current executor if it exists
            self.executor_shutdown_notifier
                .as_ref()
                .inspect(|notify| notify.notify_one());
            // Cancel health reporting task on deletion
            self.health_cancellation_token.cancel();
            return ManagementActionNotification::Deleted;
        }
        // In case this function gets cancelled the next time it is called we will process the update again.
        self.management_action_update_watcher_rx.mark_changed();
        let mut update_notification = self.management_action_update_watcher_rx.borrow().clone();

        // wait until the update has been released. If the watch sender has been dropped, this means the Asset has been deleted/dropped
        if update_notification
            .release_asset_component_notifications_rx
            .changed()
            .await
            .is_err()
        {
            self.management_action_update_watcher_rx.mark_unchanged();
            // shut down the current executor if it exists
            self.executor_shutdown_notifier
                .as_ref()
                .inspect(|notify| notify.notify_one());
            // Cancel health reporting task on deletion
            self.health_cancellation_token.cancel();
            return ManagementActionNotification::Deleted;
        }
        if update_notification.definition == self.definition {
            // Asset only update, we don't need to recreate the executor
            // Once the action update has been processed we can mark the value in the watcher as seen
            self.management_action_update_watcher_rx.mark_unchanged();
            // if nothing changed, but there was an executor error before, we need to return that error again since it still applies
            return ManagementActionNotification::AssetUpdated(
                self.previous_detected_config_status.clone(),
            );
        }

        // TODO: need to take tokens into account here? If the tokens are required, then there shouldn't be any need to evaluate the topics replaced?
        let new_topic = update_notification
            .definition
            .topic
            .as_ref()
            .or(update_notification
                .definition
                .management_group
                .default_topic
                .as_ref());
        let old_topic = self.definition.topic.as_ref().or(self
            .definition
            .management_group
            .default_topic
            .as_ref());
        if new_topic == old_topic {
            // topic won't be changed, we don't need to recreate the executor
            // update the definition in place
            self.definition = update_notification.definition;
            // Once the action update has been processed we can mark the value in the watcher as seen
            self.management_action_update_watcher_rx.mark_unchanged();
            // if nothing changed, but there was an executor error before, we need to return that error again since it still applies
            return ManagementActionNotification::Updated(
                self.previous_detected_config_status.clone(),
            );
        }

        // no await points beyond this point, so this is safe
        self.definition = update_notification.definition;

        // shut down the current executor if it exists and create new executor, since topics have changed
        self.executor_shutdown_notifier
            .as_ref()
            .inspect(|notify| notify.notify_one());
        let executor_res = management_action_executor::try_executor_topic_from_management_topics(
            self.definition.topic.as_ref(),
            self.definition.management_group.default_topic.as_ref(),
        )
        .and_then(|topic| {
            ManagementActionExecutor::new(
                topic,
                &self.management_action_ref,
                &self.connector_context,
            )
        });

        self.executor_shutdown_notifier = executor_res
            .as_ref()
            .ok()
            .map(ManagementActionExecutor::get_shutdown_notifier);
        self.previous_detected_config_status =
            executor_res.as_ref().map(|_| ()).map_err(Clone::clone);

        // Once the action definition has been updated we can mark the value in the watcher as seen
        self.management_action_update_watcher_rx.mark_unchanged();
        ManagementActionNotification::UpdatedWithNewExecutor(executor_res)
    }

    /// Creates a new status reporter for this [`ManagementActionClient`]
    /// The reporter's version snapshot is initialized to the current asset specification version.
    ///
    /// # Panics
    /// if the asset specification mutex has been poisoned, which should not be possible
    #[must_use]
    pub fn get_status_reporter(&self) -> ManagementActionStatusReporter {
        let snapshotted_version = self
            .asset_specification
            .read()
            .unwrap()
            .version
            .unwrap_or(0);
        ManagementActionStatusReporter {
            connector_context: self.connector_context.clone(),
            asset_status: self.asset_status.clone(),
            asset_specification: self.asset_specification.clone(),
            health_sender: self.health_sender.clone(),
            asset_component_ref: self.management_action_ref.clone(),
            asset_ref: self.asset_ref.clone(),
            snapshotted_version,
        }
    }

    // ~~~~~~~ fns to get various fields ~~~~~~~

    /// Returns a clone of this Management Action's Request [`MessageSchemaReference`] from
    /// the `AssetStatus`, if it exists
    #[must_use]
    pub async fn request_message_schema_reference(&self) -> Option<MessageSchemaReference> {
        self.asset_status
            .read()
            .await
            .management_groups
            .as_ref()?
            .iter()
            .find(|management_group| {
                management_group.name == *self.management_action_ref.management_group_name
            })?
            .actions
            .as_ref()?
            .iter()
            .find(|action| action.name == *self.management_action_ref.management_action_name)?
            .request_message_schema_reference
            .clone()
    }

    /// Returns a clone of this Management Action's Response [`MessageSchemaReference`] from
    /// the `AssetStatus`, if it exists
    #[must_use]
    pub async fn response_message_schema_reference(&self) -> Option<MessageSchemaReference> {
        self.asset_status
            .read()
            .await
            .management_groups
            .as_ref()?
            .iter()
            .find(|management_group| {
                management_group.name == *self.management_action_ref.management_group_name
            })?
            .actions
            .as_ref()?
            .iter()
            .find(|action| action.name == *self.management_action_ref.management_action_name)?
            .response_message_schema_reference
            .clone()
    }

    /// Management action, management group, asset, device, and inbound endpoint names
    #[must_use]
    pub fn management_action_ref(&self) -> &ManagementActionRef {
        &self.management_action_ref
    }

    /// Management Action and Group Definitions
    #[must_use]
    pub fn definition(&self) -> &ManagementActionSpecification {
        &self.definition
    }

    /// Returns a clone of the current asset specification
    /// # Panics
    /// if the asset specification mutex has been poisoned, which should not be possible
    #[must_use]
    pub fn asset_specification(&self) -> AssetSpecification {
        (*self.asset_specification.read().unwrap()).clone()
    }

    /// Returns a clone of the current asset status, if it exists
    #[must_use]
    pub async fn asset_status(&self) -> adr_models::AssetStatus {
        (*self.asset_status.read().await).clone()
    }

    /// Returns a clone of the current device specification
    /// # Panics
    /// if the device specification mutex has been poisoned, which should not be possible
    #[must_use]
    pub fn device_specification(&self) -> DeviceSpecification {
        (*self.device_specification.read().unwrap()).clone()
    }

    /// Returns a clone of the current device status
    #[must_use]
    pub async fn device_status(&self) -> DeviceEndpointStatus {
        (*self.device_status.read().await).clone()
    }

    // ~~~~~~ Internal fns ~~~~~~~

    /// Helper function to update the specific event status within the asset status
    fn update_action_status(
        asset_status_to_update: &mut adr_models::AssetStatus,
        management_group_name: &str,
        action_name: &str,
        action_status: Result<(), AdrConfigError>,
    ) {
        if let Some(curr_management_group_status) = asset_status_to_update
            .management_groups
            .as_mut()
            .and_then(|management_groups| {
                management_groups
                    .iter_mut()
                    .find(|management_group| management_group.name == management_group_name)
            })
        {
            if let Some(curr_action_status) = curr_management_group_status
                .actions
                .as_mut()
                .and_then(|actions| actions.iter_mut().find(|action| action.name == action_name))
            {
                // If the action already has a status, update the existing action with the new error
                curr_action_status.error = action_status.err();
            } else {
                // If the action doesn't exist in the current management group status, then add it
                curr_management_group_status
                    .actions
                    .get_or_insert_with(Vec::new)
                    .push(adr_models::ActionStatus {
                        name: action_name.to_string(),
                        request_message_schema_reference: None,
                        response_message_schema_reference: None,
                        error: action_status.err(),
                    });
            }
        } else {
            // If the management group doesn't exist in the current status, then add it
            asset_status_to_update
                .management_groups
                .get_or_insert_with(Vec::new)
                .push(adr_models::ManagementGroupStatus {
                    name: management_group_name.to_string(),
                    actions: Some(vec![adr_models::ActionStatus {
                        name: action_name.to_string(),
                        request_message_schema_reference: None,
                        response_message_schema_reference: None,
                        error: action_status.err(),
                    }]),
                });
        }
    }

    async fn report_message_schema_reference_if_modified<F>(
        &mut self,
        schema_side: ActionSchema,
        modify: F,
    ) -> Result<SchemaModifyResult, MessageSchemaError>
    where
        F: Fn(Option<&MessageSchemaReference>) -> Option<MessageSchemaReference>,
    {
        // Get the current version of the asset specification
        let cached_version = self.asset_specification.read().unwrap().version;

        {
            // Get the current asset status
            let status_read_guard = self.asset_status.read().await;
            let current_asset_status =
                AssetClient::get_current_asset_status(&status_read_guard, cached_version);

            let modify_input =
                self.get_schema_reference_modify_input(&current_asset_status, &schema_side);

            match modify(modify_input) {
                Some(_) => {
                    // A modification was made, we proceed to report schema
                }
                None => {
                    // No modification was made, so no need to report schema
                    return Ok(SchemaModifyResult::NotModified);
                }
            }
        };

        // To modify, we need to acquire the write lock
        let mut status_write_guard = self.asset_status.write().await;

        // We can continue here because the asset status will not change. The specification might
        // update but we will be reporting for an old version.

        if cached_version != self.asset_specification.read().unwrap().version {
            // Our modify is no longer valid
            log::debug!(
                "Reporting for an out-of-date asset specification from a Management Action client, will not modify"
            );
            return Ok(SchemaModifyResult::NotModified);
        }

        // Get the current asset status in case it has changed
        let current_asset_status =
            AssetClient::get_current_asset_status(&status_write_guard, cached_version);

        let modify_input =
            self.get_schema_reference_modify_input(&current_asset_status, &schema_side);

        let Some(new_message_schema_reference) = modify(modify_input) else {
            // No modification was made, so no need to report schema
            return Ok(SchemaModifyResult::NotModified);
        };

        let mut asset_status_to_report = current_asset_status.into_owned();

        asset_status_to_report.config = match asset_status_to_report.config {
            Some(mut config) => {
                config.last_transition_time = Some(Utc::now());
                Some(config)
            }
            None => {
                // If the config is None, we need to create a new one to report along with the asset status
                Some(azure_device_registry::ConfigStatus {
                    version: cached_version,
                    last_transition_time: Some(Utc::now()),
                    ..Default::default()
                })
            }
        };

        Self::internal_report_message_schema_reference(
            &self.connector_context,
            &self.asset_ref,
            &self.management_action_ref,
            // &mut self.forwarder,
            asset_status_to_report,
            &mut status_write_guard,
            &new_message_schema_reference,
            &schema_side,
            &format!("ManagementActionClient::report_{schema_side:?}_message_schema_reference_if_modified"),
        )
        .await?;

        Ok(SchemaModifyResult::Reported(new_message_schema_reference))
    }

    async fn report_message_schema_if_modified<F>(
        &mut self,
        schema_side: ActionSchema,
        modify: F,
    ) -> Result<SchemaModifyResult, MessageSchemaError>
    where
        F: Fn(Option<&MessageSchemaReference>) -> Option<MessageSchema>,
    {
        // Get the current version of the asset specification
        let cached_version = self.asset_specification.read().unwrap().version;

        {
            // Get the current asset status
            let status_read_guard = self.asset_status.read().await;
            let current_asset_status =
                AssetClient::get_current_asset_status(&status_read_guard, cached_version);

            let modify_input =
                self.get_schema_reference_modify_input(&current_asset_status, &schema_side);

            if modify(modify_input).is_some() {
                // A modification was made, we proceed to report schema
            } else {
                // No modification was made, so no need to report schema
                return Ok(SchemaModifyResult::NotModified);
            }
        };

        // To modify, we need to acquire the write lock
        let mut status_write_guard = self.asset_status.write().await;

        // We can continue here because the asset status will not change. The specification might
        // update but we will be reporting for an old version.

        if cached_version != self.asset_specification.read().unwrap().version {
            // Our modify is no longer valid
            log::debug!(
                "Reporting for an out-of-date asset specification from a Management Action client, will not modify"
            );
            return Ok(SchemaModifyResult::NotModified);
        }

        // Get the current asset status in case it has changed
        let current_asset_status =
            AssetClient::get_current_asset_status(&status_write_guard, cached_version);

        let modify_input =
            self.get_schema_reference_modify_input(&current_asset_status, &schema_side);

        let Some(new_message_schema) = modify(modify_input) else {
            // No modification was made, so no need to report schema
            return Ok(SchemaModifyResult::NotModified);
        };

        let mut asset_status_to_report = current_asset_status.into_owned();

        asset_status_to_report.config = match asset_status_to_report.config {
            Some(mut config) => {
                config.last_transition_time = Some(Utc::now());
                Some(config)
            }
            None => {
                // If the config is None, we need to create a new one to report along with the asset status
                Some(azure_device_registry::ConfigStatus {
                    version: cached_version,
                    last_transition_time: Some(Utc::now()),
                    ..Default::default()
                })
            }
        };

        // First put the schema in the schema registry
        let message_schema_reference = Retry::spawn(
            RETRY_STRATEGY.map(tokio_retry2::strategy::jitter),
            async || -> Result<schema_registry::Schema, RetryError<schema_registry::Error>> {
                self.connector_context
                    .schema_registry_client
                    .put(
                        new_message_schema.clone(),
                        self.connector_context.schema_registry_timeout,
                    )
                    .await
                    .map_err(|e| {
                        match e.kind() {
                            // network/retriable
                            schema_registry::ErrorKind::AIOProtocolError(_) => {
                                log::warn!(
                                    "Reporting {schema_side:?} message schema failed for {:?}. Retrying: {e}",
                                    self.management_action_ref
                                );
                                RetryError::transient(e)
                            }
                            schema_registry::ErrorKind::ServiceError(service_error) => {
                                if let schema_registry::ErrorCode::InternalError =
                                    service_error.code
                                {
                                    log::warn!(
                                        "Reporting {schema_side:?} message schema failed for {:?}. Retrying: {e}",
                                        self.management_action_ref
                                    );
                                    RetryError::transient(e)
                                } else {
                                    RetryError::permanent(e)
                                }
                            }
                            // indicates an error in the provided message schema, return to caller so they can fix
                            schema_registry::ErrorKind::InvalidRequestArgument(_) => {
                                RetryError::permanent(e)
                            }
                        }
                    })
            },
        )
        .await
        .map(|schema| MessageSchemaReference {
            name: schema.name,
            version: schema.version,
            registry_namespace: schema.namespace,
        })?;

        Self::internal_report_message_schema_reference(
            &self.connector_context,
            &self.asset_ref,
            &self.management_action_ref,
            // &mut self.forwarder,
            asset_status_to_report,
            &mut status_write_guard,
            &message_schema_reference,
            &schema_side,
            &format!("ManagementActionClient::report_{schema_side:?}_message_schema_if_modified"),
        )
        .await?;

        Ok(SchemaModifyResult::Reported(message_schema_reference))
    }

    #[allow(clippy::too_many_arguments)]
    async fn internal_report_message_schema_reference(
        connector_context: &Arc<ConnectorContext>,
        asset_ref: &AssetRef,
        management_action_ref: &ManagementActionRef,
        asset_status_to_report: adr_models::AssetStatus,
        status_write_guard: &mut adr_models::AssetStatus,
        message_schema_reference: &MessageSchemaReference,
        schema_side: &ActionSchema,
        log_identifier: &str,
    ) -> Result<(), MessageSchemaError> {
        // Use the provided asset_status_to_report instead of creating a new one
        let mut new_status = asset_status_to_report;

        // if management action is already in the current status, then update the existing management action with the new message schema
        // Otherwise if the management action isn't present, or no management actions have been reported yet, then add it with the new message schema
        if let Some(management_group_status) =
            new_status
                .management_groups
                .as_mut()
                .and_then(|management_groups| {
                    management_groups.iter_mut().find(|management_group| {
                        management_group.name == *management_action_ref.management_group_name
                    })
                })
        {
            if let Some(action_status) =
                management_group_status
                    .actions
                    .as_mut()
                    .and_then(|actions| {
                        actions.iter_mut().find(|action| {
                            action.name == *management_action_ref.management_action_name
                        })
                    })
            {
                // If the action already has a status, update the existing action with the new message schema
                match schema_side {
                    ActionSchema::Request => {
                        action_status.request_message_schema_reference =
                            Some(message_schema_reference.clone());
                    }
                    ActionSchema::Response => {
                        action_status.response_message_schema_reference =
                            Some(message_schema_reference.clone());
                    }
                }
            } else {
                // If the action doesn't exist in the current status, then add it
                management_group_status
                    .actions
                    .get_or_insert_with(Vec::new)
                    .push(adr_models::ActionStatus {
                        name: management_action_ref.management_action_name.clone(),
                        request_message_schema_reference: match schema_side {
                            ActionSchema::Request => Some(message_schema_reference.clone()),
                            ActionSchema::Response => None,
                        },
                        response_message_schema_reference: match schema_side {
                            ActionSchema::Request => None,
                            ActionSchema::Response => Some(message_schema_reference.clone()),
                        },
                        error: None,
                    });
            }
        } else {
            // If the management group doesn't exist in the current status, then add it
            new_status
                .management_groups
                .get_or_insert_with(Vec::new)
                .push(adr_models::ManagementGroupStatus {
                    name: management_action_ref.management_group_name.clone(),
                    actions: Some(vec![adr_models::ActionStatus {
                        name: management_action_ref.management_action_name.clone(),
                        request_message_schema_reference: match schema_side {
                            ActionSchema::Request => Some(message_schema_reference.clone()),
                            ActionSchema::Response => None,
                        },
                        response_message_schema_reference: match schema_side {
                            ActionSchema::Request => None,
                            ActionSchema::Response => Some(message_schema_reference.clone()),
                        },
                        error: None,
                    }]),
                });
        }

        // send status update to the service
        log::debug!(
            "Reporting management action {management_action_ref:?} {schema_side:?} message schema from app"
        );
        AssetStatusReporter::internal_report_status(
            new_status,
            connector_context,
            asset_ref,
            status_write_guard,
            log_identifier,
        )
        .await?;

        Ok(())
    }

    fn get_schema_reference_modify_input<'a>(
        &self,
        asset_status: &'a adr_models::AssetStatus,
        schema_side: &ActionSchema,
    ) -> Option<&'a adr_models::MessageSchemaReference> {
        asset_status
            .management_groups
            .as_ref()
            .and_then(|management_groups| {
                management_groups.iter().find(|mg_status| {
                    mg_status.name == *self.management_action_ref.management_group_name
                })
            })
            .and_then(|mg_status| {
                mg_status.actions.as_ref().and_then(|actions| {
                    actions.iter().find(|a_status| {
                        a_status.name == *self.management_action_ref.management_action_name
                    })
                })
            })
            .and_then(|a_status| match schema_side {
                ActionSchema::Request => a_status.request_message_schema_reference.as_ref(),
                ActionSchema::Response => a_status.response_message_schema_reference.as_ref(),
            })
    }
}

impl Drop for ManagementActionClient {
    fn drop(&mut self) {
        // shut down the current executor if it exists
        self.executor_shutdown_notifier
            .as_ref()
            .inspect(|notify| notify.notify_one());
        self.health_cancellation_token.cancel();
    }
}

// Client Structs
// Device

#[derive(Debug, Clone)]
/// Represents the specification of a device in the Azure Device Registry service.
pub struct DeviceSpecification {
    /// The 'attributes' Field.
    pub attributes: HashMap<String, String>,
    /// The 'discoveredDeviceRef' Field.
    pub discovered_device_ref: Option<String>,
    /// The 'enabled' Field.
    pub enabled: Option<bool>,
    /// The 'endpoints' Field.
    pub endpoints: DeviceEndpoints, // different from adr
    /// The 'externalDeviceId' Field.
    pub external_device_id: Option<String>,
    /// The 'lastTransitionTime' Field.
    pub last_transition_time: Option<DateTime<Utc>>,
    /// The 'manufacturer' Field.
    pub manufacturer: Option<String>,
    /// The 'model' Field.
    pub model: Option<String>,
    /// The 'operatingSystem' Field.
    pub operating_system: Option<String>,
    /// The 'operatingSystemVersion' Field.
    pub operating_system_version: Option<String>,
    /// The 'uuid' Field.
    pub uuid: Option<String>,
    /// The 'version' Field.
    pub version: Option<u64>,
}

impl DeviceSpecification {
    pub(crate) fn new(
        device_specification: adr_models::Device,
        device_endpoint_credentials_mount_path: Option<&PathBuf>,
        inbound_endpoint_name: &str,
    ) -> Result<Self, String> {
        // convert the endpoints to the new format with only the one specified inbound endpoint
        // if the inbound endpoint isn't in the specification, return an error
        let recvd_endpoints = device_specification
            .endpoints
            .ok_or("Endpoints not found on Device specification")?;

        let recvd_inbound = recvd_endpoints
            .inbound
            .get(inbound_endpoint_name)
            .cloned()
            .ok_or("Inbound endpoint not found on Device specification")?;

        // update authentication to include the full file path for the credentials
        let authentication = match recvd_inbound.authentication {
            adr_models::Authentication::Anonymous => Authentication::Anonymous,
            adr_models::Authentication::Certificate {
                certificate_secret_name,
                intermediate_certificates_secret_name,
                key_secret_name,
            } => Authentication::Certificate {
                certificate_path: device_endpoint_credentials_mount_path.expect("device_endpoint_credentials_mount_path must be present if Authentication is Certificate").as_path().join(certificate_secret_name),
                intermediate_certificates_path: intermediate_certificates_secret_name.map(|intermediate_cert_secret_name| {
                    device_endpoint_credentials_mount_path.expect("device_endpoint_credentials_mount_path must be present if Authentication is Certificate").as_path().join(intermediate_cert_secret_name)
                }),
                key_path: key_secret_name.map(|key_secret_name| {
                    device_endpoint_credentials_mount_path.expect("device_endpoint_credentials_mount_path must be present if Authentication is Certificate").as_path().join(key_secret_name)
                }),

            },
            adr_models::Authentication::UsernamePassword {
                password_secret_name,
                username_secret_name,
            } => Authentication::UsernamePassword {
                password_path: device_endpoint_credentials_mount_path.expect("device_endpoint_credentials_mount_path must be present if Authentication is UsernamePassword").as_path().join(password_secret_name),
                username_path: device_endpoint_credentials_mount_path.expect("device_endpoint_credentials_mount_path must be present if Authentication is UsernamePassword").as_path().join(username_secret_name),
            },
        };

        let endpoints = DeviceEndpoints {
            inbound: InboundEndpoint {
                name: inbound_endpoint_name.to_string(),
                additional_configuration: recvd_inbound.additional_configuration,
                address: recvd_inbound.address,
                authentication,
                endpoint_type: recvd_inbound.endpoint_type,
                trust_settings: recvd_inbound.trust_settings,
                version: recvd_inbound.version,
            },
            outbound: recvd_endpoints.outbound,
        };

        Ok(DeviceSpecification {
            attributes: device_specification.attributes,
            discovered_device_ref: device_specification.discovered_device_ref,
            enabled: device_specification.enabled,
            endpoints,
            external_device_id: device_specification.external_device_id,
            last_transition_time: device_specification.last_transition_time,
            manufacturer: device_specification.manufacturer,
            model: device_specification.model,
            operating_system: device_specification.operating_system,
            operating_system_version: device_specification.operating_system_version,
            uuid: device_specification.uuid,
            version: device_specification.version,
        })
    }
}

#[derive(Debug, Clone)]
/// Represents the endpoints of a device in the Azure Device Registry service.
pub struct DeviceEndpoints {
    /// The 'inbound' Field.
    pub inbound: InboundEndpoint, // different from adr
    /// The 'outbound' Field.
    pub outbound: Option<adr_models::OutboundEndpoints>,
}
/// Represents an inbound endpoint of a device in the Azure Device Registry service.
#[derive(Debug, Clone)]
pub struct InboundEndpoint {
    /// name
    pub name: String,
    /// The 'additionalConfiguration' Field.
    pub additional_configuration: Option<String>,
    /// The 'address' Field.
    pub address: String,
    /// The 'authentication' Field.
    pub authentication: Authentication, // different from adr
    /// The 'endpointType' Field.
    pub endpoint_type: String,
    /// The 'trustSettings' Field.
    pub trust_settings: Option<adr_models::TrustSettings>,
    /// The 'version' Field.
    pub version: Option<String>,
}

#[derive(Debug, Clone, Default)]
/// Represents the authentication method for an endpoint.
pub enum Authentication {
    #[default]
    /// Represents anonymous authentication.
    Anonymous,
    /// Represents authentication using a certificate.
    Certificate {
        /// The path to a file containing containing the client certificate in PEM format.
        certificate_path: PathBuf, // different from adr
        /// The path to a file containing the combined intermediate certificates in PEM format (if any).
        intermediate_certificates_path: Option<PathBuf>, // different from adr
        /// The path to a file containing the private key in PEM or DER format.
        key_path: Option<PathBuf>, // different from adr
    },
    /// Represents authentication using a username and password.
    UsernamePassword {
        /// The 'passwordSecretName' Field.
        password_path: PathBuf, // different from adr
        /// The 'usernameSecretName' Field.
        username_path: PathBuf, // different from adr
    },
}

#[derive(Clone, Debug, Default, PartialEq)]
/// Represents the observed status of a Device and endpoint in the ADR Service.
pub struct DeviceEndpointStatus {
    /// Defines the status for the Device.
    pub config: Option<azure_device_registry::ConfigStatus>,
    /// Defines the status for the inbound endpoint.
    /// Specifies whether the inbound endpoint status has been reported, and if so, whether it was successful or not.
    pub inbound_endpoint_status: Option<Result<(), AdrConfigError>>, // different from adr
}

impl DeviceEndpointStatus {
    pub(crate) fn new(recvd_status: adr_models::DeviceStatus, inbound_endpoint_name: &str) -> Self {
        let inbound_endpoint_status = recvd_status.endpoints.get(inbound_endpoint_name).cloned();
        DeviceEndpointStatus {
            config: recvd_status.config,
            // inbound_endpoint_status: inbound_endpoint_status.map(|status| status.map(|e| Err(e))),
            inbound_endpoint_status: inbound_endpoint_status.map(|status| match status {
                Some(e) => Err(e),
                None => Ok(()),
            }),
        }
    }

    /// Internal helper to get an [`DeviceEndpointStatus`] that can be used as a starting place
    /// to modify the current status with whatever new things we want to report.
    ///
    /// Note that it returns a `Cow`. The reason is that most of the times that we are reporting
    /// a status we will not end up modifying it. `Cow` allows us to only clone when we are going to
    /// modify.
    pub(crate) fn get_current_device_endpoint_status(
        &self,
        adr_version: Option<u64>,
    ) -> Cow<'_, Self> {
        match &self.config {
            Some(config) => {
                if config.version == adr_version {
                    // If the version in our config matches the one in ADR we return our current
                    // device endpoint status
                    Cow::Borrowed(self)
                } else {
                    // If the version doesn't match, the config version we are holding is stale so
                    // we pass a cleared version
                    Cow::Owned(DeviceEndpointStatus::default())
                }
            }
            None => {
                // If there is no config, then the status is not set, we can pass our current device
                // endpoint status
                Cow::Borrowed(self)
            }
        }
    }
}

/// Represents the specification of an Asset in the Azure Device Registry service.
#[derive(Debug, Clone)]
pub struct AssetSpecification {
    /// URI or type definition ids.
    pub asset_type_refs: Vec<String>, // if None, we can represent as empty vec. Can currently only be length of 1
    /// A set of key-value pairs that contain custom attributes
    pub attributes: HashMap<String, String>, // if None, we can represent as empty hashmap
    /// Default configuration for datasets.
    pub default_datasets_configuration: Option<String>,
    /// Default destinations for datasets.
    pub default_datasets_destinations: Vec<adr_models::DatasetDestination>, // if None, we can represent as empty vec.  Can currently only be length of 1
    /// Default configuration for events.
    pub default_events_configuration: Option<String>,
    /// Default destinations for events.
    pub default_events_destinations: Vec<adr_models::EventStreamDestination>, // if None, we can represent as empty vec.  Can currently only be length of 1
    /// Default configuration for management groups.
    pub default_management_groups_configuration: Option<String>,
    /// Default configuration for streams.
    pub default_streams_configuration: Option<String>,
    /// Default destinations for streams.
    pub default_streams_destinations: Vec<adr_models::EventStreamDestination>, // if None, we can represent as empty vec. Can currently only be length of 1
    /// The description of the asset.
    pub description: Option<String>,
    /// A reference to the Device and Endpoint within the device
    pub device_ref: adr_models::DeviceRef,
    /// Reference to a list of discovered assets
    pub discovered_asset_refs: Vec<String>, // if None, we can represent as empty vec
    /// The display name of the asset.
    pub display_name: Option<String>,
    /// Reference to the documentation.
    pub documentation_uri: Option<String>,
    /// Enabled/Disabled status of the asset.
    pub enabled: Option<bool>, // TODO: just bool?
    /// Asset id provided by the customer.
    pub external_asset_id: Option<String>,
    /// Revision number of the hardware.
    pub hardware_revision: Option<String>,
    /// The last time the asset has been modified.
    pub last_transition_time: Option<DateTime<Utc>>,
    /// The name of the manufacturer.
    pub manufacturer: Option<String>,
    /// The URI of the manufacturer.
    pub manufacturer_uri: Option<String>,
    /// The model of the asset.
    pub model: Option<String>,
    /// The product code of the asset.
    pub product_code: Option<String>,
    /// The revision number of the software.
    pub serial_number: Option<String>,
    /// The revision number of the software.
    pub software_revision: Option<String>,
    ///  Globally unique, immutable, non-reusable id.
    pub uuid: Option<String>,
    /// The version of the asset.
    pub version: Option<u64>,
}

impl From<adr_models::Asset> for AssetSpecification {
    fn from(value: adr_models::Asset) -> Self {
        AssetSpecification {
            asset_type_refs: value.asset_type_refs,
            attributes: value.attributes,
            default_datasets_configuration: value.default_datasets_configuration,
            default_datasets_destinations: value.default_datasets_destinations,
            default_events_configuration: value.default_events_configuration,
            default_events_destinations: value.default_events_destinations,
            default_management_groups_configuration: value.default_management_groups_configuration,
            default_streams_configuration: value.default_streams_configuration,
            default_streams_destinations: value.default_streams_destinations,
            description: value.description,
            device_ref: value.device_ref,
            discovered_asset_refs: value.discovered_asset_refs,
            display_name: value.display_name,
            documentation_uri: value.documentation_uri,
            enabled: value.enabled,
            external_asset_id: value.external_asset_id,
            hardware_revision: value.hardware_revision,
            last_transition_time: value.last_transition_time,
            manufacturer: value.manufacturer,
            manufacturer_uri: value.manufacturer_uri,
            model: value.model,
            product_code: value.product_code,
            serial_number: value.serial_number,
            software_revision: value.software_revision,
            uuid: value.uuid,
            version: value.version,
        }
    }
}

/// Represents the specification of a Management Action and its Management Group in the Azure Device Registry service.
#[derive(Debug, Clone, PartialEq)]
pub struct ManagementActionSpecification {
    /// Configuration for the action.
    pub action_configuration: Option<String>,
    /// Type of the action.
    pub action_type: adr_models::ActionType,
    /// Name of the action.
    pub name: String,
    /// The target URI on which a client can invoke the specific action.
    pub target_uri: String,
    /// Response timeout for the action.
    pub timeout_in_seconds: Option<u64>,
    /// The MQTT topic path on which a client will receive the request for the action.
    pub topic: Option<String>,
    /// URI or type definition ID.
    pub type_ref: Option<String>,
    /// The Management Group this action belongs to.
    pub management_group: ManagementGroupSpecification,
    name_tuple: (String, String),
}

impl ManagementActionSpecification {
    pub(crate) fn command_name(&self) -> String {
        format!("{}::{}", self.management_group.name, self.name)
    }

    pub(crate) fn hash_name(&self) -> &(String, String) {
        &self.name_tuple
    }
}

impl
    From<(
        adr_models::ManagementGroup,
        adr_models::ManagementGroupAction,
    )> for ManagementActionSpecification
{
    fn from(
        val: (
            adr_models::ManagementGroup,
            adr_models::ManagementGroupAction,
        ),
    ) -> Self {
        ManagementActionSpecification {
            name_tuple: (val.0.name.clone(), val.1.name.clone()),
            action_configuration: val.1.action_configuration,
            action_type: val.1.action_type,
            target_uri: val.1.target_uri,
            timeout_in_seconds: val.1.timeout_in_seconds,
            topic: val.1.topic,
            name: val.1.name,
            type_ref: val.1.type_ref,
            management_group: ManagementGroupSpecification {
                default_timeout_in_seconds: val.0.default_timeout_in_seconds,
                default_topic: val.0.default_topic,
                management_group_configuration: val.0.management_group_configuration,
                data_source: val.0.data_source,
                name: val.0.name,
                type_ref: val.0.type_ref,
            },
        }
    }
}

/// Represents the specification of a Management Group in the Azure Device Registry service.
#[derive(Debug, Clone, PartialEq)]
pub struct ManagementGroupSpecification {
    /// Reference to a data source for a given management group.
    pub data_source: Option<String>,
    /// Default response timeout for all actions that are part of the management group.
    pub default_timeout_in_seconds: Option<u64>,
    /// Default MQTT topic path on which a client will receive the request for all actions that are part of the management group.
    pub default_topic: Option<String>,
    /// Stringified JSON that contains connector-specific configuration for the management group.
    pub management_group_configuration: Option<String>,
    /// Name of the management group.
    pub name: String,
    /// URI or type definition ID.
    pub type_ref: Option<String>,
}

/// Represents the specification of an Event and its Event Group in the Azure Device Registry service.
#[derive(Debug, Clone, PartialEq)]
pub struct EventSpecification {
    /// Destinations for an event.
    pub destinations: Vec<adr_models::EventStreamDestination>, // if None on generated model, we can represent as empty vec. Can currently only be length of 1
    /// Stringified JSON that contains connector-specific configuration for the specific event.
    pub event_configuration: Option<String>,
    /// Reference to a data source for a given event.
    pub data_source: Option<String>,
    /// The name of the event.
    pub name: String,
    /// URI or type definition ID.
    pub type_ref: Option<String>,
    /// The Event Group this event belongs to.
    pub event_group: EventGroupSpecification,
    name_tuple: (String, String),
}

impl EventSpecification {
    pub(crate) fn hash_name(&self) -> &(String, String) {
        &self.name_tuple
    }
}

impl From<(adr_models::EventGroup, adr_models::Event)> for EventSpecification {
    fn from(val: (adr_models::EventGroup, adr_models::Event)) -> Self {
        EventSpecification {
            name_tuple: (val.0.name.clone(), val.1.name.clone()),
            destinations: val.1.destinations,
            event_configuration: val.1.event_configuration,
            data_source: val.1.data_source,
            name: val.1.name,
            type_ref: val.1.type_ref,
            event_group: EventGroupSpecification {
                default_destinations: val.0.default_destinations,
                event_group_configuration: val.0.event_group_configuration,
                data_source: val.0.data_source,
                name: val.0.name,
                type_ref: val.0.type_ref,
            },
        }
    }
}

/// Represents the specification of an Event Group in the Azure Device Registry service.
#[derive(Debug, Clone, PartialEq)]
pub struct EventGroupSpecification {
    /// Default destinations for an event on this Event Group.
    pub default_destinations: Vec<adr_models::EventStreamDestination>, // if None on generated model, we can represent as empty vec. Can currently only be length of 1
    /// Stringified JSON that contains connector-specific configuration for the specific event group.
    pub event_group_configuration: Option<String>,
    /// The address of the notifier of the event in the asset (e.g. URL) so that a client can access the notifier on the asset.
    pub data_source: Option<String>,
    /// The name of the event group.
    pub name: String,
    /// URI or type definition ID.
    pub type_ref: Option<String>,
}

/// Holds the `DataOperation`'s definition, regardless of the type
#[derive(Debug, Clone, PartialEq)]
pub enum DataOperationDefinition {
    /// Dataset definition
    Dataset(adr_models::Dataset),
    /// Event definition
    Event(EventSpecification),
    /// Stream definition
    Stream(adr_models::Stream),
}

impl DataOperationDefinition {
    /// Returns the [`DataOperationName`] of the data operation
    #[must_use]
    pub fn name(&self) -> DataOperationName {
        match self {
            DataOperationDefinition::Dataset(dataset) => DataOperationName::Dataset {
                name: dataset.name.clone(),
            },
            DataOperationDefinition::Event(event) => DataOperationName::Event {
                name: event.name.clone(),
                event_group_name: event.event_group.name.clone(),
            },
            DataOperationDefinition::Stream(stream) => DataOperationName::Stream {
                name: stream.name.clone(),
            },
        }
    }

    /// Returns the [`DataOperationKind`] of the data operation definition
    #[must_use]
    pub fn kind(&self) -> DataOperationKind {
        match self {
            DataOperationDefinition::Dataset(_) => DataOperationKind::Dataset,
            DataOperationDefinition::Event(_) => DataOperationKind::Event,
            DataOperationDefinition::Stream(_) => DataOperationKind::Stream,
        }
    }
}

/// A trait for handling different types of data operations generically.
///
/// The `DataOperation` trait provides a way to interact with data operations
/// (datasets, events, and streams) without needing to know their specific type.
/// This abstraction is useful for scenarios where operations need to be performed
/// uniformly across different data operation types.
///
/// Unlike directly implementing methods on the `DataOperationDefinition` enum,
/// this trait allows individual data operation types (e.g., `Dataset`, `Event`, `Stream`)
/// to define their own behavior while still conforming to a common interface.
trait DataOperation: Into<DataOperationDefinition> {
    type HashName: PartialEq + Eq + Hash + Clone + std::fmt::Debug;
    fn kind() -> DataOperationKind;
    fn hash_name(&self) -> &Self::HashName;
    fn data_operation_name(&self) -> DataOperationName;
}

impl From<adr_models::Dataset> for DataOperationDefinition {
    fn from(val: adr_models::Dataset) -> Self {
        DataOperationDefinition::Dataset(val)
    }
}
impl DataOperation for adr_models::Dataset {
    type HashName = String;
    fn kind() -> DataOperationKind {
        DataOperationKind::Dataset
    }
    fn hash_name(&self) -> &Self::HashName {
        &self.name
    }
    fn data_operation_name(&self) -> DataOperationName {
        DataOperationName::Dataset {
            name: self.name.clone(),
        }
    }
}

impl From<EventSpecification> for DataOperationDefinition {
    fn from(val: EventSpecification) -> Self {
        DataOperationDefinition::Event(val)
    }
}

impl DataOperation for EventSpecification {
    type HashName = (String, String);
    fn kind() -> DataOperationKind {
        DataOperationKind::Event
    }
    fn hash_name(&self) -> &Self::HashName {
        self.hash_name()
    }

    fn data_operation_name(&self) -> DataOperationName {
        DataOperationName::Event {
            name: self.name.clone(),
            event_group_name: self.event_group.name.clone(),
        }
    }
}

impl From<adr_models::Stream> for DataOperationDefinition {
    fn from(val: adr_models::Stream) -> Self {
        DataOperationDefinition::Stream(val)
    }
}
impl DataOperation for adr_models::Stream {
    type HashName = String;
    fn kind() -> DataOperationKind {
        DataOperationKind::Stream
    }
    fn hash_name(&self) -> &Self::HashName {
        &self.name
    }
    fn data_operation_name(&self) -> DataOperationName {
        DataOperationName::Stream {
            name: self.name.clone(),
        }
    }
}

fn adr_error_into_retry_error(
    e: azure_device_registry::Error,
    operation_for_log: &str,
) -> RetryError<azure_device_registry::Error> {
    match e.kind() {
        // network/retriable
        azure_device_registry::ErrorKind::AIOProtocolError(_) => {
            log::warn!("{operation_for_log} failed. Retrying: {e}");
            RetryError::transient(e)
        }
        // indicates an error in the configuration, might be transient in the future depending on what it can indicate
        azure_device_registry::ErrorKind::ServiceError(_)
        // DuplicateObserve indicates an sdk bug where we called observe more than once. Only possible for unobserve calls.
        // Because we unobserve on deletion, this should not happen, but because unobserve can possibly fail, don't panic on this
        | azure_device_registry::ErrorKind::DuplicateObserve(_) => RetryError::permanent(e),
        _ => {
            // InvalidRequestArgument shouldn't be possible since timeout is already validated
            // ValidationError shouldn't be possible since we should never have an empty asset name. It's not possible to be returned for device calls.
            // ShutdownError isn't possible for this fn to return
            unreachable!()
        }
    }
}

#[cfg(test)]
mod tests {
    use test_case::test_case;

    use super::*;

    #[test_case(None, 1, false, true; "new")]
    #[test_case(Some(azure_device_registry::ConfigStatus {
            version: Some(2),
            ..Default::default()
        }), 2, false, true; "version_match")]
    #[test_case(Some(azure_device_registry::ConfigStatus {
            version: Some(1),
            ..Default::default()
        }), 2,
        false,
        false; "version_mismatch")]
    #[test_case(Some(azure_device_registry::ConfigStatus {
            version: Some(2),
            ..Default::default()
        }), 2, true, true; "version_match_with_dataset")]
    #[test_case(Some(azure_device_registry::ConfigStatus {
            version: Some(1),
            ..Default::default()
        }), 2,
        true,
        false; "version_mismatch_with_dataset")]
    fn get_current_asset_status_asset(
        config: Option<azure_device_registry::ConfigStatus>,
        new_spec_version: u64,
        datasets_set: bool,
        expect_keep_received: bool,
    ) {
        let current_status = if datasets_set {
            adr_models::AssetStatus {
                config,
                datasets: Some(vec![adr_models::DatasetEventStreamStatus {
                    name: "test".to_string(),
                    message_schema_reference: None,
                    error: None,
                }]),
                ..Default::default()
            }
        } else {
            adr_models::AssetStatus {
                config,
                ..Default::default()
            }
        };

        let new_status_base =
            AssetClient::get_current_asset_status(&current_status, Some(new_spec_version))
                .into_owned();
        if expect_keep_received {
            assert_eq!(new_status_base, current_status);
        } else {
            assert_eq!(new_status_base, adr_models::AssetStatus::default());
        }
    }

    #[test_case(None, 1, Some(1), true; "no_data_operation")]
    #[test_case(Some(adr_models::DatasetEventStreamStatus {
            name: "test".to_string(),
            message_schema_reference: None,
            error: None,
    }), 1, Some(1), true; "version_match")]
    #[test_case(Some(adr_models::DatasetEventStreamStatus {
            name: "test".to_string(),
            message_schema_reference: None,
            error: None,
    }), 1, Some(2), false; "version_mismatch")]
    fn get_current_asset_status_data_operation(
        data_operation_status: Option<adr_models::DatasetEventStreamStatus>,
        current_version: u64,
        new_spec_version: Option<u64>,
        expect_keep_received: bool,
    ) {
        let datasets = data_operation_status.map(|status| vec![status]);
        let current_status = adr_models::AssetStatus {
            config: Some(azure_device_registry::ConfigStatus {
                version: Some(current_version),
                ..Default::default()
            }),
            datasets,
            ..Default::default()
        };

        let new_status_base =
            AssetClient::get_current_asset_status(&current_status, new_spec_version).into_owned();

        if expect_keep_received {
            assert_eq!(new_status_base, current_status);
        } else {
            assert_eq!(new_status_base, adr_models::AssetStatus::default());
        }
    }

    #[test_case(None, None, true; "new")]
    #[test_case(Some(azure_device_registry::ConfigStatus {
            version: Some(1),
            ..Default::default()
        }), Some(1), true; "version_match")]
    #[test_case(Some(azure_device_registry::ConfigStatus {
            version: Some(1),
            ..Default::default()
        }), Some(2), false; "version_mismatch")]
    fn get_current_device_endpoint_status_device(
        config: Option<azure_device_registry::ConfigStatus>,
        adr_version: Option<u64>,
        expect_keep_received: bool,
    ) {
        let current_status = DeviceEndpointStatus {
            config,
            inbound_endpoint_status: None,
        };

        let new_status_base = current_status
            .get_current_device_endpoint_status(adr_version)
            .into_owned();

        if expect_keep_received {
            assert_eq!(new_status_base, current_status);
        } else {
            assert_eq!(new_status_base, DeviceEndpointStatus::default());
        }
    }

    #[test_case(None, 1, Some(1), true; "no_endpoint")]
    #[test_case(Some(Ok(())), 1, Some(1), true; "version_match")]
    #[test_case(Some(Ok(())), 1, Some(2), false; "version_mismatch")]
    fn get_current_device_endpoint_status_endpoint(
        endpoint_status: Option<Result<(), AdrConfigError>>,
        current_version: u64,
        new_spec_version: Option<u64>,
        expect_keep_received: bool,
    ) {
        let current_status = DeviceEndpointStatus {
            config: Some(azure_device_registry::ConfigStatus {
                version: Some(current_version),
                ..Default::default()
            }),
            inbound_endpoint_status: endpoint_status,
        };

        let new_status_base = current_status
            .get_current_device_endpoint_status(new_spec_version)
            .into_owned();

        if expect_keep_received {
            assert_eq!(new_status_base, current_status);
        } else {
            assert_eq!(new_status_base, DeviceEndpointStatus::default());
        }
    }
}
