import logging
import time
from datetime import datetime
from typing import Any, Optional

import boto3
import botocore
from botocore.exceptions import ClientError
from typeguard import typechecked

from psi_kinesis import get_logger

DEFAULT_MAX_RETRY_COUNT = 5  # default max number of retries for Kinesis operations
DEFAULT_RETRY_BACKOFF_RATE = 3.0  # default retry backoff multiplier
PROVISIONED_THROUGHPUT_EXCEPTION = "ProvisionedThroughputExceededException"
VALID_SHARD_ITERATOR_TYPES = [
    "AT_SEQUENCE_NUMBER",
    "AFTER_SEQUENCE_NUMBER",
    "TRIM_HORIZON",
    "LATEST",
    "AT_TIMESTAMP",
]


@typechecked
class KinesisReader:
    """KinesisReader provides read-only utility functions for interacting with Kinesis resources"""

    def __init__(
        self,
        aws_session: Optional[Any] = None,
        max_retry_count: Optional[int] = None,
        retry_backoff_rate: Optional[float] = None,
        logger: Optional[logging.Logger] = None,
    ):
        """Initializes a KinesisReader

        Parameters
        ----------
        aws_session : Any
            (optional) A boto3.session.Session() object used to create boto3 clients.

        max_retry_count : int
            (optional) The max number of retries for Kinesis operations like
            get_records. Default value of 5.

        retry_backoff_rate : float
            (optional) The retry backoff rate multiplier for Kinesis operations like
            get_records. Default value of 3.0.

        logger : logging.Logger
            (optional) The logger used by a KinesisReader object. If unspecified,
            a new logger will be created.
        """
        assert (
            max_retry_count is None or max_retry_count >= 0
        ), f"Max retry count must be greater than or equal to 0"
        assert (
            retry_backoff_rate is None or retry_backoff_rate > 0
        ), f"Retry backoff rate must be greater than 0"
        self.__aws_session = aws_session or boto3.session.Session()
        self.__original_aws_session = self.aws_session
        self.__client = self.aws_session.client("kinesis")
        self.__max_retry_count = (
            DEFAULT_MAX_RETRY_COUNT if max_retry_count is None else max_retry_count
        )
        self.__retry_backoff_rate = retry_backoff_rate or DEFAULT_RETRY_BACKOFF_RATE
        self.__logger = logger or get_logger()

    @property
    def aws_session(self) -> Any:
        return self.__aws_session

    @property
    def client(self) -> Any:
        return self.__client

    @property
    def max_retry_count(self) -> int:
        return self.__max_retry_count

    @property
    def retry_backoff_rate(self) -> float:
        return self.__retry_backoff_rate

    @property
    def logger(self) -> logging.Logger:
        return self.__logger

    def assume_role(
        self,
        role_name: str,
        account_id: str,
        region: str = "us-east-1",
        role_session_name: str = "CrossAccountKinesisRole",
    ):
        """Assume an IAM role to perform subsequent Kinesis actions

        Parameters
        ----------
        role_name : str
            Name of the IAM role to assume

        account_id : str
            AWS account ID where the IAM role resides

        region : str
            (optional) AWS region code to use as the default
            region for subsequent Kinesis actions

        role_session_name : str
            (optional) An identifier for the assumed role session.
            Default value "CrossAccountKinesisRole"
        """
        self.logger.info(
            f"Assuming role arn:aws:iam::{account_id}:role/{role_name} with default region {region}"
        )
        try:
            sts_client = self.__original_aws_session.client("sts")
            role_creds = sts_client.assume_role(
                RoleArn=f"arn:aws:iam::{account_id}:role/{role_name}",
                RoleSessionName=role_session_name,
            )["Credentials"]

            self.__aws_session = boto3.Session(
                aws_access_key_id=role_creds["AccessKeyId"],
                aws_secret_access_key=role_creds["SecretAccessKey"],
                aws_session_token=role_creds["SessionToken"],
                region_name=region,
            )
            self.__client = self.aws_session.client("kinesis", region_name=region)
        except Exception as e:
            self.logger.error(f"Failed to assume role with exception: {e}")
            raise e

    def does_stream_exist(self, stream_name: str) -> bool:
        """Checks if a Kinesis stream exists and returns true/false

        Parameters
        ----------
        stream_name : str
            The name of the stream to look for

        Returns
        -------
        bool
            Whether or not the stream exists
        """
        try:
            self.logger.info(f'Describing stream with name "{stream_name}"')
            resp = self.client.describe_stream(StreamName=stream_name)
            return resp.get("StreamDescription", {}).get("StreamName") == stream_name
        except botocore.exceptions.ClientError as e:
            if e.response["Error"]["Code"] == "ResourceNotFoundException":
                return False
            self.logger.error(f"Failed to describe Kinesis stream with exception: {e}")
            raise e
        except Exception as e:
            self.logger.error(f"Failed to describe Kinesis stream with exception: {e}")
            raise e

    def list_shards(
        self, stream_name: str, shard_filter: Optional[dict[str, Any]] = None
    ) -> list[dict]:
        """list shards for a given Kinesis stream. Handles paginiation.

        Parameters
        ----------
        stream_name : str
            The name of the Kinesis stream to retrieve shards

        shard_filter : dict[str, Any]
            (optional) Allows filtering the response of the
            ListShards API with a dictionary like:

            {
                'Type': 'string',
                'ShardId': 'string',
                'Timestamp': datetime(2015, 1, 1)
            }

        Returns
        -------
        list[dict]
            The list of shards for the given Kinesis stream.
            Each shard is a dictionary like:

            {
                'ShardId': 'string',
                'ParentShardId': 'string',
                'AdjacentParentShardId': 'string',
                'HashKeyRange': {
                    'StartingHashKey': 'string',
                    'EndingHashKey': 'string'
                },
                'SequenceNumberRange': {
                    'StartingSequenceNumber': 'string',
                    'EndingSequenceNumber': 'string'
                }
            }
        """
        try:
            self.logger.info(f'Listing shards for Kinesis stream "{stream_name}"')
            list_shards_input = {"StreamName": stream_name}
            if shard_filter:
                list_shards_input["ShardFilter"] = shard_filter

            shards, has_more_shards = [], True
            while has_more_shards:
                resp = self.client.list_shards(**list_shards_input)
                shards.extend(resp.get("Shards") or [])
                list_shards_input["NextToken"] = resp.get("NextToken")
                has_more_shards = bool(resp.get("NextToken"))
            return shards
        except Exception as e:
            self.logger.error(f"Failed to list shards for Kinesis stream with exception: {e}")
            raise e

    def get_records(
        self,
        stream_name: str,
        shard_id: Optional[str] = None,
        shard_iterator_type: Optional[str] = None,
        starting_sequence_number: Optional[str] = None,
        timestamp: Optional[datetime] = None,
        shard_filter: Optional[dict[str, Any]] = None,
    ) -> dict[str, list]:
        """Retrieve records from a Kinesis stream. By default, retrieves all records
        in the stream.
        You can limit which shards to retrieve records from using
        either shard_id or shard_filter (but not both).
        You can control how records are iterated with shard_iterator_type,
        starting_sequence_number, and timestamp. For more information, see:
        https://docs.aws.amazon.com/kinesis/latest/APIReference/API_GetShardIterator.html

        This function handles Backoff and retry in the event of a
        ProvisionedThroughputExceededException.

        Parameters
        ----------
        stream_name : str
            The name of the Kinesis stream to get records from

        shard_id : str
            (optional) The ID of the shard to get records from.
            If set, shard_filter cannot be specified.

        shard_iterator_type : str
            (optional) Determines how the shard iterator is used to start reading records
            from the shard.
            Defaults to "TRIM_HORIZON"

        starting_sequence_number : str
            (optional) The sequence number of the record in the shard from which to start
            reading.
            If set, timestamp cannot be specified.

        timestamp : datetime
            (optional) The timestamp of the record from which to start reading.
            If set, starting_sequence_number cannot be specified.

        shard_filter : dict[str, Any]
            (optional) Allows filtering the response of the
            ListShards API with a dictionary like:

            {
                'Type': 'string',
                'ShardId': 'string',
                'Timestamp': datetime(2015, 1, 1)
            }

            If set, shard_id cannot be specified

        Returns
        -------
        dict[str, list]
            Dictionary mapping shard ID to list of records with the form:

            {
                'string': [
                    {
                        'SequenceNumber': 'string',
                        'ApproximateArrivalTimestamp': datetime(2015, 1, 1),
                        'Data': b'bytes',
                        'PartitionKey': 'string',
                        'EncryptionType': 'NONE'|'KMS'
                    }
                ]
            }
        """
        self.logger.info(f"Getting records from Kinesis stream: {stream_name}")

        # Validate input
        shard_iterator_type = shard_iterator_type or "TRIM_HORIZON"
        assert not (
            shard_id and shard_filter
        ), "Cannot get records with both a shard ID and a shard filter. Specify one or the other"
        assert not (
            starting_sequence_number and timestamp
        ), "Cannot specify both timestamp and starting sequence number. Specify one or the other"
        assert (
            shard_iterator_type in VALID_SHARD_ITERATOR_TYPES
        ), f'Invalid shard iterator type "{shard_iterator_type}". Must be one of: {VALID_SHARD_ITERATOR_TYPES}'
        if shard_iterator_type in ["AT_SEQUENCE_NUMBER", "AFTER_SEQUENCE_NUMBER"]:
            assert (
                starting_sequence_number
            ), "Starting sequence number must be specified if shard iterator type is one of [AT_SEQUENCE_NUMBER,AFTER_SEQUENCE_NUMBER]"
        elif shard_iterator_type in ["AT_TIMESTAMP"]:
            assert (
                timestamp
            ), "A timestamp must be specified if shard iterator type is one of [AT_TIMESTAMP]"
        else:
            assert not (
                starting_sequence_number or timestamp
            ), "Cannot specify a starting sequence number or timestamp if shard iteratory type is one of [TRIM_HORIZON,LATEST]"

        # Get shard IDs
        if shard_id:
            shard_ids = [shard_id]
        else:
            shard_ids = [
                s["ShardId"] for s in self.list_shards(stream_name, shard_filter=shard_filter)
            ]

        # For each shard, get a shard iterator and retrieve records using the iterator
        records = {}
        for sid in shard_ids:
            shard_iterator = self.__get_shard_iterator(
                stream_name=stream_name,
                shard_id=sid,
                shard_iterator_type=shard_iterator_type,
                starting_sequence_number=starting_sequence_number,
                timestamp=timestamp,
            )
            records[sid] = self.get_records_with_shard_iterator(shard_iterator) or []
        return records

    def get_records_with_shard_iterator(self, shard_iterator: Optional[str] = None) -> list[dict]:
        """Retrieve records from a Kinesis stream's shard beginning at shard_iterator.
        This function handles backoff and retry in the event of a
        ProvisionedThroughputExceededException.

        Parameters
        ----------
        shard_iterator : str
            The position in the shard from which to start sequentially reading records.

        Returns
        -------
        list[dict]
            list of records of the form

            [
                {
                    'SequenceNumber': 'string',
                    'ApproximateArrivalTimestamp': datetime(2015, 1, 1),
                    'Data': b'bytes',
                    'PartitionKey': 'string',
                    'EncryptionType': 'NONE'|'KMS'
                }
            ]
        """
        self.logger.info(f"Getting records with shard iterator: {shard_iterator}")
        records, has_more_records = [], True
        while has_more_records:
            self.logger.debug(f"Getting records with shard iterator: {shard_iterator}")
            resp = self.__get_records_with_shard_iterator(shard_iterator)
            records.extend(resp.get("Records") or [])
            shard_iterator = resp.get("NextShardIterator")
            has_more_records = resp.get("MillisBehindLatest") > 0 and not resp.get("ChildShards")
        return records

    def __get_shard_iterator(
        self,
        stream_name: str,
        shard_id: Optional[str] = None,
        shard_iterator_type: Optional[str] = None,
        starting_sequence_number: Optional[str] = None,
        timestamp: Optional[datetime] = None,
        wait_seconds: Optional[float] = None,
        retry_count: Optional[int] = None,
    ) -> str:
        """Retrieves a shard iterator. Handles backoff and retry in
        the event of a ProvisionedThroughputExceededException.
        """
        self.logger.debug(
            f'Retrieving shard iterator for shard "{shard_id}" with type "{shard_iterator_type}"'
        )
        try:
            assert int(retry_count or 0) <= self.max_retry_count, "Retries have been exhausted"
            if wait_seconds:
                self.logger.info(f"Waiting {wait_seconds} second(s)...")
                time.sleep(round(wait_seconds, 2))

            get_shard_iterator_input = {
                "StreamName": stream_name,
                "ShardId": shard_id,
                "ShardIteratorType": shard_iterator_type,
            }
            if starting_sequence_number:
                get_shard_iterator_input["StartingSequenceNumber"] = starting_sequence_number
            elif timestamp:
                get_shard_iterator_input["Timestamp"] = timestamp

            return self.client.get_shard_iterator(**get_shard_iterator_input).get("ShardIterator")
        except ClientError as e:
            if e.response["Error"]["Code"] == PROVISIONED_THROUGHPUT_EXCEPTION:
                self.logger.warning(
                    "Encountered a ProvisionedThroughputExceededException. Retrying."
                )
                return self.__get_shard_iterator(
                    stream_name=stream_name,
                    shard_id=shard_id,
                    shard_iterator_type=shard_iterator_type,
                    starting_sequence_number=starting_sequence_number,
                    timestamp=timestamp,
                    wait_seconds=self._next_wait_seconds(wait_seconds),
                    retry_count=self._next_retry_count(retry_count),
                )
            self.logger.error(
                f'Failed to get shard iterator for shard "{shard_id}" with exception: {e}'
            )
            raise e
        except Exception as e:
            self.logger.error(
                f'Failed to get shard iterator for shard "{shard_id}" with exception: {e}'
            )
            raise e

    def __get_records_with_shard_iterator(
        self,
        shard_iterator: Optional[str] = None,
        wait_seconds: Optional[float] = None,
        retry_count: Optional[int] = None,
    ) -> dict:
        """Retrieve records from a Kinesis stream using a shard iterator.
        This function handles backoff and retry in the event of a
        ProvisionedThroughputExceededException.
        """
        try:
            assert int(retry_count or 0) <= self.max_retry_count, "Retries have been exhausted"
            if wait_seconds:
                self.logger.info(f"Waiting {wait_seconds} second(s)...")
                time.sleep(round(wait_seconds, 2))

            return self.client.get_records(ShardIterator=shard_iterator, Limit=10000)
        except ClientError as e:
            if e.response["Error"]["Code"] == PROVISIONED_THROUGHPUT_EXCEPTION:
                self.logger.warning(
                    "Encountered a ProvisionedThroughputExceededException. Retrying."
                )
                return self.__get_records_with_shard_iterator(
                    shard_iterator=shard_iterator,
                    wait_seconds=self._next_wait_seconds(wait_seconds),
                    retry_count=self._next_retry_count(retry_count),
                )
            self.logger.error(f"Failed to get records with exception: {e}")
            raise e
        except Exception as e:
            self.logger.error(f"Failed to get records with exception: {e}")
            raise e

    def _next_wait_seconds(self, current_wait_seconds: Optional[float] = None) -> float:
        if not current_wait_seconds:
            return 1.0
        return current_wait_seconds * self.retry_backoff_rate

    def _next_retry_count(self, current_retry_count: Optional[int] = None) -> int:
        if not current_retry_count:
            return 1
        return current_retry_count + 1
