import boto3
import csv
import os
import pandas as pd

def list_s3_files(bucket_name, folder):
    s3 = boto3.Session(profile_name='temp_cred').client('s3')
    paginator = s3.get_paginator('list_objects_v2')
    operation_parameters = {'Bucket': bucket_name, 'Prefix': folder}
    file_names = set()

    for page in paginator.paginate(**operation_parameters):
        if 'Contents' in page:
            for content in page['Contents']:
                file_name = content['Key']
                if not file_name.endswith('/'):
                    file_name_without_extension = file_name.replace(folder, '', 1)[:-4]
                    # file_name_without_extension = file_name[:-4]
                    file_names.add(file_name_without_extension)

    return file_names

def list_local_files(path):
    files = set()
    try:
        for _, _, filenames in os.walk(path):
            for filename in filenames:
                files.add(filename[:-4])
    except Exception as e:
        print(f"Error accessing directory {path}: {e}")
    return files

def compare_s3_folders(bucket_name, source1, folder1, source2, folder2, output_csv):

    if source1 == 's3':
        files_in_folder1 = list_s3_files(bucket_name, folder1)
    elif source1 == 'local':
        files_in_folder1 = list_local_files(folder1)
    elif source1 == 'csv':
        df = pd.read_csv(folder1,encoding='utf-8')
        # df['File Name'] = df['File Name'][:-4]
        files_in_folder1 = set(df['File Name'].to_list())
    elif source1 == 'xlsb':
        with pd.ExcelFile(folder1, engine='pyxlsb') as xlsb:
            first_sheet = xlsb.sheet_names[0]
            df = xlsb.parse(first_sheet)
            files_in_folder1 = set(df['Contract Name'].to_list())

    if source2 == 's3':
        files_in_folder2 = list_s3_files(bucket_name, folder2)
    elif source2 == 'local':
        files_in_folder2 = list_local_files(folder2)
    elif source2 == 'csv':
        df = pd.read_csv(folder2,encoding='utf-8')
        # df['File Name'] = df['File Name'][:-4]
        files_in_folder2 = set(df['File Name'].to_list())
    elif source2 == 'xlsb':
        with pd.ExcelFile(folder2, engine='pyxlsb') as xlsb:
            first_sheet = xlsb.sheet_names[0]
            df = xlsb.parse(first_sheet)
            files_in_folder2 = set(df['Contract Name'].to_list())


    common_files = files_in_folder1 & files_in_folder2
    only_in_folder1 = files_in_folder1 - files_in_folder2
    only_in_folder2 = files_in_folder2 - files_in_folder1

    with open(output_csv, 'w', newline='', encoding='utf-8') as csvfile:
        csv_writer = csv.writer(csvfile)
        csv_writer.writerow(['Common Files', f'Only in {folder1}', f'Only in {folder2}'])

        max_length = max(len(common_files), len(only_in_folder1), len(only_in_folder2))
        for i in range(max_length):
            row = [
                list(common_files)[i] if i < len(common_files) and list(common_files)[i] else '',
                list(only_in_folder1)[i] if i < len(only_in_folder1) and list(only_in_folder1)[i] else '',
                list(only_in_folder2)[i] if i < len(only_in_folder2) and list(only_in_folder2)[i]  else ''
            ]
            csv_writer.writerow(row)

bucket_name = 'centene-national-contracting-files'
source1_type = 'local' # local / s3 / csv / xlsb
source1 = 'Batch 6 TXT Files'
source2_type = 'csv'   # local / s3 / csv / xlsb
source2 = 'Batch 6 Outputs'
output_csv = 'batch6/batch6_outputs_diff_with_tracker.csv'

compare_s3_folders(bucket_name, source1_type, source1, source2_type, source2, output_csv)
print(f'Comparison results have been written to {output_csv}')