"""Script to listen kinesis stream and invoke lambda."""
import base64
from concurrent.futures import ThreadPoolExecutor
import logging
import os
from queue import Queue
import time

import boto3
import requests
from requests.exceptions import RequestException


# Environment variables
KINESIS_STREAM_NAME = os.getenv('KINESIS_STREAM_NAME')
LAMBDA_ENDPOINT = os.getenv('LAMBDA_ENDPOINT')
AWS_REGION = os.getenv('AWS_REGION', 'us-east-1')

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s'
)

EVENT_SOURCE_ARN = f'arn:aws:kinesis:us-east-1:123456789012:stream/{KINESIS_STREAM_NAME}'
EVENT_ID = 'shardId-000000000000:49647831234567890123456789012345678901234567890'

# Create Kinesis client
kinesis_client = boto3.client('kinesis', region_name=AWS_REGION)

# Shared queue for events
event_queue = Queue()


def invoke_lambda_sequential():
    """Sequentially invoke Lambda with events from the queue."""
    while True:
        # Wait for events to be added to the queue
        event_batch = event_queue.get()
        if event_batch is None:  # Sentinel value to stop the thread
            break

        kinesis_event = {'Records': event_batch}
        logging.info(f'Invoking Lambda with payload: {kinesis_event}')

        retries = 3  # Retry 3 times before failing
        for attempt in range(retries):
            try:
                response = requests.post(LAMBDA_ENDPOINT, json=kinesis_event)
                response.raise_for_status()
                logging.info(
                    f'Lambda response: {response.status_code}, {response.text}'
                )
                break
            except RequestException as e:
                logging.error(
                    f'Error invoking Lambda: {e}, Attempt {attempt + 1}/{retries}'
                )
                if attempt < retries - 1:
                    time.sleep(2 ** attempt)  # Exponential backoff
                else:
                    logging.error(
                        'Max retry attempts reached. Skipping Lambda invocation.'
                    )

        event_queue.task_done()


def process_shard(shard_id):
    """Listen to a specific shard."""
    kinesis_client = boto3.client('kinesis', region_name=AWS_REGION)

    # Get the shard iterator
    shard_iterator_response = kinesis_client.get_shard_iterator(
        StreamName=KINESIS_STREAM_NAME,
        ShardId=shard_id,
        ShardIteratorType='LATEST'  # Start reading from the latest data
    )
    shard_iterator = shard_iterator_response['ShardIterator']

    while True:
        # Get records from the shard
        records_response = kinesis_client.get_records(
            ShardIterator=shard_iterator,
            Limit=10
        )
        records = records_response['Records']

        # Check for records
        if records:
            event_batch = []
            for record in records:
                # Decode the event from bytes to string
                event_data = record['Data']
                logging.info(f'Received event data from shard {shard_id}: {event_data}')

                # Construct Kinesis Event
                kinesis_event_record = {
                    'eventID': EVENT_ID,
                    'eventName': 'aws:kinesis:record',
                    'eventSourceARN': EVENT_SOURCE_ARN,
                    'eventSource': 'aws:kinesis',
                    'awsRegion': AWS_REGION,
                    'kinesis': {
                        'kinesisSchemaVersion': '1.0',
                        'partitionKey': record.get('PartitionKey', 'partitionKey-1'),
                        'sequenceNumber': record['SequenceNumber'],
                        'data': base64.b64encode(record['Data']).decode('utf-8'),
                        'approximateArrivalTimestamp': time.time()
                    }
                }
                event_batch.append(kinesis_event_record)

            # Add the batch of events to the queue
            event_queue.put(event_batch)

        # Update the shard iterator for the next request
        shard_iterator = records_response['NextShardIterator']

        # Wait for a short period before polling again
        if not records:  # If no records, sleep longer
            time.sleep(5)
        else:
            time.sleep(1)


def main():
    """Get all shards in the stream and listen to each shard."""
    # Get all shards in the stream
    stream_description = kinesis_client.describe_stream(
        StreamName=KINESIS_STREAM_NAME
    )
    shards = stream_description['StreamDescription']['Shards']
    shard_ids = [shard['ShardId'] for shard in shards]
    logging.info(f'Found {len(shard_ids)} shard(s). Listening to all shards...')

    # Start the Lambda invocation thread
    lambda_thread = ThreadPoolExecutor(max_workers=1).submit(invoke_lambda_sequential)

    # Use a ThreadPoolExecutor to process multiple shards concurrently
    with ThreadPoolExecutor(max_workers=len(shard_ids)) as executor:
        futures = [
            executor.submit(process_shard, shard_id) for shard_id in shard_ids
        ]

        # Wait for all threads to complete
        for future in futures:
            future.result()

    # Stop the Lambda invocation thread
    event_queue.put(None)  # Sentinel value to stop the thread
    lambda_thread.result()


if __name__ == '__main__':
    """Entrypoint of the script."""
    main()