from django.apps import apps
from lru import LRU
from .. import settings
from . import RecordMetadata, Record, get_offset_backend
from botocore.exceptions import ClientError
import boto3
import collections
import logging
import time

logger = logging.getLogger(__name__)


class KinesisBase(object):
    _client = None

    @property
    def client(self):
        if not self._client:
            kwargs = self._get_client_config()
            self._client = boto3.client("kinesis", **kwargs)
        return self._client

    def _get_client_config(self):
        return {
            "region_name": settings.get_aws_region(),
        }


class ModelOffsetStore(object):
    def commit(self, consumer, message):
        KinesisOffset = apps.get_model(app_label="logpipe", model_name="KinesisOffset")
        region = settings.get_aws_region()
        logger.debug(
            'Commit offset "%s" for region "%s", stream "%s", shard "%s" to %s'
            % (
                message.offset,
                region,
                message.topic,
                message.partition,
                self.__class__.__name__,
            )
        )
        obj, created = KinesisOffset.objects.get_or_create(
            region=region, stream=message.topic, shard=message.partition
        )
        obj.sequence_number = message.offset
        obj.save()

    def seek(self, consumer, stream, shard):
        KinesisOffset = apps.get_model(app_label="logpipe", model_name="KinesisOffset")
        region = settings.get_aws_region()
        try:
            obj = KinesisOffset.objects.get(
                region=settings.get_aws_region(), stream=stream, shard=shard
            )
            logger.debug(
                'Seeking to offset "%s" on region "%s", stream "%s", partition "%s"'
                % (obj.sequence_number, region, stream, shard)
            )
            consumer.seek_to_sequence_number(shard, obj.sequence_number)
        except KinesisOffset.DoesNotExist:
            logger.debug(
                'Seeking to beginning of region "%s", stream "%s", partition "%s"'
                % (region, stream, shard)
            )
            consumer.seek_to_sequence_number(shard, None)


class Consumer(KinesisBase):
    def __init__(self, topic_name, **kwargs):
        self.topic_name = topic_name
        self.client_kwargs = kwargs

        self.shards = collections.deque()
        self.records = collections.deque()
        self.shard_iters = {}

        shards = self._list_shard_ids()
        logger.debug("Found {} kinesis shards.".format(len(shards)))
        backend = get_offset_backend()
        for shard in shards:
            self.shards.append(shard)
            backend.seek(self, self.topic_name, shard)

    def seek_to_sequence_number(self, shard, sequence_number=None):
        if sequence_number is None:
            resp = self.client.get_shard_iterator(
                StreamName=self.topic_name,
                ShardId=shard,
                ShardIteratorType="TRIM_HORIZON",
            )
        else:
            resp = self.client.get_shard_iterator(
                StreamName=self.topic_name,
                ShardId=shard,
                ShardIteratorType="AFTER_SEQUENCE_NUMBER",
                StartingSequenceNumber=sequence_number,
            )
        self.shard_iters[shard] = resp["ShardIterator"]

    def __iter__(self):
        return self

    def __next__(self):
        # Try and load records. Keep trying until either (1) we have some records or (2) current_lag drops to 0
        while len(self.records) <= 0:
            # Load a page from each shard and sum the shard lags
            current_lag = 0
            for i in range(len(self.shards)):
                current_lag += self._load_next_page()

            # If all shards report 0 lag, then give up trying to load records
            if current_lag <= 0:
                break

        # If we've tried all the shards and still don't have any records, stop iteration
        if len(self.records) == 0:
            raise StopIteration()

        # Return the left most record in the queue
        return self.records.popleft()

    def _load_next_page(self):
        # Load a page from the left-most shard in the queue
        try:
            shard = self.shards.popleft()
        except IndexError:
            return 0

        # Get the next shard iterator for the shard
        shard_iter = self.shard_iters.pop(shard, None)
        if not shard_iter:
            return 0

        # Fetch the records from Kinesis
        logger.debug(
            "Loading page of records from {}.{}".format(self.topic_name, shard)
        )
        fetch_limit = settings.get("KINESIS_FETCH_LIMIT", 25)
        response = self._get_records(shard_iter, fetch_limit)
        if response is None:
            return 0

        # This default value is mostly just for testing with Moto. Real Kinesis should always return a value for MillisBehindLatest.
        num_records = len(response["Records"])
        if "MillisBehindLatest" in response:
            current_stream_lag = response["MillisBehindLatest"]
        else:
            current_stream_lag = 0 if num_records == 0 else 1
        logger.debug(
            "Loaded {} records from {}.{}. Currently {}ms behind stream head.".format(
                num_records, self.topic_name, shard, current_stream_lag
            )
        )

        # Add the records page into the queue
        timestamp = (time.time() * 1000) - current_stream_lag
        for r in response["Records"]:
            record = Record(
                topic=self.topic_name,
                partition=shard,
                offset=r["SequenceNumber"],
                timestamp=timestamp,
                key=r["PartitionKey"],
                value=r["Data"],
            )
            self.records.append(record)

        # Add the shard back to the right of the queue and save the shard iterator for next time we need
        # to get records from this shard. If NextShardIterator is None, the shard has been closed and
        # we should remove it from the pool.
        if response.get("NextShardIterator", None):
            self.shard_iters[shard] = response["NextShardIterator"]
            self.shards.append(shard)
        else:
            logger.info(
                "Shard {}.{} has been closed. Removing it from the fetch pool.".format(
                    self.topic_name, shard
                )
            )

        return current_stream_lag

    def _get_records(self, shard_iter, fetch_limit, retries=1):
        i = 0
        while i <= retries:
            try:
                response = self.client.get_records(
                    ShardIterator=shard_iter, Limit=fetch_limit
                )
                return response
            except ClientError as e:
                if (
                    e.response["Error"]["Code"]
                    == "ProvisionedThroughputExceededException"
                ):
                    logger.warning(
                        "Caught ProvisionedThroughputExceededException. Sleeping for 5 seconds."
                    )
                    time.sleep(5)
                else:
                    logger.warning(
                        "Received {} from AWS API: {}".format(
                            e.response["Error"]["Code"], e.response["Error"]["Message"]
                        )
                    )
            i += 1
        logger.warning(
            "After {} attempts, couldn't get records from Kinesis. Giving up.".format(i)
        )
        return None

    def _list_shard_ids(self):
        resp = self.client.describe_stream(StreamName=self.topic_name)
        return [shard["ShardId"] for shard in resp["StreamDescription"]["Shards"]]


class Producer(KinesisBase):
    _last_sequence_numbers = LRU(settings.get("KINESIS_SEQ_NUM_CACHE_SIZE", 1000))

    def send(self, topic_name, key, value):
        kwargs = {
            "StreamName": topic_name,
            "Data": value,
            "PartitionKey": key,
        }

        if topic_name not in self._last_sequence_numbers:
            self._last_sequence_numbers[topic_name] = {}
        last_seq_num = self._last_sequence_numbers[topic_name].get(key)
        if last_seq_num:
            kwargs["SequenceNumberForOrdering"] = last_seq_num

        metadata = self._send_and_retry(kwargs)

        shard_id = metadata["ShardId"]
        seq_num = str(metadata["SequenceNumber"])
        self._last_sequence_numbers[topic_name][key] = seq_num

        return RecordMetadata(topic=topic_name, partition=shard_id, offset=seq_num)

    def _send_and_retry(self, data, retries=1):
        i = 0
        while i <= retries:
            try:
                metadata = self.client.put_record(**data)
                return metadata
            except ClientError as e:
                if (
                    e.response["Error"]["Code"]
                    == "ProvisionedThroughputExceededException"
                ):
                    logger.warning(
                        "Caught ProvisionedThroughputExceededException. Sleeping for 5 seconds."
                    )
                    time.sleep(5)
                else:
                    logger.warning(
                        "Received {} from AWS API: {}".format(
                            e.response["Error"]["Code"], e.response["Error"]["Message"]
                        )
                    )
            i += 1
        logger.warning(
            "After {} attempts, couldn't send message to Kinesis. Giving up.".format(i)
        )
