import logging
from multiprocessing import Process
import time

import boto3
import pymysql.cursors

import config

logger = logging.getLogger()
logger.setLevel(config.LOGGER_LEVEL)
logging.basicConfig(
    filename=f"{config.MYSQL_HOST}.log",
    filemode="a",
    format="%(asctime)s,%(msecs)d %(name)s %(levelname)s %(message)s",
    datefmt="%H:%M:%S",
    level=logging.INFO,
)


def get_mysql_connection():
    connection = pymysql.connect(
        host=config.MYSQL_HOST,
        user=config.MYSQL_USER,
        password=config.MYSQL_PASSWORD,
        database=config.MYSQL_DATABASE,
        charset="utf8mb4",
        cursorclass=pymysql.cursors.DictCursor,
    )

    return connection


def get_s3_objects():
    client = boto3.client("s3")
    response = client.list_objects_v2(
        Bucket=config.S3_BUCKET,
        MaxKeys=1000,
        Prefix=config.S3_OBJECT_PREFIX,
    )

    if "Contents" in response:
        object_keys = [s3_object["Key"] for s3_object in response["Contents"]]
        logger.info(
            f"Found {response['KeyCount']} objects in {config.S3_BUCKET}/{config.S3_OBJECT_PREFIX}"
        )
        return object_keys
    else:
        logger.error(
            f"No objects found in {config.S3_BUCKET}/{config.S3_OBJECT_PREFIX}"
        )
        return None


def load_data(object_key):
    connection = get_mysql_connection()
    logger.info("Got a MySQL connection")
    with connection.cursor() as cursor:
        logger.info(f"Loading data from {object_key}")
        now = time.time()
        sql = f"LOAD DATA FROM S3 FILE 's3://{object_key}' INTO TABLE accountingflat.dig_sales_detail_moose FIELDS ESCAPED BY '\\\\' TERMINATED BY '\\t' OPTIONALLY ENCLOSED BY '\"' LINES TERMINATED BY '\\n' (statement_id,date,upc,cd,track_id,isrc,track_name,qty,unit_price,total,trans_type,retail_price,original_price,discount);"
        cursor.execute(sql)
        post_query = time.time()
        logger.info(f"Query time for is {object_key} is {post_query - now}\n")
    connection.commit()
    connection.close()


def main():
    s3_objects = get_s3_objects()
    logger.info(f"objects are {s3_objects}")
    batch_size = config.MYSQL_CONCURRENT_LOADS
    for batch_start in range(0, len(s3_objects), batch_size):
        s3_object_batch = s3_objects[batch_start : batch_start + batch_size]
        logger.info(f"Smaller list is {s3_object_batch}")
        processes = [
            Process(target=load_data, args=(f"{config.S3_BUCKET}/{s3_object}",))
            for s3_object in s3_object_batch
        ]
        for process in processes:
            process.start()
        for process in processes:
            process.join()


if __name__ == "__main__":
    main()

# main()