"""
Multithreading script to manually delete all the items from _totals
DynamoDB tables of the Source of Streams.

Requirements:
    boto3, pynamodb, source-of-streams-dynamodb-models (from The Orchard pypi)
    AWS creds in the env

Approx. execution time for one day:
    vendor_totals – 2 min (for ~16-18K items)
    track_totals – 50 mins (for ~3 mln items)

PLEASE ADJUST THROUGHPUT MANUALLY BEFORE AND AFTER! Recommended values:
{env}_streams_vendor_totals: RCU – 5000, WCU - 1000.
{env}_streams_track_totals: RCU – 15000, WCU - 2000.
"""

import datetime
import queue
import random
import time
import threading

import boto3
import boto3.session
from pynamodb.exceptions import PutError, ScanError, VerboseClientError
from sosmodels.streams_track_totals import StreamsTrackTotals
from sosmodels.streams_vendor_totals import StreamsVendorTotals


vendors_totals_model = StreamsVendorTotals
track_totals_model = StreamsTrackTotals


class DeleteItemsTask(threading.Thread):
    """Thread to delete items from DynamoDB table.

    Parallel scan segment of the table by sort key (date), and deletes items 
    by batch writes (25 items in batch).
    """
    def __init__(self, date, table_name):
        super(DeleteItemsTask, self).__init__()
        self.date = date
        self.table_name = table_name
        self.deleted_count = 0
        session = boto3.session.Session()
        session.resource('dynamodb')

    def run(self):
        self.delete_items()

    def delete_items(self):
        if 'track_totals' in self.table_name:
            table_model = track_totals_model
        elif 'vendor_totals' in self.table_name:
            table_model = vendors_totals_model
        else:
            table_model = None
        if not table_model:
            print('Incorrect table name')
            return
        table_model.Meta.table_name = self.table_name

        print('Throughput of this table was already decreased {} today'.format(
            table_model.describe_table().get('ProvisionedThroughput').get(
                'NumberOfDecreasesToday')))

        with table_model.batch_write() as batch:
            segment_num = segments_q.get()
            if 'track_totals' in self.table_name:
                scan_obj = table_model.scan(
                    date__eq=self.date,
                    segment=segment_num,
                    total_segments=40)
            elif 'vendor_totals' in self.table_name:
                scan_obj = table_model.scan(
                    date__eq=self.date,
                    segment=segment_num,
                    total_segments=40)
            while True:
                try:
                    item = next(scan_obj)
                    batch.delete(item)
                    self.deleted_count += 1
                    if not self.deleted_count % 250:
                        counter_q.put(250)
                except StopIteration:
                    print(
                        'No more items in segment {}, exit from thread'.format(
                            segment_num))
                    counter_q.put('end')
                    print('{count} items were deleted by thread {n}'.format(
                        count=self.deleted_count, n=segment_num))
                    break
                except (ScanError, VerboseClientError, PutError) as e:
                    secs = random.randint(20, 30)
                    print(
                        'Throughput exceeded due to {e}, sleep thread {seg}'
                        'for '.format(e=e, seg=segment_num), secs)
                    time.sleep(secs)
                    continue


def print_count(count_queue, number_of_threads):
    count = 0
    ended = 0
    while True:
        item = count_queue.get()
        if item == 'end':
            ended += 1
        else:
            count += item
        if ended == number_of_threads:
            print('{} items were deleted overall, exit'.format(count))
            break
        if not count % 5000:
            print(ended, ' threads ended')
            print('Approximately {} items were deleted'.format(count))


if __name__ == '__main__':
    start_time = datetime.datetime.now()

    date = input(
        '\n\n===IMPORTANT: please increase throughput of _totals tables. '
        'Deleting will be performed in parallel threads.=== \n\n'
        'So if you are ready to delete items from _totals DynamoDB tables for '
        'a specific day, enter date in YYYY-MM-DD format:\n')

    table_name = input(
        '\n\nEnter table name (e.g., prod_streams_track_totals):\n')

    number_of_segments_and_threads = int(input(
        '\n\nEnter number of threads for parallel scanning the segments of '
        'the table (recommended value: 40):\n'))

    counter_q = queue.Queue()
    segments_q = queue.Queue()
    for segment in range(0, number_of_segments_and_threads):
        segments_q.put(segment)

    count_thread = threading.Thread(
        target=print_count, args=(counter_q, number_of_segments_and_threads))
    count_thread.start()

    threads = []

    for i in range(0, number_of_segments_and_threads):
        new_thread = DeleteItemsTask(date=date, table_name=table_name)
        threads.append((new_thread, i))
        new_thread.start()
        print('Started thread number ' + str(i))
        time.sleep(random.randint(2, 4))

    for thread, i in threads:
        thread.join()
        print('Joined thread number ' + str(i))

    count_thread.join()
    print('Joined counter thread')

    print(
        'Execution time' + str(
            (datetime.datetime.now() - start_time).total_seconds()))