import logging
from datetime import datetime
from io import StringIO
from typing import Optional

import boto3
import pandas

from pcm_vulnerabilities_api import get_logger, log_it


class QualysAthenaClient:
    def __init__(
        self, logger: Optional[logging.Logger] = None, session: Optional[boto3.Session] = None
    ):
        self._logger = logger or get_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: Optional[str] = None,
        catalog: Optional[str] = None,
        workgroup: Optional[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: Optional[int] = None,
        database: Optional[str] = None,
        catalog: Optional[str] = None,
        workgroup: Optional[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: Optional[int] = None,
        database: Optional[str] = None,
        catalog: Optional[str] = None,
        workgroup: Optional[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: Optional[dict] = None,
        filter_date: Optional[str] = None,
        limit_results: Optional[int] = None,
        database: Optional[str] = None,
        catalog: Optional[str] = None,
        workgroup: Optional[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: Optional[dict] = None,
        filter_date: Optional[str] = None,
        limit_results: Optional[int] = None,
        database: Optional[str] = None,
        catalog: Optional[str] = None,
        workgroup: Optional[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: Optional[int] = None,
        database: Optional[str] = None,
        catalog: Optional[str] = None,
        workgroup: Optional[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: Optional[datetime] = None,
        limit_results: Optional[int] = None,
        database: Optional[str] = None,
        catalog: Optional[str] = None,
        workgroup: Optional[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: Optional[datetime] = None,
        limit_results: Optional[int] = None,
        database: Optional[str] = None,
        catalog: Optional[str] = None,
        workgroup: Optional[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: Optional[dict] = None,
        filter_date: Optional[str] = None,
        filter_date_field: Optional[str] = None,
        order_by: Optional[str] = None,
        order_desc: bool = False,
        limit_results: Optional[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
