import base64
import json
import logging
import os
import pandas
import requests
import sys
import time
import uuid

from datetime import datetime
from io import StringIO
from typing import Dict, List, Any

PCM_VULNERABILITIES_API_URL = os.environ.get('PCM_VULNERABILITIES_API_URL', 'https://api.pearsoncloudmanagement.com/vulnerabilities')


class PcmVulnerabilitiesApi:

    def __init__(
        self,
        user_group_name: str,
        user_group_token: str,
        logger: logging.Logger = None,
        max_wait_time: int = 300,
        wait_interval: int = 2
    ):
        self._logger = logger if logger else self._get_logger()
        self._max_wait_time = max_wait_time  # how long to wait for a response
        self._wait_interval = wait_interval  # how long to wait between each request when polling
        self._auth_token = self.api_login(user_group_name, user_group_token)

    def execute_query(self, path: str, data: Dict[str, Any] = {}):
        """Executes a series of calls to the PCM vulnerabilities API to 
        (1) start a new query and (2) poll for the query status until it completes.
        Returns a query dict of a successful query.

        Parameters
        ----------
        path : str
            Resource path to query in the PCM vulnerabilities API such that
            the final url is: https://api.pearsoncloudmanagement.com/vulnerabilities{path}

        data : dict
            Optional data dictionary to use with the POST request

        Returns
        -------
        dict
            A query dict of the form:
            {
                'query_execution_id': '92c28d6e-236d-47cd-9efa-e9ad1068ed22',
                'query_execution_status': 'SUCCEEDED',
                'query_execution_start': '2020-12-17T22:26:21Z',
                'query_execution_end': '2020-12-17T22:26:25Z',
                'query_results': 'https://some.signed.s3.url/athena/signature'
            }
        """
        url = f'{PCM_VULNERABILITIES_API_URL}{path}'
        headers = {'Authorization': f'Bearer {self._auth_token}'}

        # Start the query execution
        try:
            resp = requests.post(url, json=data, headers=headers)
            resp.raise_for_status()
            next_link = json.loads(resp.text)['next_link']
        except Exception as e:
            self._logger.error(f'Failed to start PCM vulnerabilities API query execution with exception: {e}')
            raise e

        # Poll for the query status.
        # Exit loop upon query completion or max wait time exceeded
        retry_count = 0
        query_execution_status = 'QUEUED'
        while retry_count < (self._max_wait_time / self._wait_interval):
            try:
                resp = requests.get(next_link, headers=headers)
                resp.raise_for_status()
                query = json.loads(resp.text)
                query_execution_status = query['query_execution_status']
            except Exception as e:
                self._logger.error(f'Failed to retrieve PCM vulnerabilities API query_execution_status with exception: {e}')
                raise e

            if query_execution_status in ['SUCCEEDED', 'FAILED', 'CANCELLED']:
                break
            retry_count += 1
            time.sleep(wait_time)

        self._logger.debug(f'Query for path "{path}": {query}')
        return query

    def api_login(self, user_group_name: str, user_group_token: str) -> str:
        """Takes a user_group_name and user_group_token, base64 encodes them,
        and executes a request using the encoded values to retrieve a temporary
        Bearer auth token for use with the PCM API.
        Returns the temporary Bearer auth token.

        Parameters
        ----------
        user_group_name : str
            The user_group_name, like 'finops_ro'

        user_group_token : str
            The secret token value associated with the user_group_name

        Returns
        -------
        str
            A temporary Bearer auth token for use with the PCM API.
        """
        path = 'authtoken'
        url = f'{PCM_VULNERABILITIES_API_URL}/{path}'
        data = {
            'user_group_name': base64.b64encode(user_group_name.encode('utf-8')).decode('utf-8'),
            'user_group_token': base64.b64encode(user_group_token.encode('utf-8')).decode('utf-8')
        }

        try:
            resp = requests.post(url, json=data)
            resp.raise_for_status()
            auth_token = json.loads(resp.text)['api_access_token']
        except Exception as e:
            self._logger.error(f'Failed to retrieve an auth token for the PCM API with error {e}')
            raise e

        return auth_token

    def download_results(self, s3_url: str, filename: str):
        """Downloads the file at the given s3 location to the local file system.

        Parameters
        ----------
        s3_url : str
            The S3 URL location of the file to download
        filename : str
            Filename to name the file with when written to the local file system.
        """
        with requests.get(s3_url, stream=True) as resp:
            resp.raise_for_status()
            with open(filename, 'wb') as f:
                for chunk in resp.iter_content(chunk_size=8192):
                    f.write(chunk)

    def download_results_as_dict(self, s3_url: str) -> Dict[str, Any]:
        """Downloads the file at the given s3 location to the local file system.
        Transforms the data into a dict and returns that dict. Deletes the local file.

        Parameters
        ----------
        s3_url : str
            The S3 URL location of the file to download
        dtype : Dict[str, Any]
            Optional description of the data types of the dict to return.
            Each CSV row may be parsed into a dict using the data types defined.
            Example: { 'account_id': str, 'cvss_v2_base': float }
        """
        filename = f'/tmp/{uuid.uuid4().hex}.csv'
        try:
            self.download_results(s3_url, filename)
        except Exception as e:
            self._logger.error(f'Failed to download results at {s3_url} with error: {e}')
            raise e

        with open(filename, 'r') as f:
            df = pandas.read_csv(StringIO(f.read()), dtype=str)
        os.remove(filename)
        return df.to_dict('records')

    def _get_logger(self, log_level: str = 'info') -> logging.Logger:
        """Return a logger with the given logging level set"""
        handler = logging.StreamHandler(sys.stdout)
        log = logging.getLogger('prometheus_exporter')
        log_level = log_level.lower()
        if log_level == 'debug':
            log.setLevel(logging.DEBUG)
            handler.setLevel(logging.DEBUG)
        elif log_level == 'warning':
            log.setLevel(logging.WARNING)
            handler.setLevel(logging.WARNING)
        elif log_level == 'error':
            log.setLevel(logging.ERROR)
            handler.setLevel(logging.ERROR)
        elif log_level == 'critical':
            log.setLevel(logging.CRITICAL)
            handler.setLevel(logging.CRITICAL)
        else:
            log.setLevel(logging.INFO)
            handler.setLevel(logging.INFO)

        formatter = logging.Formatter(
            '%(asctime)s [%(levelname)s] %(filename)s:%(lineno)d - %(message)s')
        handler.setFormatter(formatter)
        log.addHandler(handler)
        log.propagate = False
        return log


def main():
    api = PcmVulnerabilitiesApi('', '')
    # resp = api.execute_query('/aws/qualys-agent-coverage')
    resp = api.execute_query('/aws/vpd')
    if resp['query_execution_status'] != 'SUCCEEDED':
        raise Exception(f'Query failed with response: {resp}')

    vpd_data = api.download_results_as_dict(s3_url=resp['query_results'])

    print(vpd_data)


if __name__ == '__main__':
    main()
