import boto3
import logging
import zipfile

from datetime import date
from io import BytesIO
from pcm_vulnerabilities_api import DEFAULT_LOGGER, log_it
from typing import Tuple


class QualysS3Client(object):
    def __init__(
        self,
        logger: logging.Logger = None,
        session: boto3.Session = None,
    ) -> None:
        self._logger = logger or DEFAULT_LOGGER
        self._session = session or boto3.Session()

    @log_it
    def get_report_index_for_date(
        self,
        bucket: str = None,
        report_prefix: str = None,
        mask: str = None
    ) -> Tuple:
        bucket = bucket or 'qualys-rollout-reports'
        report_prefix = report_prefix or 'qualys-rollout-report'
        mask = mask or f'{date.today().year}-{str(date.today().month).zfill(2)}-{str(date.today().day).zfill(2)}'

        self._logger.debug(f'searching date mask: "{mask}"')
        client = self._session.client('s3')
        response = client.list_objects_v2(
            Bucket=bucket,
            EncodingType='url',
            Prefix=f'{report_prefix}_{mask}'
        )
        if not response.get('Contents'):
            return tuple()
        return response.get('Contents')

    @log_it
    def get_basic_report_data(self, key: str, bucket: str = None) -> str:
        bucket = bucket or 'qualys-rollout-reports'
        client = self._session.client('s3')
        obj = client.get_object(Bucket=bucket, Key=key)

        with zipfile.ZipFile(BytesIO(obj['Body'].read())) as report:
            mask = key.split('.')[0].split('_')[1]
            return report.read(f"reports/basic-qualys-rollout-report_{mask}.csv").decode('utf-8')

    @log_it
    def get_detailed_report_data(self, key: str, bucket: str = 'qualys-rollout-reports') -> str:
        bucket = bucket or 'qualys-rollout-reports'
        client = self._session.client('s3')
        obj = client.get_object(Bucket=bucket, Key=key)

        with zipfile.ZipFile(BytesIO(obj['Body'].read())) as report:
            mask = key.split('.')[0].split('_')[1]
            return report.read(f'reports/detailed-qualys-rollout-report_{mask}.csv').decode('utf-8')
