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

from typeguard import typechecked

from psi_kinesis.kinesis_reader import KinesisReader, PROVISIONED_THROUGHPUT_EXCEPTION

MAX_FAILED_RECORDS = 25  # max number of failed put records to raise
MAX_INVALID_RECORDS = 10  # max number of invalid put records to raise
MAX_PARTITION_KEY_LEN = 256  # max number of characters in a partition key
MAX_PUT_RECORDS_BATCH = 500  # max number of records in put records operation


@typechecked
class KinesisWriter(KinesisReader):
    """KinesisWriter provides read/write 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 KinesisWriter

        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
            put_records or get_records. Default value of 5.

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

        logger : logging.Logger
            (optional) The logger used by a KinesisWriter object. If unspecified,
            a new logger will be created.
        """
        super().__init__(aws_session, max_retry_count, retry_backoff_rate, logger)
        # placeholder for __init__ differences between a KinesisReader and a KinesisWriter

    def put_records(
        self, stream_name: str, records: Optional[list] = None, quit_on_error: bool = False
    ) -> list[dict]:
        """Puts records into a Kinesis stream.

        This function handles the following:
          1. Validating each record's format
          2. Backoff and retry in the event of a ProvisionedThroughputExceededException
          3. Batching (the PutRecords API has a maximum of 500 records per operation)
          4. Partitioning (see "Partitioning" below)
          5. Gracefully handles failed records and exceptions

        ### Partitioning

        Kinesis streams divide data among shards using either a partition key or an
        explicit hash key to determine which shard a record goes to.
        See: https://docs.aws.amazon.com/kinesis/latest/APIReference/API_PutRecords.html

        Either a "PartitionKey" or an "ExplicitHashKey" may be placed on each record
        dictionary to control which shard the record goes to.
        If neither the "PartitionKey" nor "ExplicitHashKey" are set on the record, this
        function will attempt to automatically distribute the records uniformly among shards.

        Parameters
        ----------
        stream_name : str
            The name of the Kinesis stream where the records will be sent

        records : list
            list of records, where each record is a dictionary like:

            {
                'PartitionKey': 'string',
                'ExplicitHashKey': 'string',
                'Data': b'bytes'
            }

            The "Data" field is required and must be bytes.
            The "PartitionKey" and "ExplicitHashKey" are optional. If one is set,
            do not include the other.

        quit_on_error : bool
            (optional) Whether to quit putting records if an error occurs or to
            attempt writing all records (even if an error occurs). Default False.

        Returns
        -------
        list[dict]
            A list of locations of the new records with the form:

            [
                {
                    'ShardId': 'string',
                    'SequenceNumber': 'string'
                }
            ]
        """
        # Validate records
        invalid_records = self.__validate_put_records(records)
        assert not invalid_records, f"Some records are invalid: {invalid_records}"
        self.logger.info(f'Putting {len(records)} record(s) into Kinesis stream "{stream_name}"')

        # Batch records
        batches = self.__batch_put_records(records)

        # Write each batch
        results, exceptions = [], []
        for batch in batches:
            try:
                batch = self.__partition_put_records(batch)
                resp = self.__put_records_batch(batch, stream_name)
                results.extend(resp.get("Records") or [])
                if quit_on_error and resp.get("FailedRecordCount"):
                    break
            except Exception as e:
                self.logger.error(f"Failed to put Kinesis record with exception: {e}")
                if quit_on_error:
                    raise e
                exceptions.append(e)

        # Raise any exceptions that occurred
        if exceptions:
            self.logger.error("At least one exception occurred while putting records")
            raise Exception(exceptions)

        # Assert no records failed to write
        failed_records = [r for r in results if r.get("ErrorMessage")]
        assert (
            not failed_records
        ), f"{len(failed_records)} record(s) failed to write: {failed_records[:MAX_FAILED_RECORDS]}"

        # Return the successfully written record locations
        return [
            {"SequenceNumber": r.get("SequenceNumber"), "ShardId": r.get("ShardId")}
            for r in results
        ]

    def __put_records_batch(
        self,
        batch: list[dict],
        stream_name: str,
        wait_seconds: Optional[float] = None,
        retry_count: Optional[int] = None,
    ) -> dict[str, Any]:
        """Performs the actual boto3 put_records. Recursively handles backoff
        and retry in the event of a ProvisionedThroughputExceededException(s)."""
        self.logger.debug(f"Putting batch with {len(batch)} record(s)")
        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))

            resp = self.client.put_records(Records=batch, StreamName=stream_name)

            has_throughput_exception, has_other_exception = (
                self.__put_records_resp_has_throughput_exception(resp=resp)
            )
            if has_other_exception:
                self.logger.error("Encountered an unexpected Exception. Skipping retry.")
                return resp
            if has_throughput_exception:
                self.logger.warning(
                    "Encountered a ProvisionedThroughputExceededException. Retrying."
                )
                return self.__put_records_batch(
                    batch=batch,
                    stream_name=stream_name,
                    wait_seconds=self._next_wait_seconds(wait_seconds),
                    retry_count=self._next_retry_count(retry_count),
                )
            return resp
        except Exception as e:
            self.logger.error(f"Failed to put Kinesis records batch with exception: {e}")
            raise e

    def __partition_put_records(self, records: list[dict]) -> list[dict]:
        """Assigns a new partition key to records without an existing
        partition key or explicit hash key"""
        for r in records:
            if not ("PartitionKey" in r or "ExplicitHashKey" in r):
                sha256 = hashlib.sha3_256()
                sha256.update(r["Data"])
                r["PartitionKey"] = sha256.hexdigest()
        return records

    def __batch_put_records(
        self, records: list[dict], batch_size: int = MAX_PUT_RECORDS_BATCH
    ) -> Any:
        """Chunks the list of records into chunks of size batch_size"""
        for i in range(0, len(records), batch_size):
            yield records[i : i + batch_size]

    def __validate_put_records(self, records: list) -> list[dict]:
        """Validates records to write to a Kinesis stream.
        Returns a list of dictionaries of invalid records with the form:

        [
            {
                'record': Any,     # The invalid record
                'reason': 'string' # Why the record is invalid
            }
        ]

        The number of invalid records returned is limited by MAX_INVALID_RECORDS.
        """
        invalid_records = []
        for r in records:
            is_valid, invalid_reason = self.__validate_put_record(r)
            if not is_valid:
                invalid_records.append({"record": r, "invalid_reason": invalid_reason})
            if len(invalid_records) >= MAX_INVALID_RECORDS:
                break
        return invalid_records

    def __validate_put_record(self, record) -> tuple[bool, str]:
        """Validates a record to write to a Kinesis stream.
        Returns a tuple[bool, str] with two elements:
          1. bool : Whether the record is valid
          2. str : Reason for why the record is invalid. Empty if the record is valid
        """
        try:
            assert isinstance(record, dict), "Expected record to be a dictionary"
            assert isinstance(record.get("Data"), bytes), "Expected record.Data to be bytes"

            has_partiton_key = "PartitionKey" in record
            has_explicit_hash_key = "ExplicitHashKey" in record
            assert not (
                has_partiton_key and has_explicit_hash_key
            ), "Expected either record.PartitionKey or record.ExplicitHashKey, not both"
            if has_partiton_key:
                assert isinstance(
                    record["PartitionKey"], str
                ), "Expected record.PartitionKey to be a string"
                assert (
                    len(record["PartitionKey"]) <= MAX_PARTITION_KEY_LEN
                ), f"Expected record.PartitionKey to have fewer than {MAX_PARTITION_KEY_LEN + 1} characters"
            if has_explicit_hash_key:
                assert isinstance(
                    record["ExplicitHashKey"], str
                ), "Expected record.ExplicitHashKey to be a string"
            return True, ""
        except AssertionError as e:
            return False, str(e)

    def __put_records_resp_has_throughput_exception(self, resp: Any) -> tuple[bool, bool]:
        """Checks for ProvisionedThroughputExceededException(s) in a put_records response"""
        has_throughput_exception, has_other_exception = False, False
        if isinstance(resp, dict) and resp.get("FailedRecordCount"):
            for r in resp.get("Records") or []:
                if r.get("ErrorCode") == PROVISIONED_THROUGHPUT_EXCEPTION:
                    has_throughput_exception = True
                elif r.get("ErrorCode"):
                    has_other_exception = True
        return has_throughput_exception, has_other_exception
