import boto3
import logging
import pandas

from datetime import datetime
from io import StringIO
from pcm_vulnerabilities_api import DEFAULT_LOGGER, log_it
from typing import Dict, Tuple


class QualysAthenaClient:
    def __init__(
        self,
        logger: logging.Logger = None,
        session: boto3.Session = None
    ):
        self._logger = logger or DEFAULT_LOGGER
        self._session = session or boto3.Session()
        self._default_athena_database = 'vulnerability_management'
        self._default_athena_catalog = 'AwsDataCatalog'
        self._default_athena_workgroup = 'vulnerability-management'

    @log_it
    def start_query_execution(
        self,
        query: str,
        database: str = None,
        catalog: str = None,
        workgroup: str = None
    ) -> str:
        '''Start an Athena query execution

        Parameters
        ----------
        query : str
            SQL query to execute on athena

        database : str
            SQL query against this database

        catalog : str
            SQL query against this catalog

        workgroup : str
            SQL query against this workgroup

        Returns
        -------
        QueryExecutionId : str
            Athena query execution id
        '''
        assert query, 'Cannot start_query_execution. No query provided'
        athena_client = self._session.client('athena')
        response = athena_client.start_query_execution(
            QueryString=query,
            QueryExecutionContext={
                'Database': database or self._default_athena_database,
                'Catalog': catalog or self._default_athena_catalog
            },
            WorkGroup=workgroup or self._default_athena_workgroup
        )
        return response['QueryExecutionId']

    @log_it
    def start_current_instances_query_execution(
        self,
        parameters: Dict = {},
        limit_results: int = None,
        database: str = None,
        catalog: str = None,
        workgroup: str = None
    ) -> str:
        '''Start an Athena query execution for a current instances report.

        Parameters
        ----------

        parameters : Dict
            Map of query parameter filters.
            E.g. {'id':'svc02222','environment':'prd','account': '012597783113'}

        limit_results : int
            Number of records to return

        database : str
            SQL query against this database.
            Defaults to 'vulnerability_management'. Note the underscore.

        catalog : str
            SQL query against this catalog
            Defaults to 'AwsDataCatalog'. Note the underscore.

        workgroup : str
            SQL query against this workgroup
            Defaults to 'vulnerability-management'. Note the hyphen.

        Returns
        -------
        QueryExecutionId : str
            Athena query execution id
        '''
        query = self.get_query_string(
            query='SELECT * FROM vulnerability_management.swa_current_instances_vw',
            parameters=parameters,
            limit_results=limit_results
        )
        return self.start_query_execution(query, database, catalog, workgroup)

    @log_it
    def start_current_vulnerabilities_query_execution(
        self,
        parameters: Dict = {},
        limit_results: int = None,
        database: str = None,
        catalog: str = None,
        workgroup: str = None
    ) -> str:
        '''Start an Athena query execution for a current vulnerabilities report.'''
        query = self.get_query_string(
            query='SELECT * FROM vulnerability_management.swa_current_vulnerabilities_vw',
            parameters=parameters,
            limit_results=limit_results
        )
        return self.start_query_execution(query, database, catalog, workgroup)

    @log_it
    def start_historical_instances_query_execution(
        self,
        parameters: dict = None,
        filter_date: str = None,
        limit_results: int = None,
        database: str = None,
        catalog: str = None,
        workgroup: str = None
    ) -> str:
        '''Start an Athena query execution for a historical instances report.

        Parameters
        ----------

        parameters : Dict
            Map of query parameter filters.
            E.g. {'id':'svc02222','environment':'prd','account': '012597783113'}

        filter_date : str
            Format YYYY-MM-DD as generated by datetime.date.today().
            Defaults to datetime.date.today().

       limit_results : int
            Number of records to return

        database : str
            SQL query against this database.
            Defaults to 'vulnerability_management'. Note the underscore.

        catalog : str
            SQL query against this catalog
            Defaults to 'AwsDataCatalog'. Note the underscore.

        workgroup : str
            SQL query against this workgroup
            Defaults to 'vulnerability-management'. Note the hyphen.

        Returns
        -------
        QueryExecutionId : str
            Athena query execution id
        '''
        assert filter_date, f'A filter date is required to search historical records'
        query = self.get_query_string(
            query='SELECT * FROM vulnerability_management.swa_historical_instances_vw',
            parameters=parameters,
            filter_date=filter_date,
            limit_results=limit_results
        )
        return self.start_query_execution(query, database, catalog, workgroup)

    @log_it
    def start_historical_vulnerabilities_query_execution(
        self,
        parameters: dict = None,
        filter_date: str = None,
        limit_results: int = None,
        database: str = None,
        catalog: str = None,
        workgroup: str = None
    ) -> str:
        '''Start an Athena query execution for a historical vulnerabilities report.

        Parameters
        ----------

        parameters : Dict
            Map of query parameter filters.
            E.g. {'id':'svc02222','environment':'prd','account': '012597783113'}

        filter_date : str
            Format YYYY-MM-DD as generated by datetime.date.today().
            Defaults to datetime.date.today().

       limit_results : int
            Number of records to return

        database : str
            SQL query against this database.
            Defaults to 'vulnerability_management'. Note the underscore.

        catalog : str
            SQL query against this catalog
            Defaults to 'AwsDataCatalog'. Note the underscore.

        workgroup : str
            SQL query against this workgroup
            Defaults to 'vulnerability-management'. Note the hyphen.

        Returns
        -------
        QueryExecutionId : str
            Athena query execution id
        '''
        assert filter_date, f'A filter date is required to search historical records'
        query = self.get_query_string(
            query='SELECT * FROM vulnerability_management.swa_historical_vulnerabilities_vw',
            parameters=parameters,
            filter_date=filter_date,
            limit_results=limit_results
        )
        return self.start_query_execution(query, database, catalog, workgroup)

    @log_it
    def start_vpd_trends_query_execution(
        self,
        parameters: Dict = {},
        limit_results: int = None,
        database: str = None,
        catalog: str = None,
        workgroup: str = None
    ) -> str:
        '''Start an Athena query execution for a current VPD trends report.

        Parameters
        ----------
        parameters : Dict
            Map of query parameter filters.
            E.g. {'ids':['svc02222'],'environments':['prd'],'accounts': ['012597783113']}

        limit_results : int
            Number of records to return

        database : str
            SQL query against this database.
            Defaults to 'vulnerability_management'. Note the underscore.

        catalog : str
            SQL query against this catalog
            Defaults to 'AwsDataCatalog'. Note the underscore.

        workgroup : str
            SQL query against this workgroup
            Defaults to 'vulnerability-management'. Note the hyphen.

        Returns
        -------
        QueryExecutionId : str
            Athena query execution id
        '''
        query = self.get_query_string(
            query='SELECT * FROM swa_vpd_30day_change_vw',
            parameters=parameters,
            order_by='t_appid',
            limit_results=limit_results
        )
        return self.start_query_execution(query, database, catalog, workgroup)

    @log_it
    def start_historical_qualys_agent_coverage_query_execution(
        self,
        parameters: Dict = {},
        filter_date: datetime = None,
        limit_results: int = None,
        database: str = None,
        catalog: str = None,
        workgroup: str = None
    ):
        '''Start an Athena query execution for a historical qualys agent coverage report.

        Parameters
        ----------

        parameters : Dict
            Map of query parameter filters.
            E.g. {'id':'svc02222','environment':'prd','account': '012597783113'}

        filter_date : str
            Format YYYY-MM-DD as generated by datetime.date.today().
            Defaults to datetime.date.today().

       limit_results : int
            Number of records to return

        database : str
            SQL query against this database.
            Defaults to 'vulnerability_management'. Note the underscore.

        catalog : str
            SQL query against this catalog
            Defaults to 'AwsDataCatalog'. Note the underscore.

        workgroup : str
            SQL query against this workgroup
            Defaults to 'vulnerability-management'. Note the hyphen.

        Returns
        -------
        QueryExecutionId : str
            Athena query execution id
        '''
        assert filter_date, f'A filter date is required to search historical records'
        query = self.get_query_string(
            query='SELECT * FROM vulnerability_management.historical_agent_coverage_vw',
            parameters=parameters,
            filter_date=filter_date,
            limit_results=limit_results
        )
        return self.start_query_execution(query, database, catalog, workgroup)

    @log_it
    def start_historical_vpd_query_execution(
        self,
        parameters: Dict = {},
        filter_date: datetime = None,
        limit_results: int = None,
        database: str = None,
        catalog: str = None,
        workgroup: str = None
    ):
        '''Start an Athena query execution for a historical VPD report.

        Parameters
        ----------

        parameters : Dict
            Map of query parameter filters.
            E.g. {'id':'svc02222','environment':'prd','account': '012597783113'}

        filter_date : str
            Format YYYY-MM-DD as generated by datetime.date.today().
            Defaults to datetime.date.today().

       limit_results : int
            Number of records to return

        database : str
            SQL query against this database.
            Defaults to 'vulnerability_management'. Note the underscore.

        catalog : str
            SQL query against this catalog
            Defaults to 'AwsDataCatalog'. Note the underscore.

        workgroup : str
            SQL query against this workgroup
            Defaults to 'vulnerability-management'. Note the hyphen.

        Returns
        -------
        QueryExecutionId : str
            Athena query execution id
        '''
        assert filter_date, f'A filter date is required to search historical records'
        query = self.get_query_string(
            query='SELECT * FROM vulnerability_management.swa_vpd_vw',
            parameters=parameters,
            filter_date=filter_date,
            limit_results=limit_results
        )
        return self.start_query_execution(query, database, catalog, workgroup)

    @log_it
    def get_query_execution_status(self, execution_id: str):
        ''' Fetch query execution status

        Parameters
        ----------
        execution_id : str
            Athena query execution id

        Returns
        -------
        State : str
            Athena query execution status state
        '''
        assert execution_id, 'Cannot retrieve execution status without a valid execution ID'
        athena_client = self._session.client('athena')
        response = athena_client.get_query_execution(
            QueryExecutionId=execution_id
        )
        return response['QueryExecution']['Status']['State']

    @log_it
    def get_query_results(self, execution_id: str) -> str:
        '''Retrieve athena query results from S3

        Parameters
        ----------
        execution_id : str
            Athena query execution id

        Returns
        -------
        results : str
            Athena query execution results as a CSV string
        '''
        assert execution_id, 'Cannot retrieve execution results without a valid execution ID'
        athena_client = self._session.client('athena')
        resp = athena_client.get_query_execution(
            QueryExecutionId=execution_id
        )
        query_execution = resp.get('QueryExecution', {})
        if not query_execution:
            return None
        s3_path = query_execution['ResultConfiguration']['OutputLocation']

        client = self._session.client('s3')
        self._logger.debug(f'Fetching object from S3: {s3_path}')
        obj = client.get_object(
            Bucket=s3_path[len('s3://'):].split('/', 1)[0],
            Key=s3_path[len('s3://'):].split('/', 1)[1]
        )
        return obj['Body'].read()

    @log_it
    def get_query_results_as_dict(self, execution_id: str) -> str:
        '''Retrieve athena query results from S3

        Parameters
        ----------
        execution_id : str
            Athena query execution id

        Returns
        -------
        results : str
            Athena query execution results as a dict
        '''
        results = self.get_query_results(execution_id)
        df = pandas.read_csv(StringIO(results.decode('utf-8')), dtype=str)
        return df.to_dict('records')

    @log_it
    def get_query_string(
        self,
        query: str,
        parameters: dict = None,
        filter_date: str = None,
        filter_date_field: str = None,
        order_by: str = None,
        order_desc: bool = False,
        limit_results: int = None,
    ) -> str:
        '''Build a SQL query string

        Parameters
        ----------
        query : str
            BASE SQL 'FROM' query claus to execute on athena. Will be combined
            with parameters and filter_date.

        parameters : Dict
            Map of query parameter filters.
            E.g. {'id':'svc02222','environment':'prd','account': '012597783113'}

        filter_date : str
            Format YYYY-MM-DD as generated by datetime.date.today()

        filter_date_field : str
            The data field containing the dates to filter by.
            Defaults to "report_date"

        limit_results : int
            Number of records to return

        order_by: str
            Which attribute to order the results by

        order_desc: bool
            Whether or not to sort in descending order.
            Only applies if "order_by" is specified.
            Defaults to ascending order.

        Returns
        -------
        query : str
            SQL query string
        '''
        if filter_date:
            filter_date_field = filter_date_field or 'report_date'
            query += f" WHERE {filter_date_field} = DATE('{filter_date}')"

        parameters = parameters or {}
        for k, v in parameters.items():
            where_or_and = 'WHERE' if 'WHERE' not in query else 'AND'
            if isinstance(v, list):
                comma_separated = "','".join(v)
                query += f" {where_or_and} {k} IN ('{comma_separated}')"
            elif isinstance(v, str) and k == 't_environment':
                query += f" {where_or_and} {k} LIKE '%{v}'"
            elif isinstance(v, str):
                query += f" {where_or_and} {k} = '{v}'"
            else:
                self._logger.warning(f'Invalid SQL query parameter with type "{type(v)}": v')

        if order_by:
            query += f' ORDER BY {order_by} {"DESC" if order_desc else "ASC"}'
        if limit_results:
            query += f' LIMIT {limit_results}'
        return f'{query};'

    @log_it
    def data_to_row(self, data) -> Tuple:
        row = [
            datum.get('VarCharValue', '')
            for datum in data['Data']
        ]
        return row
