import datetime
import gzip
import logging
import os
from functools import partial
from pathlib import Path
from tempfile import TemporaryDirectory
from typing import Optional

from airflow.providers.amazon.aws.hooks.s3 import S3Hook

from flows.youtube_monthly import config

THIS_DIR = Path(__file__).parent
logger = logging.getLogger(__name__)

def bootstrap(report_name, report_config, mcn_account, data_interval_start, data_interval_end, **kwargs):

    first_day, last_day = data_interval_start.date(), data_interval_end.date()
    date_obj: datetime.date = first_day

    logger.info(f'Bootstrap report {report_name} for {date_obj}')

    file_template = report_config['file_pattern']
    filename = file_template.format(
        start_date=first_day,
        end_date=last_day,
        account=mcn_account,
    )

    drop_s3_key = f'{config.DROP_S3_KEY_PATH}{filename}'
    drop_s3_url = f's3://{config.DROP_S3_BUCKET}/{drop_s3_key}'
    archive_s3_path = config.ARCHIVE_S3_KEY_PATH_TEMPLATE.format(
        report_name=report_name,
        date=date_obj,
    )
    archive_s3_path_url = 's3://{bucket}/{key}'.format(
        bucket=config.ARCHIVE_S3_BUCKET,
        key=archive_s3_path,
    )
    report_date = first_day.strftime('%Y-%m-%d')
    return dict(
        date=report_date,
        drop_filename=filename,
        drop_s3_key=drop_s3_key,
        drop_s3_url=drop_s3_url,
        archive_s3_path=archive_s3_path,
        archive_s3_path_url=archive_s3_path_url,
    )

def extract_reports_from_multireport_file(
        source_filename,
        local_temp_path,
        report_filename_map
):
    """Extract multiple reports stored in single source "CSV" file.

    The reports in youtube's pseudo-CSV file are separated by empty line.
    Next line is report name like 'Music Streams'
    Then it has normal CSV content with header line.

    Args:
        activity (ActivityWorker): The activity worker.
        source_filename (str): report file (csv or gz).
        local_temp_path (str): path to temp directory for reports,
        report_filename_map (dict): key: report name, value: cvs filename,
    """
    if source_filename.endswith('.gz'):
        open_func = partial(gzip.open, source_filename, mode='rt')
    else:
        open_func = partial(open, source_filename)

    processed_reports = []
    processed_files = []
    with open_func() as source_file:
        try:
            while True:
                line = next(source_file)
                report_name = line.strip()
                if report_name not in report_filename_map:
                    raise ValueError(f'Unexpected report {report_name}')
                result_filename = report_filename_map[report_name]
                processed_reports.append(report_name)
                processed_files.append(result_filename)
                if not result_filename.endswith('.csv') \
                        and not result_filename.endswith('.csv.gz'):
                    raise ValueError(
                        f'Non-CSV/GZ report file name {result_filename} '
                        f'for report {report_name}')

                target_file_path = f'{local_temp_path}/{result_filename}'
                if result_filename.endswith('.gz'):
                    open_write_func = partial(
                        gzip.open, target_file_path, mode='wt')
                else:
                    open_write_func = partial(
                        open, target_file_path, mode='w')
                with open_write_func() as result_file:
                    while True:
                        line = next(source_file)
                        if line == '\n':
                            # reached empty line, going to next report
                            break
                        result_file.write(line)
        except StopIteration:
            # all good. The source file has ended
            pass
    missed_reports = set(report_filename_map) - set(processed_reports)
    if missed_reports:
        raise ValueError(
            f'Missed report(s) {missed_reports} in file {source_filename}')
    return {'files': processed_files}


def move_and_extract_files(
        report_name: str,
        filename: str,
        drop_s3_url: str,
        archive_s3_path_url: str,
        subreports: Optional[dict],
):
    s3_hook_drop = S3Hook(aws_conn_id=config.AWS_DROP_BUCKET_CONN_ID)
    s3_hook_archive = S3Hook(aws_conn_id=config.AWS_ARCHIVE_BUCKET_CONN_ID)

    drop_s3_path_full = drop_s3_url
    archive_s3_path_full = f'{archive_s3_path_url}{filename}'

    drop_s3_bucket, drop_s3_key = s3_hook_drop.parse_s3_url(drop_s3_url)
    archive_s3_bucket, archive_s3_path = s3_hook_archive.parse_s3_url(archive_s3_path_url)

    with TemporaryDirectory() as local_temp_path:
        # download from drop
        download_file_kwargs = dict(
            bucket_name=drop_s3_bucket,
            key=drop_s3_key,
        )
        logger.info(f'Download {download_file_kwargs}')
        downloaded_file = s3_hook_drop.download_file(
            local_path=local_temp_path,
            **download_file_kwargs
        )
        logger.info(f'Downloaded {downloaded_file}')
        path = Path(downloaded_file).resolve()
        logger.info(f'Resolved {path}')
        # rename from temp filename to original filename

        new_location = os.path.join(local_temp_path, filename)
        os.rename(downloaded_file, new_location)
        downloaded_file = new_location

        path = Path(downloaded_file).resolve()
        logger.info(f'Resolved after rename: {path}')

        if filename.endswith('.zip'):
            filename = filename.replace('.zip', '.gz')
            archive_temp_s3_path = f'{archive_s3_path_url}{filename}'
            # update file_info in order to return updated file_name to flow
            logger.info(
                f'Convert ZIP {drop_s3_path_full} '
                f'to {archive_temp_s3_path}')
            s3.convert_zip_to_gzip_on_s3(
                zip_s3_path=drop_s3_path_full,
                gz_s3_path=archive_temp_s3_path,
                local_temp_dir=local_temp_path)
        else:
            logger.info(
                f'Copy {drop_s3_path_full} to {archive_s3_path_full}')
            s3_hook_archive.load_file(
                filename=downloaded_file,
                bucket_name=archive_s3_bucket,
                key=f'{archive_s3_path}{filename}',
                replace=True,
            )
        if subreports:
            report_filename_map = {}
            for subreport_name, info in subreports.items():
                subreport_filename = config.subreport_csv_filename(report_name, subreport_name)
                report_filename_map[info['report_title']] = subreport_filename
            logger.info(
                f'Going to extracted subreports: {report_filename_map}')
            extracted_files = extract_reports_from_multireport_file(
                source_filename=downloaded_file,
                local_temp_path=local_temp_path,
                report_filename_map=report_filename_map
            )
            logger.info(f'Extracted subreports {extracted_files}')
            for subreport_file_name in extracted_files['files']:
                load_file_kwargs = dict(
                    filename=f'{local_temp_path}/{subreport_file_name}',
                    bucket_name=archive_s3_bucket,
                    key=f'{archive_s3_path}{subreport_file_name}',
                )
                logger.info(
                    f'Upload subreport file to {load_file_kwargs}')
                s3_hook_archive.load_file(
                    **load_file_kwargs,
                    replace=True,
                )