import gzip
import logging
from datetime import timedelta, datetime
from itertools import chain
from typing import Dict, Sequence

from airflow.exceptions import AirflowSkipException
from airflow.providers.amazon.aws.hooks.s3 import S3Hook
from airflow.sensors.base import BaseSensorOperator
from airflow.utils import timezone
from requests import HTTPError, Session

from flows.spotify import api, config
from utils import aws_utils
from utils.context_util import get_context_values

MEGABYTE = 1024 ** 2

logger = logging.getLogger(__name__)

def bootstrap(report_name, licensor, dag_run, ds_nodash, ds, **kwargs):
    reports = dag_run.conf.get('reports', "")
    reports_list = get_context_values(reports, list(config.reports.keys()))

    archive_s3_path_template = f's3://{config.ARCHIVE_S3_BUCKET}/{config.ARCHIVE_S3_KEY_PATH_TEMPLATE}'
    archive_s3_path = archive_s3_path_template.format(
        licensor=licensor,
        report_name=report_name,
        date=ds,
    )
    temp_table = config.temp_staging_raw_table.format(
        licensor=licensor,
        report_name=report_name,
        date_nodash=ds_nodash,
    )
    return {
        'temp_table': temp_table,
        'archive_s3_path': archive_s3_path,
    }


class SpotifyAPISensor(BaseSensorOperator):
    template_fields: Sequence[str] = ['archive_s3_path']

    def __init__(
            self,
            *,
            report_name: str,
            licensor: str,
            archive_s3_path: str,
            country: str = None,
            aws_conn_id = 'aws_default',
            dev_mode = False,
            **kwargs) -> None:
        super().__init__(**kwargs)
        self.report_name = report_name
        self.licensor = licensor
        self.country = country
        self.archive_s3_path = archive_s3_path
        self.aws_conn_id = aws_conn_id
        self.dev_mode = dev_mode

    def poke(self, context: Dict) -> bool:
        date = context['logical_date']

        creds = config.spotify_api_credentials[self.licensor]
        client = api.SpotifyAPI(
            client_id=creds['client_id'],
            client_secret=creds['client_secret'],
            licensor_name=creds['licensor'],
            version=creds['version'],
        )

        s3_hook = S3Hook(aws_conn_id=self.aws_conn_id)
        s3_client = s3_hook.get_session().client('s3')

        filename = config.file_pattern.format(
            licensor=self.licensor,
            report_name=self.report_name,
            date=date,
            country_code=self.country
        )
        archive_s3_url = f'{self.archive_s3_path}{filename}'
        logger.info(f'Uploading to {archive_s3_url}')
        bucket, key = s3_hook.parse_s3_url(archive_s3_url)
        mpu = aws_utils.S3MultipartUpload(
            bucket=bucket,
            key=key,
            s3_client=s3_client,
        )

        resource_name = config.reports[self.report_name]['api_resource_name']
        request_to_api = client._prepare_request(resource_name, date, creds['version'], self.country)

        logger.info(f'Request to Spotify API: {request_to_api.url}')
        res = Session().send(request_to_api, stream=True)
        res.raise_for_status()

        chunk_size = MEGABYTE * 50 if not self.dev_mode else MEGABYTE * 10
        logger.info(f'Using chunk size: {aws_utils.human_readable_size((chunk_size))}')
        data_chunks = res.iter_content(chunk_size=chunk_size)

        if self.dev_mode:
            LIMIT_LINES_IN_DEV_MODE = 100000
            logger.info(f'In DEV mode currently. Only loads first: {LIMIT_LINES_IN_DEV_MODE} lines of data')

            first_chunk = next(data_chunks)
            data = aws_utils.gunzip_head(first_chunk, limit_lines=LIMIT_LINES_IN_DEV_MODE)
            data = gzip.compress(data)
            s3_hook.load_bytes(
                bytes_data=data,
                bucket_name=bucket,
                key=key,
                replace=True,
            )
        else:
            # We stick last two chunks in order ensure they all above S3 multiupload minimum
            ready_to_multiupload_chunk = aws_utils.stick_last_two(data_chunks)
            first_chunk = next(ready_to_multiupload_chunk)
            if len(first_chunk) <= aws_utils.S3MultipartUpload.PART_MINIMUM:
                logger.info(f'Small file {len(first_chunk)}. Load at once')
                s3_hook.load_bytes(
                    bytes_data=first_chunk,
                    bucket_name=bucket,
                    key=key,
                    replace=True,
                )
            else:
                result = mpu.upload(chain([first_chunk], ready_to_multiupload_chunk))
                self.log.info(result)

        return True


def split_by_prefixes(prefixes, data):
    splits = {}
    removed = set()
    for prefix in reversed(prefixes):
        splits[prefix] = []
        for item in data:
            if item in removed:
                continue
            if item[0] >= prefix:
                splits[prefix].append(item)
                removed.add(item)
    for item in reversed(data):
        if item not in removed:
            splits[prefixes[0]].insert(0, item)
    return splits