from multiprocessing import Process
import os

import config
from connectors import mysql

logger = config.logger


def main() -> None:
    """
    Get the subpartition count of processed_dig_sales.
    Then for each subpartition, select into outfile to S3.
    """
    database = "accountingflat"
    table = "processed_dig_sales"
    subpartitions = mysql.get_table_subpartition_count_for_period_id(
        database, table, config.PERIOD_ID
    )
    batch_size = config.MYSQL_CONCURRENT_LOADS
    for batch_start in range(0, len(subpartitions), batch_size):
        logger.info(f"Starting from subpartition {batch_start}")
        subpartition_batch = subpartitions[
            batch_start : batch_start + batch_size
        ]
        processes = []
        for subpartition in subpartition_batch:
            with open(
                os.path.realpath(os.path.dirname(__name__))
                + "/accounting_run_utils/queries/select_processed_dig_sales_into_outfile.sql"
            ) as sql_file:
                prepared_statement = sql_file.read()
            """
            Partition names in queries cannot contain quotes, so replace 
            the subpartition name in the prepared statement
            with the actual subpartition to work around that limitation.
            """
            sql = prepared_statement.replace("SUBPARTITION_NAME", subpartition)
            object_key = f"s3://{config.S3_DIM_IMPORTS_BUCKET}/{config.S3_DIM_IMPORTS_BUCKET_PREFIX}/{subpartition}.txt"
            keywords = {
                "sql_args": {
                    "period_id": config.PERIOD_ID,
                    "s3_object_path": object_key,
                }
            }
            processes.append(
                Process(
                    target=mysql.select_data_into_outfile,
                    args=(sql, object_key),
                    kwargs=keywords,
                )
            )
        for process in processes:
            process.start()
        for process in processes:
            process.join()


if __name__ == "__main__":
    main()