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

from datetime import datetime, timedelta
from io import StringIO
from pcm_vulnerabilities_api import DEFAULT_LOGGER, log_it
from typing import Dict, 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,
        token_refresh: int = 30,
    ):
        self._logger = logger or DEFAULT_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._token_refresh = token_refresh  # how often to refresh the auth token (in minutues)
        self.api_login(user_group_name, user_group_token)

    @log_it
    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'
            }
        """
        self._logger.info(f'Executing query against path: {path}')
        self._refresh_token()
        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

    @log_it
    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.
        """
        self._logger.info('Retrieving new auth token')
        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

        self._user_group_name = user_group_name
        self._user_group_token = user_group_token
        self._auth_token = auth_token
        self._auth_token_dt = datetime.now().astimezone()

    @log_it
    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.
        """
        self._logger.info('Downloading results')
        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)

    @log_it
    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')

    @log_it
    def _refresh_token(self):
        if datetime.now().astimezone() > self._auth_token_dt + timedelta(minutes=self._token_refresh):
            self.api_login(self._user_group_name, self._user_group_token)
