import logging
import zipfile
from datetime import date
from io import BytesIO
from typing import Optional

import boto3

from pcm_vulnerabilities_api import get_logger, log_it


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

    @log_it
    def get_report_index_for_date(
        self,
        bucket: Optional[str] = None,
        report_prefix: Optional[str] = None,
        mask: Optional[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: Optional[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")
