#!/usr/bin/env python3

import csv
import getopt
import os
import sys
import time

import boto3

MAX_POLLS_FOR_RESULTS = 300


class Configurator:
    """
    Responsible for collecting all the necessary configuration from the
    environment.
    """

    # defaults
    region = "us-east-1"
    catalog = "AwsDataCatalog"
    database = None
    output_location = None
    table_name = None
    work_group = "primary"
    logs_type = "cloudfront"
    url_filter_prefix = ""
    no_header_row = False
    outfile_path = None
    lookback_days = 3
    result_limit = 1000

    def process_args(self, argv):

        name = sys.argv[0]
        args = sys.argv[1:]

        usage_msg = " ".join([name, "-o <outfile> -d <date>"])
        short_opts = "hHl:o:d:s:"
        long_opts = ["help", "no-header-row", "limit=", "outfile=", "date=", "shift="]

        try:
            opts, formal_args = getopt.gnu_getopt(args, short_opts, long_opts)
        except getopt.GetoptError:
            print(usage_msg)
            sys.exit(2)

        for opt, arg in opts:
            if opt in ("-h", "--help"):
                print(usage_msg)
                sys.exit()
            elif opt in ("-H", "--no-header-row"):
                self.no_header_row = True
            elif opt in ("-l", "--limit"):
                self.result_limit = int(arg)
            elif opt in ("-o", "--outfile"):
                self.outfile_path = arg
            # query_date
            elif opt in ("-d", "--date"):
                self.query_date = arg
            # offset days
            elif opt in ("-s", "--shift"):
                self.lookback_days = arg
        if not (self.outfile_path or self.query_date):
            print("Requires an argument of the date and outfile.")
            sys.exit(2)

    def collect_env_vars(self, environ):
        self.region = environ.get("AWS_REGION") or self.region
        self.catalog = environ.get("ATHENA_CATALOG") or self.catalog
        self.database = environ["ATHENA_DATABASE"]
        self.output_location = environ["ATHENA_OUTPUT_LOCATION"]
        self.table_name = environ["ATHENA_TABLE"]
        self.work_group = environ.get("ATHENA_WORK_GROUP") or self.work_group
        self.logs_type = environ.get("ATHENA_LOGS_TYPE") or self.logs_type
        self.url_filter_prefix = (
            environ.get("ATHENA_URL_FILTER_PREFIX") or self.url_filter_prefix
        )


class AthenaPorcelain:
    """
    A wrapper around the basic Athena client that is equipped with query
    completion polling. It also tracks some of the necessary
    configuration.
    """

    def __init__(self, config, session):
        self.config = config
        self.session = session
        self.athena_client = session.client("athena", config.region)

    def execute_query(self, query, max_polls=MAX_POLLS_FOR_RESULTS):
        """
        Executes a query and then waits for the query to complete before
        returning the query result.
        """
        execution_id = self.start_query(query)
        self.wait_for_query_completion(execution_id, max_polls)
        results = self.athena_client.get_query_results(QueryExecutionId=execution_id)
        return results

    def start_query(self, query):
        """
        Runs a given query using the config’s values and returns the
        execution ID of the running query.
        """
        params = {
            "QueryString": query,
            "QueryExecutionContext": {
                "Database": self.config.database,
                "Catalog": self.config.catalog,
            },
            "ResultConfiguration": {"OutputLocation": self.config.output_location},
            "WorkGroup": self.config.work_group,
        }
        response = self.athena_client.start_query_execution(**params)
        execution_id = response["QueryExecutionId"]
        if not execution_id:
            print(response)
            sys.exit("No execution_id on query start")
        return execution_id

    def wait_for_query_completion(self, execution_id, max_polls=MAX_POLLS_FOR_RESULTS):
        while max_polls > 0:
            max_polls = max_polls - 1
            response = self.athena_client.get_query_execution(
                QueryExecutionId=execution_id
            )
            status = None
            if (
                "QueryExecution" in response
                and "Status" in response["QueryExecution"]
                and "State" in response["QueryExecution"]["Status"]
            ):

                status = response["QueryExecution"]["Status"]["State"]
            if status is not None and status in ["SUCCEEDED", "FAILED", "CANCELLED"]:
                break
            time.sleep(1)
            print(f"Waiting for query execution complete: {max_polls}")

        if status != "SUCCEEDED":
            print(response)
            sys.exit("Query execution was not successful.")


class OutputFormatter:
    """Output the Athena query results as a CSV."""

    def __init__(self, config, results):
        self.no_header_row = config.no_header_row
        self.headers = self._headers(results)
        self.rows = self._clean_rows(results, self.headers)

    def output_csv(self, outfile_path):
        """
        Write CSV to the given outfile path. If no path is given,
        defaults to STDOUT.
        """

        if outfile_path is None:
            file_handle = sys.stdout
        else:
            file_handle = open(outfile_path, "w")

        csv_writer = csv.writer(file_handle, quoting=csv.QUOTE_ALL)
        csv_writer.writerow(self.headers)
        for row in self.rows:
            csv_writer.writerow(row)

    def _headers(self, results):
        """Gets a list of headers from the results metadata."""

        column_info = results["ResultSet"]["ResultSetMetadata"]["ColumnInfo"]
        return [h["Name"] for h in column_info]

    def _clean_rows(self, results, headers):
        """
        Return a simpler form of the results table.

        Additional cleaning:

        1. Remove a potential header row from the body of the table.

        2. The Athena results produce a hyphen character in the
        query_string column when no query string was included in the
        logged request.
        """
        rows = list()

        for row in results["ResultSet"]["Rows"]:
            # https://forums.aws.amazon.com/thread.jspa?threadID=256505
            if row["Data"][0].get("VarCharValue", None) == headers[0]:
                continue  # skip header
            rows.append(self._clean_row(row))

        return rows

    def _clean_row(self, row):
        simple_row = [d.get("VarCharValue", None) for d in row["Data"]]
        if simple_row[2] == "-":
            simple_row[2] = None
        return simple_row


class QueryBuilder:
    def __init__(self, config):
        self.config = config

    def make_query(self):
        if self.config.logs_type == "elb":
            return self._make_query_for_elb_logs()
        else:
            return self._make_query_for_cloudfront_logs()

    def _make_query_for_cloudfront_logs(self):
        query_template = """
            SELECT method, uri, query_string, count(*) AS "count"
            FROM {0}
            WHERE method = 'GET' AND "date" between date_add('day', -{2}, CAST('{1}' AS DATE)) and CAST('{1}' AS DATE)
            AND (uri like '{3}%' or '{3}' = '')
            GROUP BY method, uri, query_string
        """
        if self.config.result_limit is None:
            query_template = query_template + ";"
            query = query_template.format(
                self.config.table_name,
                self.config.query_date,
                self.config.lookback_days,
            )
        else:
            query_template = query_template + "LIMIT {4};"
            query = query_template.format(
                self.config.table_name,
                self.config.query_date,
                self.config.lookback_days,
                self.config.url_filter_prefix,
                self.config.result_limit,
            )
        return query

    def _make_query_for_elb_logs(self):
        query_template = """
            SELECT request_verb as  method,
                    url_extract_path(request_url) as uri,
                    url_extract_query(request_url) query_string,
                    count(*) AS "count"
            FROM {0}
            WHERE  request_verb = 'GET' AND  CAST(SUBSTR(time, 1, 10) AS DATE) between date_add('day', -{2}, CAST('{1}' AS DATE)) and CAST('{1}' AS DATE)
            AND (request_url like '{3}%' or '{3}' = '')
            GROUP BY request_verb, url_extract_path(request_url), url_extract_query(request_url)
        """
        if self.config.result_limit is None:
            query_template = query_template + ";"
            query = query_template.format(
                self.config.table_name,
                self.config.query_date,
                self.config.lookback_days,
                self.config.url_filter_prefix,
            )
        else:
            query_template = query_template + "LIMIT {4};"
            query = query_template.format(
                self.config.table_name,
                self.config.query_date,
                self.config.lookback_days,
                self.config.url_filter_prefix,
                self.config.result_limit,
            )
        return query


def main():
    config = Configurator()
    config.process_args(sys.argv)
    config.collect_env_vars(os.environ)

    client = AthenaPorcelain(config, boto3.Session())

    query_builder = QueryBuilder(config)
    query = query_builder.make_query()
    results = client.execute_query(query)

    formatter = OutputFormatter(config, results)
    formatter.output_csv(config.outfile_path)


if __name__ == "__main__":
    main()