import json
import boto3
import psycopg2
import time
import pandas
from sqlalchemy import create_engine
from sqlalchemy import text
from concurrent.futures import ThreadPoolExecutor
from concurrent.futures import as_completed
from urllib.parse import quote_plus as urlquote
import urllib
import re
import os

SQL_SCRIPT_S3_PATH = os.environ["SQL_SCRIPT_S3_PATH"]
REDSHIFT_CLUSTER_ENDPOINT = os.environ["REDSHIFT_CLUSTER_ENDPOINT"]
REDSHIFT_IAM_ROLE = os.environ["REDSHIFT_IAM_ROLE"]
BUCKET_NAME = os.environ["SQL_SCRIPT_S3_PATH"]
REDSHIFT_USER_NAME = os.environ["REDSHIFT_USER_NAME"]
NUMBER_OF_PARALLEL_SESSIONS_LIST = os.environ["NUMBER_OF_PARALLEL_SESSIONS_LIST"]
DISABLE_RESULT_CACHE = os.environ["DISABLE_RESULT_CACHE"]
DEFAULT_OUTPUT_LIMIT = os.environ["DEFAULT_OUTPUT_LIMIT"]
MAX_NUMBER_OF_QUERIES = os.environ["MAX_NUMBER_OF_QUERIES"]
MAX_PARALLEL_SESSIONS = os.environ["MAX_PARALLEL_SESSIONS"]
QUERY_LABEL_PREFIX = os.environ["QUERY_LABEL_PREFIX"]


def connect_to_redshift(host, username):
    client = boto3.client("redshift")
    cluster_creds = client.get_cluster_credentials(
        DbUser=username,
        DbName=REDSHIFT_CLUSTER_ENDPOINT.split("/")[1],
        ClusterIdentifier=REDSHIFT_CLUSTER_ENDPOINT.split(".")[0],
    )

    connection_string = (
        "postgresql://"
        + urlquote(cluster_creds["DbUser"])
        + ":"
        + urlquote(cluster_creds["DbPassword"])
        + "@"
        + REDSHIFT_CLUSTER_ENDPOINT
    )
    return create_engine(connection_string, pool_size=0, max_overflow=-1)


def get_json_config_from_s3(script_s3_path):
    bucket, key = script_s3_path.replace("s3://", "").split("/", 1)
    obj = boto3.client("s3").get_object(Bucket=bucket, Key=key)
    return json.loads(obj["Body"].read().decode("utf-8"))


def get_sql_scripts_from_s3():

    bucket, key = SQL_SCRIPT_S3_PATH.replace("s3://", "").split("/", 1)
    obj = boto3.client("s3").get_object(Bucket=bucket, Key=key)
    script = obj["Body"].read().decode("utf-8")
    script = script.format(redshift_iam_role=REDSHIFT_IAM_ROLE, bucket_name=BUCKET_NAME)
    split_scripts = script.split(";")[:-1]
    if len(split_scripts) > int(MAX_NUMBER_OF_QUERIES):
        split_scripts = split_scripts[0 : int(MAX_NUMBER_OF_QUERIES)]
    return split_scripts


def get_sql(engine, number_of_parallel_sessions):
    sql_script = ""

    pattern = re.compile(r"limit[\s|\t|\n]+[\d]+[\s]*$", re.IGNORECASE)
    for query in get_sql_scripts_from_s3():
        if not re.search(pattern, query):
            query += " limit " + DEFAULT_OUTPUT_LIMIT
        sql_script += query + ";\n"

    if DISABLE_RESULT_CACHE == "true":
        sql_script = "set enable_result_cache_for_session to false;\n" + sql_script

    sql_script = (
        "set query_group to '"
        + QUERY_LABEL_PREFIX
        + str(number_of_parallel_sessions)
        + "';\n"
        + sql_script
    )

    df = pandas.read_sql(text(sql_script), engine)
    return df


def run_concurrency_test(number_of_parallel_sessions):
    engine = connect_to_redshift(REDSHIFT_CLUSTER_ENDPOINT, REDSHIFT_USER_NAME)
    start_time = time.time()
    try:
        with ThreadPoolExecutor(max_workers=number_of_parallel_sessions) as executor:
            futures = []
            for _ in range(number_of_parallel_sessions):
                futures.append(executor.submit(get_sql, engine, number_of_parallel_sessions))
            for future in as_completed(futures):
                rs = future.result()

    except Exception as e:
        raise e
    elapsed_time_in_secs = time.time() - start_time
    print("--- %s seconds ---" % elapsed_time_in_secs)
    return elapsed_time_in_secs


print(
    f"script:{SQL_SCRIPT_S3_PATH}, cluster:{REDSHIFT_CLUSTER_ENDPOINT},role:{REDSHIFT_IAM_ROLE},bucket:{BUCKET_NAME},user:{REDSHIFT_USER_NAME},sessions:{NUMBER_OF_PARALLEL_SESSIONS_LIST}"
)
for sessions in NUMBER_OF_PARALLEL_SESSIONS_LIST.split(","):
    number_of_parallel_sessions = int(sessions)
    if number_of_parallel_sessions <= int(MAX_PARALLEL_SESSIONS):
        print(f"running {number_of_parallel_sessions} parallel threads ..")
        run_concurrency_test(number_of_parallel_sessions)
    else:
        print(
            f"parallel sessions {number_of_parallel_sessions} exceeds maximum allowed {MAX_PARALLEL_SESSIONS} .."
        )