# SPDX-License-Identifier: MIT
# Copyright (C) 2024 Avnet
# Authors: Nikola Markovic <nikola.markovic@avnet.com> and Zackary Andraka <zackary.andraka@avnet.com> et al.

import random
import time
from dataclasses import dataclass, field
from datetime import datetime, timezone
from pathlib import Path
from ssl import SSLError
from typing import Callable, Optional, List, Dict, TYPE_CHECKING

from avnet.iotconnect.sdk.sdklib.dra import DeviceRestApi, AwsCredentialsResponse, DeviceIdentityData
from avnet.iotconnect.sdk.sdklib.error import C2DDecodeError, NotSupportedError, ClientError
from avnet.iotconnect.sdk.sdklib.mqtt import C2dOta, C2dMessage, C2dCommand, C2dAck, TelemetryRecord, TelemetryValueType, encode_telemetry_records, encode_c2d_ack, decode_c2d_message
from avnet.iotconnect.sdk.sdklib.util import Timing
from paho.mqtt.client import CallbackAPIVersion, MQTTErrorCode, DisconnectFlags, MQTTMessageInfo
from paho.mqtt.client import Client as PahoClient
from paho.mqtt.reasoncodes import ReasonCode

from .config import DeviceConfig

if TYPE_CHECKING:
    import boto3
    from mypy_boto3_sts import STSClient
    from botocore.exceptions import BotoCoreError
else:
    try:
        import boto3
    except ImportError:
        boto3 = None

class Callbacks:
    def __init__(
            self,
            command_cb: Optional[Callable[[C2dCommand], None]] = None,
            ota_cb: Optional[Callable[[C2dOta], None]] = None,
            disconnected_cb: Optional[Callable[[str, bool], None]] = None,
            generic_message_callbacks: dict[int, Callable[[C2dMessage, dict], None]] = None,
            vs_cb: Optional[Callable[['KvsClient'], None]] = None
    ):
        """
        Specify callbacks for C2D command, OTA, Video Streaming or MQTT disconnection.

        :param command_cb: Callback function with first parameter being C2dCommand object.
            Use this callback to process commands sent by the back end.

        :param ota_cb: Callback function with first parameter being C2dOta object.
            Use this callback to process OTA updates sent by the back end.

        :param disconnected_cb: Callback function with first parameter being string with reason for disconnect
            and second a boolean indicating whether a disconnect request was received from the server.
            Use this callback to asynchronously react to the back end disconnection event rather than polling Client.is_connected.

        :param generic_message_callbacks: A dictionary of callbacks indexed by the message type.

        :param vs_cb: Callback function for video streaming related events.
        """

        self.disconnected_cb = disconnected_cb
        self.command_cb = command_cb
        self.ota_cb = ota_cb
        self.generic_message_callbacks = generic_message_callbacks or dict[int, C2dMessage]()
        self.vs_cb = vs_cb

class AwsCredentialsProvider:
    def __init__(self, dra: DeviceRestApi, credentials_endpoint: str, verbose: bool = False):
        """
        This class provides a way to obtain and refresh AWS temporary credentials.

        If auto_obtain is set to True, the first call to get_credentials() will automatically
        obtain the credentials if they are not already obtained.

        If refresh_before_expiry_secs is set, get_credentials() will be refrsh accordingly.
        """
        self._dra = dra
        self.credentials_endpoint = credentials_endpoint
        self.verbose = False
        self.credentials: Optional[AwsCredentialsResponse] = None

    def obtain_credentials(self) -> AwsCredentialsResponse:
        if self.verbose and self.credentials is not None:
            print("Refreshing AWS credentials. Current expiry secs: %d" % self.get_secs_to_expiry())
        self.credentials = self._dra.get_aws_credentials(credentials_endpoint=self.credentials_endpoint)
        return self.credentials

    def get_secs_to_expiry(self) -> float:
        if self.credentials is None:
            return 0.0
        delta = self.credentials.expiration - datetime.now(timezone.utc)
        return delta.total_seconds()


    def get_credentials(self, refresh_if_secs_to_expiry: Optional[int] = None, env: Optional[Dict[str, str]] = None) -> Optional[AwsCredentialsResponse]:
        """
        Returns the current AWS temporary credentials or NONE if they expired.

        If refresh_if_secs_to_expiry is provided current credentials will be refreshed if they are expired.

        If env dictionary is provided, the AWS credentials will be set in the dictionary. This can be
        useful when invoking shell commands.
        """
        if self.credentials is None:
            return None

        if refresh_if_secs_to_expiry is not None:
            if self.get_secs_to_expiry() < refresh_if_secs_to_expiry:
                print("Refreshing AWS credentials before expiry.")
                return self.obtain_credentials()

        if self.get_secs_to_expiry() <= 0:
            if self.verbose:
                print("AWS credentials have expired.")
            return None

        if env is not None:
            env["AWS_ACCESS_KEY_ID"] = self.credentials.access_key_id
            env["AWS_SECRET_ACCESS_KEY"] = self.credentials.secret_access_key
            env["AWS_SESSION_TOKEN"] = self.credentials.session_token
        return self.credentials


class KvsClient(AwsCredentialsProvider):
    def __init__(
            self,
            dra: DeviceRestApi,
            identity_data: DeviceIdentityData,
            verbose: bool = False
    ):
        """
        This class provides KVS-related information for the device and a way to obtain AWS temporary credentials.

        The auto_start property indicates whether KVS streaming should be started automatically.
        the is_streaming property indicates the current desired status of KVS streaming.
        """
        self.identity_data = identity_data
        if not self.identity_data.vs or not self.identity_data.vs.url:
            raise NotSupportedError("KVS is not enabled for this device")

        super().__init__(dra, self.identity_data.vs.url, verbose=verbose)

        self._is_auto_start = self.identity_data.vs.as_
        self._signaling_channel_arn = self.identity_data.vs.carn
        self._is_streaming = self._is_auto_start

    def is_auto_start(self) -> bool:
        """ Indicates whether KVS streaming should be started automatically """
        return self._is_auto_start

    def is_streaming(self) -> bool:
        """ Indicates whether KVS streaming is currently streaming """
        return self._is_streaming
    def get_signaling_channel_arn(self) -> bool:
        """ Returns the signaling channel ARN for WebRTC connection to KVS if available."""
        return self._signaling_channel_arn


@dataclass
class S3BucketInfo:
    bucket_name: str = field(default=None)
    is_customer_owned: bool = field(default=False)
    role_arn: str = field(default=None)


class S3Client(AwsCredentialsProvider):
    """
    This class provides S3 access information for the device and a way to obtain AWS temporary credentials.

    When is_customer_owned for a bucket is True, role_arn will be need to be used to make an assume role request with
    the credentials from AwsCredentialsProvider. Those new temporary credentials can then be used to access the actual S3 bucket.
    Obtaining these customer bucket temporary credentials is outside of scope of this SDK at the moment.
    """

    def __init__(
            self,
            dra: DeviceRestApi,
            identity_data: DeviceIdentityData
    ):
        self.identity_data = identity_data
        if not self.identity_data.filesystem or not self.identity_data.filesystem.url:
            raise NotSupportedError("Filesystem (S3) support is not enabled for this device")

        super().__init__(dra, self.identity_data.filesystem.url)

        self.buckets = []
        self.buckets = [
            S3BucketInfo(bucket_name=bucket.bn,is_customer_owned=bucket.ca,role_arn=bucket.rarn)
            for bucket in self.identity_data.filesystem.buckets
        ]

    def get_buckets(self) -> List[S3BucketInfo]:
        return self.buckets

    def get_default_bucket(self) -> S3BucketInfo:
        """
        Returns the first non-customer owned bucket that should be used for device uploads.
        Otherwise, returns the first bucket available.
        Normally there should be at least one bucket that is not customer-owned,
        but if "Sagemaker Support" is enabled in your account settings, you will get customer owned bucket
        In this case we return the first bucket.
        """
        first_bucket = None
        for bucket in self.get_buckets():
            if first_bucket is None:
                first_bucket = bucket
            # the first bucket that's not customer owned should be the default bucket
            # that can be used to upload files and show then in telemetry UI
            if not bucket.is_customer_owned:
                return bucket
        return first_bucket # or nothing if none available

    def upload_to_bucket(self, local_path: str, bucket_path: str, bucket: Optional[S3BucketInfo] = None):
        """
        Upload a file to S3 using the device's S3 credentials.

        :param local_path: Path to the local file to upload.
        :param bucket_path: (Optional) Full object path (Key) in the S3 bucket where the file will be uploaded.
            If not provided, the file will be uploaded to the root of the bucket with the same file name as local_path with prepended epoch timestamp.
        :param bucket: (Optional) S3BucketInfo object representing the bucket to upload to.
            If not provided, the first non-customer owned bucket will be used. Otherwise, the first bucket will be used.
        """
        if boto3 is None:
            raise NotSupportedError("S3: Optional package is required. Install this package with pip3 install iotconnect-lite-sdk[aws-s3]")

        if bucket is None:
            # will raise if no suitable bucket found
            bucket = self.get_default_bucket()
            if bucket is None:
                raise ClientError("No S3 bucket available for file uploads")

        creds = self.get_credentials(refresh_if_secs_to_expiry=60)
        if creds is None:
            # will raise if unable to obtain new credentials
            creds = self.obtain_credentials()

        if bucket.is_customer_owned:
            # do STS assume role with rarn here to obtain new credentials for the customer bucket
            # the type import and forward declaration is only for type checking
            try:
                sts_client: 'STSClient' = boto3.client(
                    'sts',
                    aws_access_key_id=creds.access_key_id,
                    aws_secret_access_key=creds.secret_access_key,
                    aws_session_token=creds.session_token
                )
                assumed_role = sts_client.assume_role(
                    RoleArn=bucket.role_arn,
                    RoleSessionName="iotconnect-lite-sdk-s3-upload-" + self.identity_data.client_id + "-" + str(int(time.time()))
                )
                creds = AwsCredentialsResponse(
                    access_key_id=assumed_role['Credentials']['AccessKeyId'],
                    secret_access_key=assumed_role['Credentials']['SecretAccessKey'],
                    session_token=assumed_role['Credentials']['SessionToken'],
                    expiration_str= assumed_role['Credentials']['Expiration'].isoformat()
                )
            except BotoCoreError as e:
                raise ClientError(f"Failed to assume role {bucket.role_arn}")
        try:
            s3 = boto3.client(
                's3',
                aws_access_key_id=creds.access_key_id,
                aws_secret_access_key=creds.secret_access_key,
                aws_session_token=creds.session_token
            )
        except BotoCoreError as e:
            raise ClientError("Failed to create S3 client: " + str(e))

        try:
            s3.upload_file(
                Filename=local_path,
                Bucket=bucket.bucket_name,
                Key=bucket_path
            )
            if self.verbose:
                print(f"File {local_path} uploaded to bucket {bucket.bucket_name}")
        except BotoCoreError as e:
            raise ClientError(f"Failed to upload file to S3: {str(e)}")

class ClientSettings:
    """ Optional settings that the user can use to control the client MQTT connection behavior"""

    def __init__(
            self,
            verbose: bool = True,
            connect_timeout_secs: int = 30,
            connect_tries: int = 100,
            connect_backoff_max_secs: int = 15
    ):
        if verbose:
            from . import __version__ as SDK_VERSION
            from avnet.iotconnect.sdk.sdklib import __version__ as LIB_VERSION
            print(f"/IOTCONNECT Lite Client started with version {SDK_VERSION} and Lib version {LIB_VERSION}")
        self.verbose = verbose
        self.connect_timeout_secs = connect_timeout_secs
        self.connect_tries = connect_tries
        self.connect_backoff_max_secs = connect_backoff_max_secs
        if connect_timeout_secs < 1:
            raise ValueError("connect_timeout_secs must be greater than 1")
        if connect_tries < 1:
            raise ValueError("connect_tries must be greater than 1")


class Client:
    """
    This is an Avnet /IOTCONNECT MQTT client that provides an easy way for the user to
    connect to /IOTCONNECT, send and receive data.

    :param config: Required device configuration. See the examples or DeviceConfig class description for more details.
    :param callbacks: Optional callbacks that can be provided for the following events:
        - /IOTCONNECT C2D (Cloud To Device) commands that you can send to your device using the /IOTCONNECT
        - /IOTCONNECT OTA update events.
        - Device MQTT disconnection.
    :param settings: Tune the client behavior by providing your preferences regarding connection timeouts and logging.

    Usage (see basic-example.py or minimal.py examples at https://github.com/avnet-iotconnect/iotc-python-lite-sdk for more details):

    - Construct this client class:
        - Provide your device information and x509 credentials (certificate and private key)
            Either provide the path the downloaded iotcDeviceConfig.json and use the DeviceConfig.from_iotc_device_config_json_file()
            class method. You can download the iotcDeviceConfig.json by clicking the cog icon in the upper right of your device's info panel.

            Or you can also provide the device parameters directly:

            device_config = DeviceConfig(
                platform="aws",                             # The IoTconnect IoT platform - Either "aws" for AWS IoTCore or "az" for Azure IoTHub
                env="your-environment",                     # Your account environment. You can locate this in you /IOTCONNECT web UI at Settings -> Key Value
                cpid="ABCDEFG123456",                       # Your account CPID (Company ID). You can locate this in you /IOTCONNECT web UI at Settings -> Key Value
                duid="my-device",                           # Your device unique ID
                device_cert_path="path/to/device-cert.pem", # Path to the device certificate file
                device_pkey_path="path/to/device-pkey.pem"  # Path to the device private key file
            )
            NOTE: If you do not pass the server certificate, we will use the system's trusted certificate store, if available.
            For example, the trusted Root CA certificates from the in /etc/ssl/certs will be used on Linux.
            However, it is more secure to pass the actual server CA Root certificate in order to avoid potential MITM attacks.
            On Linux, you can use server_ca_cert_path="/etc/ssl/certs/DigiCert_Global_Root_CA.pem" for Azure
            or server_ca_cert_path="/etc/ssl/certs/Amazon_Root_CA_1.pem" for AWS

        - (Optional) provide callbacks for C2D Commands or device disconnect.
            See the basic-example.py at https://github.com/avnet-iotconnect/iotc-python-lite-sdk example for details.

        - (Optional) provide ClientSettings to tune the client behavior.

        It is recommended to surround the DeviceConfig constructor and the Client constructor
        with a try/except block catching DeviceConfigError and printing it to get more information
        to be able to troubleshoot configuration errors.


    - Call Client.connect(). The call will block until successfully connected or timed out. For example:

    - Ensure that the Client.is_connected() periodically or react to a disconnect event callback
        and attempt to reconnect or perform a different action appropriate to your application. For example:
        if not client.is_connected:
            client.connect()

    - Send messages with Client.send_telemetry() or send_telemetry_records().
        See basic-example.py example at https://github.com/avnet-iotconnect/iotc-python-lite-sdk for more info.
        For example:
            c.send_telemetry({
                'temperature': get_sensor_temperature()
            })

    """

    def __init__(
            self,
            config: DeviceConfig,
            callbacks: Callbacks = None,
            settings: ClientSettings = None
    ):
        self.user_callbacks = callbacks or Callbacks()
        self.settings = settings or ClientSettings()
        self.config = config

        self._dra = DeviceRestApi(
            config = config.to_properties(),
            tls_credentials= config.to_tls_credentials(),
            verbose=self.settings.verbose)
        self._identity_data = self._dra.get_identity_data()  # can raise DeviceConfigError

        self.mqtt = PahoClient(
            callback_api_version=CallbackAPIVersion.VERSION2,
            client_id=self._identity_data.client_id
        )
        # Investigate if we need to make this configurable:
        self.mqtt.reconnect_delay_set(min_delay=1, max_delay=int(self.settings.connect_timeout_secs / 2 + 1))
        self.mqtt.tls_set(certfile=config.device_cert_path, keyfile=config.device_pkey_path, ca_certs=config.server_ca_cert_path)
        self.mqtt.username = self._identity_data.username

        self.mqtt.on_message = self._on_mqtt_message
        self.mqtt.on_connect = self._on_mqtt_connect
        self.mqtt.on_disconnect = self._on_mqtt_disconnect
        self.mqtt.on_publish = self._on_mqtt_publish

        self.user_callbacks = callbacks or Callbacks()

        self._kvs_client: Optional[KvsClient] = None
        self._s3_client: Optional[S3Client] = None

        # If any of these are not available, then the template doesn't enable the features
        # Most of the devices will not have those supported, so fail silently.
        try:
            self._kvs_client = KvsClient(self._dra, self._identity_data)
        except NotSupportedError:
            pass

        try:
            self._s3_client = S3Client(self._dra, self._identity_data)
        except NotSupportedError:
            pass

    @classmethod
    def timestamp_now(cls) -> datetime:
        """ Returns the UTC timestamp that can be used to stamp telemetry records """
        return datetime.now(timezone.utc)

    def is_connected(self):
        return self.mqtt.is_connected()

    def connect(self):
        def wait_for_connection() -> bool:
            connect_timer = Timing()
            if self.settings.verbose:
                print("waiting to connect...")
            while True:
                if self.is_connected():
                    if self.settings.verbose:
                        print("MQTT connected")
                    return True
                time.sleep(0.5)
                if connect_timer.diff_now().seconds > self.settings.connect_timeout_secs:
                    print("Timed out.")
                    self.disconnect()
                    return False

        if self.is_connected():
            return

        for i in range(1, self.settings.connect_tries):
            try:
                t = Timing()
                mqtt_error = self.mqtt.connect(
                    host=self._identity_data.host,
                    port=8883
                )
                if mqtt_error != MQTTErrorCode.MQTT_ERR_SUCCESS:
                    print("TLS connection to the endpoint failed")
                else:
                    print("Awaiting MQTT connection establishment...")
                    self.mqtt.loop_start()
                    if wait_for_connection():
                        if self.settings.verbose:
                            print("Connected in %dms" % (t.diff_now().microseconds / 1000))
                        break
                    else:
                        continue

            except (SSLError, TimeoutError, OSError) as ex:
                # OSError includes socket.gaierror when host could not be resolved
                # This could also be temporary, so keep trying
                print("Failed to connect to host %s. Exception: %s" % (self._identity_data.host, str(ex)))

            backoff_ms = random.randrange(1000, self.settings.connect_backoff_max_secs * 1000)
            print("Retrying connection... Backing off for %d ms." % backoff_ms)
            # Jitter back off a random number of milliseconds between 1 and 10 seconds.
            time.sleep(backoff_ms / 1000)

        self.mqtt.subscribe(self._identity_data.topics.c2d, qos=1)

    def disconnect(self) -> MQTTErrorCode:
        ret = self.mqtt.disconnect()
        if self.settings.verbose:
            print("Disconnected.")
        return ret

    def send_telemetry(self, values: dict[str, TelemetryValueType], timestamp: datetime = None):
        """ Sends a single telemetry dataset.
        If you need gateway/child functionality or need to send multiple value sets in one packet,
        use the send_telemetry_records() method.

        :param TelemetryValues values:
            The name-value telemetry pairs to send. Each value can be
                - a primitive value: Maps directly to a JSON string, number or boolean
                - None: Maps to JSON null,
                - Tuple[float, float]: Used to send a lat/long geographic coordinate as decimal degrees as an
                    array of two (positive or negative) floats.
                    For example, [44.787197, 20.457273] is the geo coordinate Belgrade in Serbia,
                    where latitude 44.787197 is a positive number indicating degrees north,
                    and longitude a positive number as well, indicating degrees east.
                    Maps to JSON array of two elements.
                - Another hash with possible values above when sending an object. Maps to JSON object.
            in case when an object needs to be sent.
        :param datetime timestamp: (Optional) The timestamp corresponding to this dataset.
            If not provided, this will save bandwidth, as no timestamp will not be sent over MQTT.
             The server receipt timestamp will be applied to the telemetry values in this telemetry record.
             Supply this value (using Client.timestamp()) if you need more control over timestamps.
        """
        self.send_telemetry_records([TelemetryRecord(
            values=values,
            timestamp=timestamp
        )])

    def send_telemetry_records(self, records: list[TelemetryRecord]) -> Optional[MQTTMessageInfo]:
        """
        A complex, but more powerful way to send telemetry.
        It allows the user to send multiple sets of telemetry values
        and control the timestamp of each telemetry value set.
        Supports gateway devices with gateway-child relationship by allowing
        the user to set the parent/child unique_id ("id" in JSON)
        and tag of respective parent./child ("tg" in JSON)

        See https://docs.iotconnect.io/iotconnect/sdk/message-protocol/device-message-2-1/d2c-messages/#Device for more information.
        """

        if not self.is_connected():
            print('Message NOT sent. Not connected!')
            return None
        else:
            packet = encode_telemetry_records(records)
            ret = self.mqtt.publish(
                topic=self._identity_data.topics.rpt,
                qos=1,
                payload=packet
            )
            if self.settings.verbose:
                print(">", packet)
            return ret

    def s3_upload(self, local_path: str, custom_values: Optional[dict[str, TelemetryValueType]] = None, relative_upload_path: Optional[str] = None):
        """
        Upload a file into the default file upload bucket send a telemetry message with the file URL.

        This method is intended for general use. Use other provided methods if you need more control.

        This will call S3Client.upload_to_bucket() and then send_file_upload_message().

        :param local_path: Path to the local file to upload.
        :param custom_values: (Optional) Additional telemetry values to send along with the file URL.
            Do not populate the "url" field in it. The field is reserved.
            If you populate the "cf" key, the values will appear in /IOTCONNECT Telemetry Files
            The web UI recognize the file type and show it appropriately.
            Example:
                'cf': {
                    'class': 'dog',
                    'confidence': 0.700,
                }
            or just:
                'cf': 'dog'
        :param relative_upload_path: Sub-path in the S3 bucket where the file will be uploaded.
            This path should not contain the 'device-uploads/<client_id>' (thing name), but have only the relative path after that.
        """
        if boto3 is None:
            raise NotSupportedError("S3: Optional package is required. Install this package with pip3 install iotconnect-lite-sdk[aws-s3]")
        if self._s3_client is None:
            raise NotSupportedError("FS support is not enabled for this device. Please ensure to enable it in the device template or ensure to re-create the device after enabling it.")

        bucket = self._s3_client.get_default_bucket()
        if bucket is None:
            raise ClientError("No S3 bucket available for file uploads")

        if relative_upload_path is None or len(relative_upload_path) == 0:
            file_name = Path(local_path).name
            unix_timestamp = int(datetime.now(timezone.utc).timestamp())
            relative_upload_path = f"{unix_timestamp}-{file_name}"

        self._s3_client.upload_to_bucket(local_path, f"device-uploads/{self.get_client_id()}/{relative_upload_path}", bucket)

        self.send_file_upload_message(
            relative_file_upload_path=f"{relative_upload_path}",
            custom_values=custom_values
        )

    def send_file_upload_message(self, relative_file_upload_path: str, custom_values: Optional[dict[str, TelemetryValueType]] = None):
        """
        Send a file upload MQTT message indicating that a file has been uploaded to the device's S3 bucket.
        The Telemetry Files tab in the /IOTCONNECT web UI will show the uploaded file accordingly.

        :param relative_file_upload_path: The relative path in the device-uploads S3 bucket where the file was uploaded.
            This path should not contain the 'device-uploads/<client_id>' (thing name), but have only the relative path after that.
        :param custom_values: (Optional) Additional telemetry values to send along with the file URL.
            Do not populate the "url" field in it. The field is reserved.
            If you populate the "cf" key as value or object, the values will appear in /IOTCONNECT Telemetry Files
            The web UI recognize the file type and show it appropriately.
            Example:
                'cf': {
                    'class': 'dog',
                    'confidence': 0.700,
                }
            or just:
                'cf': 'dog'
        """
        if relative_file_upload_path is None or len(relative_file_upload_path) == 0:
            raise ValueError("relative_file_upload_path cannot be empty")

        # the custom values really become the top level telemetry record
        # with "url" added to it as a special field
        if custom_values is None:
            record = dict[str, TelemetryValueType]()
        else:
            record = custom_values

        # validate file path to make sure the user did not make some mistake
        # but only if verbose is enabled?... Maybe the user intended it?
        if "device-uploads" in relative_file_upload_path:
            print("Warning: The 'device_uploads' prefix should not generally be included in the relative_file_upload_path")
        if self.get_duid() in relative_file_upload_path:
            print("Warning: relative_file_upload_path should not generally contain the device unique ID (DUID)")

        # url goes at top level:
        if record.get("url") is not None:
            raise ValueError("The 'url' key in custom_values is reserved and cannot be used in custom_values")
        record["url"] = relative_file_upload_path
        records = [TelemetryRecord(record)]
        if not self.is_connected():
            print('Message NOT sent. Not connected!')
            return None
        else:
            packet = encode_telemetry_records(records)
            ret = self.mqtt.publish(
                topic=self._identity_data.topics.fu,
                qos=1,
                payload=packet
            )
            if self.settings.verbose:
                print("file: >", packet)
            return ret

    def send_command_ack(self, original_message: C2dCommand, status: int, message_str=None):
        """
        Send Command acknowledgement.

        :param original_message: The original message that was received in the callback.
        :param status: C2dAck.CMD_FAILED or C2dAck.CMD_SUCCESS_WITH_ACK.
        :param message_str: (Optional) For example: 'LED color now "blue"', or 'LED color "red" not supported'
        """

        if original_message.type != C2dMessage.COMMAND:
            print('Error: Called send_command_ack(), but message is not a command!')
            return
        self.send_ack(
            ack_id=original_message.ack_id,
            message_type=original_message.type,
            status=status,
            message_str=message_str,
            original_command=original_message.command_name
        )

    def send_ota_ack(self, original_message: C2dOta, status: int, message_str = None):
        """
        Send OTA acknowledgement.
        See the C2dAck comments for best practices with OTA download ACks.

        :param original_message: The original message that was received in the callback.
        :param status: For example: C2dAck.OTA_DOWNLOAD_FAILED.
        :param message_str: (Optional) For example: "Failed to unzip the OTA package".
        """
        if original_message.type != C2dMessage.OTA:
            print('Error: Called send_ota_ack(), but message is not an OTA request!')
            return
        self.send_ack(
            ack_id=original_message.ack_id,
            message_type=original_message.type,
            status=status,
            message_str=message_str
        )

    def send_ack(self, ack_id: str, message_type: int, status: int, message_str: str = None, original_command: str = None):
        """
        Send Command or OTA ack while having only ACK ID

        :param ack_id: Recorded ack_id from the original message.
        :param message_type: For example: C2dMessage.COMMAND or C2dMessage.OTA, .
        :param status: For example: C2dAck.CMD_FAILED or C2dAck.OTA_DOWNLOAD_DONE for command/OTA respectively.
        :param message_str: (Optional) For example: "Failed to unzip the OTA package", or 'LED color "red" not supported'
        :param original_command: (Optional) If this argument is passed,
            this command name will be printed along with any potential error messages.

        While the client should generally use send_ota_ack or send_command_ack, this method can be used in cases
        where the context of the original received message is not available (after OTA restart for example)
        """
        if not self.is_connected():
            print('Message NOT sent. Not connected!')
        elif ack_id is None or len(ack_id) == 0:
            if original_command is not None:
                print('Error: Message ACK ID missing. Ensure to set "Acknowledgement Required" in the template for command %s!' % original_command)
            else:
                print('Error: Message ACK ID missing. Ensure to set "Acknowledgement Required" in the template the command!')
            return None
        elif message_type not in (C2dMessage.COMMAND, C2dMessage.OTA):
            print('Warning: Message type %d does not appear to be a valid message type!' % message_type) # let it pass, just in case we can still somehow send different kind of ack
        elif message_type == C2dMessage.COMMAND and not C2dAck.is_valid_cmd_status(status):
            print('Warning: Status %d does not appear to be a valid command ACK status!' % status) # let it pass, just in case there is a new status
        elif message_type == C2dMessage.OTA and not C2dAck.is_valid_ota_status(status):
            print('Warning: Status %d does not appear to be a valid OTA ACK status!' % status) # let it pass, just in case there is a new status

        packet = encode_c2d_ack(ack_id, message_type, status, message_str)
        ret = self.mqtt.publish(
            topic=self._identity_data.topics.ack,
            qos=1,
            payload=packet
        )
        if self.settings.verbose:
            print(">", packet)
        return ret

    def get_duid(self) -> str:
        """ Convenience method to get device unique ID """
        return self.config.duid

    def get_client_id(self) -> str:
        """ Returns MQTT/Client ID (AWS Thing Name on AWS) """
        return self._identity_data.client_id


    def _process_c2d_message(self, topic: str, payload: str) -> bool:
        # topic is ignored for now as we only subscribe to one
        # we ought to change this once we start supporting Properties (Twin/Shadow)
        try:
            # use the simplest form of ProtocolC2dMessageJson when deserializing first and
            # convert message to appropriate json later

            decoding_result = decode_c2d_message(payload)
            generic_message = decoding_result.generic_message

            # if the user wants to handle this message type, stop processing further
            generic_cb = self.user_callbacks.generic_message_callbacks.get(generic_message.type)
            if generic_cb is not None:
                generic_cb(generic_message, decoding_result.raw_message)
                return True

            if self.user_callbacks.vs_cb is not None:
                if generic_message.type in (C2dMessage.START_STREAM, C2dMessage.STOP_STREAM):
                    print(f"Received {C2dMessage.TYPES.get(generic_message.type)}")
                    self._kvs_client._is_streaming = generic_message.type == C2dMessage.START_STREAM
                    self.user_callbacks.vs_cb(self._kvs_client)
                    return True

            if decoding_result.command is not None:
                # Potential way to deal with runtime qualification, but has issues.
                # if msg.command_name == 'aws-qualification-start':
                #     self._aws_qualification_start(msg.command_args)
                # elif self.user_callbacks.command_cb is not None:
                if self.user_callbacks.command_cb is not None:
                    self.user_callbacks.command_cb(decoding_result.command)
                else:
                    if self.settings.verbose:
                        print("WARN: Unhandled command %s received!" % decoding_result.command.command_name)
            elif decoding_result.ota is not None:
                if self.user_callbacks.ota_cb is not None:
                    self.user_callbacks.ota_cb(decoding_result.ota)
                else:
                    if self.settings.verbose:
                        print("WARN: Unhandled OTA request received!")
            elif generic_message.is_fatal:
                print("Received C2D message %s from backend. Device should stop operation." % generic_message.type_description)
            elif generic_message.needs_refresh:
                print("Received C2D message %s from backend. Device should re-initialize the application." % generic_message.type_description)
            elif generic_message.heartbeat_operation is not None:
                operation_str = "start" if generic_message.heartbeat_operation == True else "stop"
                print("Received C2D message %s from backend. Device should %s heartbeat messages." % (generic_message.type_description, operation_str))
            else:
                print("C2D Message parsing for message type %d is not supported by this client. Message was: %s" % (generic_message.ct, payload))
            return True

        except C2DDecodeError:
            print('C2D Parsing Error: "%s"' % payload)
            return False

    def _on_mqtt_connect(self, mqttc: PahoClient, obj, flags, reason_code, properties):
        if self.settings.verbose:
            print("Connected. Reason Code: " + str(reason_code))

    def _on_mqtt_disconnect(self, mqttc: PahoClient, obj, flags: DisconnectFlags, reason_code: ReasonCode, properties):
        if self.user_callbacks.disconnected_cb is not None:
            # cannot send raw reason code from paho. We could technically change the backend.
            self.user_callbacks.disconnected_cb(str(reason_code), flags.is_disconnect_packet_from_server)
        else:
            print("Disconnected. Reason: %s. Flags: %s" % (str(reason_code), str(flags)))

    def _on_mqtt_message(self, mqttc: PahoClient, obj, msg):
        if self.settings.verbose:
            print(msg.topic + " " + str(msg.qos) + " " + str(msg.payload))
        self._process_c2d_message(msg.topic, msg.payload)

    def _on_mqtt_publish(self, mqttc: PahoClient, obj, mid, reason_code, properties):
        # print("mid: " + str(mid))
        pass

    def _aws_qualification_start(self, command_args: list[str]):
        t = Timing()

        def log_callback(client, userdata, level, buf):
            print("%d [%s]: %s" % (t.diff_now().microseconds / 1000, str(level), buf))
            t.reset(False)

        if len(command_args) >= 1:
            host = command_args[0]
            print("Starting AWS Device Qualification for", host)
            self._identity_data.topics.rpt = 'qualification'
            self._identity_data.topics.c2d = 'qualification'
            self._identity_data.topics.ack = 'qualification'
            self._identity_data.host = host
            self.mqtt.on_log = log_callback
            self.disconnect()
            while True:
                connected_time = Timing()
                if not self.is_connected():
                    print('(re)connecting to', self._identity_data.host)
                    self.connect()
                    connected_time.reset(False)  # reset the timer
                else:
                    if connected_time.diff_now().seconds > 60:
                        print("Stayed connected for too long. resetting the connection")
                        self.disconnect()
                        continue
                    self.send_telemetry({
                        'qualification': 'true'
                    })
                time.sleep(5)

        else:
            print("Malformed AWS qualification command. Missing command argument!")

    def get_kvs_client(self) -> Optional[KvsClient]:
        return self._kvs_client
    
    def get_s3_client(self) -> Optional[S3Client]:
        return self._s3_client
