"""API to Smart Downloader."""
from collections import Counter
import dataclasses
import json
import logging
import time
from typing import Iterable, List
import uuid

import boto3
from boto3.dynamodb.conditions import Key

from feed_ingestion.util.aws.dynamodb import query_dynamodb_paginated_generator
from feed_ingestion.util.itertools_utils import batched

logger = logging.getLogger(__name__)


@dataclasses.dataclass
class DownloadTask:
    """Represents single download task."""

    source_url: str
    destination_url: str


@dataclasses.dataclass
class DownloadRequest:
    """Represents one request for smart downloader."""

    tasks: List[DownloadTask]
    tasksMessageId: str

    def to_dict(self, **kwargs):
        """Convert self to dict."""
        return dataclasses.asdict(self)

    def to_json(self, **json_kwargs):
        """Convert self to json string."""
        return json.dumps(self.to_dict(), **json_kwargs)


def query_for_job_id(job_id, dynamodb_table):
    """Lookup all items by job-id."""
    query_args = {
        'IndexName': 'jobId-index',
        'KeyConditionExpression': Key('jobId').eq(job_id),
        'ReturnConsumedCapacity': 'TOTAL',
        'ProjectionExpression': '#v1',
        'ExpressionAttributeNames': {'#v1': 'status'}
    }
    generator = query_dynamodb_paginated_generator(
        dynamodb_table=dynamodb_table,
        query_args=query_args,
    )
    return list(generator)


def send_download_requests(
        download_requests: Iterable[DownloadRequest],
        job_id: str,
        table_name: str,
        ttl_timeout_seconds: int,
        aws_session=None
):
    """Send download requests to dynamoDB.

    It uses batch_writer for better performance.
    """
    if not aws_session:
        aws_session = boto3.Session()

    dynamodb = aws_session.resource('dynamodb')
    table = dynamodb.Table(table_name)

    message_counter = 0
    ttl = int(time.time()) + ttl_timeout_seconds
    with table.batch_writer() as batch_writer:
        for download_request in download_requests:
            item = convert_download_request_to_dynamodb_item(
                download_request=download_request,
                job_id=job_id,
                task_ttl_seconds=ttl,
            )
            batch_writer.put_item(Item=item)
            message_counter += 1

    return message_counter


def await_downloads_completion(
        job_id: str,
        expected_number_of_items: int,
        table_name: str = None,
        timeout_seconds: int = 600,
        cycle_wait_seconds: int = 10,
        aws_session=None,
):
    """Wait for final states (DONE or ERROR) for all entries of job_id.

    :param job_id:
    :param expected_number_of_items:
    :param table_name:
    :param cycle_wait_seconds: sleep before next try
    :return:

    :raises TimeoutError if duration longer than timeout
    """
    if cycle_wait_seconds > timeout_seconds:
        raise ValueError(f'cycle_wait_seconds {cycle_wait_seconds} should exceed timeout_seconds {timeout_seconds}')  # noqa:E501
    if expected_number_of_items <= 0:
        raise ValueError(f'expected_number_of_items {expected_number_of_items} should be positive')  # noqa:E501
    if not aws_session:
        aws_session = boto3.Session()

    dynamodb = aws_session.resource('dynamodb')
    table = dynamodb.Table(table_name)

    started_at = time.time()
    while True:
        logger.debug('Listing of items...')
        items = query_for_job_id(
            job_id=job_id,
            dynamodb_table=table,
        )
        number_of_items = len(items)
        status_counter = Counter(item['status'] for item in items)
        states = dict(status_counter)
        logger.info(f'Current states: {states}')
        if number_of_items != expected_number_of_items:
            logger.warning(
                (f'Number of items in DynamoDB {number_of_items}'
                 f' not equals to expected number of responses'
                 f' {expected_number_of_items}'))
        logger.debug(f'Got {number_of_items} items')
        number_of_completed = states.get('DONE', 0) + states.get('ERROR', 0)

        if expected_number_of_items == number_of_completed:
            return states
        duration = time.time() - started_at
        if duration >= timeout_seconds:
            raise TimeoutError(
                (f'Waited for {duration:.0f} seconds. '
                 f'Current states: {states} of '
                 f'total {expected_number_of_items} '
                 f'expected responses'))
        logger.info(f'Sleep for {cycle_wait_seconds}')
        time.sleep(cycle_wait_seconds)


def convert_download_request_to_dynamodb_item(
        download_request: DownloadRequest,
        job_id: str,
        task_ttl_seconds: int = None):
    """
    Generate dynamodb item ready to be used by dynamodb AWS resource.

    :param download_request:
    :param job_id: download request is a part of this job
    :param task_ttl_seconds: TTL attribute for https://docs.aws.amazon.com/amazondynamodb/latest/developerguide/TTL.html  # noqa:E501
    :return: dict ready to be used as put_item for dynamodb table
    """
    return {
        'jobId': job_id,
        'taskId': download_request.tasksMessageId,
        'request': download_request.to_json(indent=2),
        'partition': f'{job_id}-{download_request.tasksMessageId}',
        'status': 'NEW',
        'ttl': task_ttl_seconds,
    }


def download_request_batched_generator(
        tasks: Iterable[DownloadTask],
        batch_size: int
) -> List[DownloadRequest]:
    """Generate download requests from download tasks."""
    for batch in batched(tasks, batch_size):
        download_request = DownloadRequest(
            tasksMessageId=str(uuid.uuid4()),
            tasks=batch,
        )
        yield download_request