import boto3
import os
import pandas as pd
from datetime import datetime
from concurrent.futures import ThreadPoolExecutor, as_completed
from io import StringIO

PROFILE_NAME = 'temp_cred'
SOURCE_BUCKET = 'doczyai-use2-u-cn1-s3-textract-processing-001'
DESTINATION_BUCKET = 'healthnet-json-files'
ERROR_LOG_FILE = "error_log_files.txt"
MAX_WORKERS = 200
CSV_FILE_PATH = "processed_batch_ids/large-pdf-files-processed-3.csv"
DESTINATION_PATHS = {
    "text": 'new-batch6/large-pdf-txt-files/',
    "stuck": 'new-batch6/stuck-files/',
    "pdf": 'new-batch6/large-pdf-processed-files/',
    "invalid": 'new-batch6/invalid-files/',
    "unprocessed": 'new-batch6/unprocessed-files/'
}

session = boto3.Session(profile_name=PROFILE_NAME)
s3_client = session.client('s3')

def log_error_to_file(key, error_message):
    with open(ERROR_LOG_FILE, 'a') as f:
        f.write(f"Error with file: {key} - {error_message}\n")

def copy_single_file(key, destination_prefix):
    destination_prefix = DESTINATION_PATHS[destination_prefix]
    try:
        copy_source = {'Bucket': SOURCE_BUCKET, 'Key': key}
        destination_key = os.path.join(destination_prefix, os.path.basename(key))
        s3_client.copy(copy_source, DESTINATION_BUCKET, destination_key)
        print(f"Copied {key} to {destination_key}")
    except Exception as e:
        print(f"Error copying {key} to {destination_key}")
        log_error_to_file(key, str(e))

def main():
    response = s3_client.get_object(Bucket=DESTINATION_BUCKET, Key=CSV_FILE_PATH)
    csv_content = response['Body'].read().decode('utf-8')
    batch_ids_df = pd.read_csv(StringIO(csv_content))
    batch_ids = batch_ids_df['batch_id'].unique()
    contract_names = batch_ids_df['contract_name']
    
    print(f"Total number of batches: {len(batch_ids)}")
    file_keys = []
    for batch_id, contract_name in zip(batch_ids_df['batch_id'], contract_names):
        contract_name = contract_name[:-4]
        txt_file_key = f'contract-text-file/{batch_id}/{contract_name}.txt'
        stuck_file_key = f'all-pdfs/{batch_id}/{contract_name}.pdf'
        pdf_file_key = f'textract-receiver-processed-pdfs/{batch_id}/{contract_name}.pdf'
        invalid_file_key = f'invalid-pdfs/{batch_id}/{contract_name}.pdf'
        unprocessed_file_key = f'textract-sender-staging-pdfs/{batch_id}/{contract_name}.pdf'
        file_keys.append((txt_file_key, 'text'))
        file_keys.append((stuck_file_key, 'stuck'))
        file_keys.append((pdf_file_key, 'pdf'))
        file_keys.append((invalid_file_key, 'invalid'))
        file_keys.append((unprocessed_file_key, 'unprocessed'))

    with ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor:
        futures = [executor.submit(copy_single_file, key, dest) for key, dest in file_keys]
        for future in as_completed(futures):
            future.result()
    
    print("Copy operation completed.")

if __name__ == '__main__':
    main()