import sys
import csv
from io import StringIO

import queries
from config import S3_BUCKET_NAME, S3_UPLOAD_FOLDER, db_config, logger
from s3_handler import S3Client
from sql_handler import SqlClient


def generate_csv_data(results, headers):
    csv_buffer = StringIO()
    writer = csv.writer(csv_buffer)
    writer.writerow([i[0] for i in headers])
    writer.writerows(results)
    return csv_buffer.getvalue()


def main():
    try:
        s3_client = S3Client(S3_BUCKET_NAME, S3_UPLOAD_FOLDER)
        sql_client = SqlClient(db_config)

        roles_and_ids, headers = sql_client.execute_query(queries.role_ids)

        # define role_id and role_name dictionary
        role_dict = dict(roles_and_ids)

        # execute user_query and retrieve results
        user_results, user_results_headers = sql_client.execute_query(
            queries.all_roles_query(role_dict),
            list(role_dict.values()) + list(role_dict.keys())
        )

        all_roles_data = generate_csv_data(user_results, user_results_headers)
        s3_client.upload_to_s3('all_roles.csv', all_roles_data)

        # iterate over role_dict to execute queries.specific_role_query for each role name
        for role_name in role_dict:

            # execute queries.specific_role_query with the current role name
            specific_role_results, specific_role_results_headers = sql_client.execute_query(
                queries.specific_role_query,
                (role_name,)
            )

            if not specific_role_results:
                logger.warning(f"No results found for {role_name}")
                continue

            role_name = role_name.lower().replace(" ", "_")
            specific_role_data = generate_csv_data(specific_role_results, specific_role_results_headers)
            s3_client.upload_to_s3(f'{role_name}_role.csv', specific_role_data)

        sql_client.close()

    except Exception as e:
        logger.error(f"Error while executing the script: {e}")
        sys.exit(1)


if __name__ == "__main__":
    main()