import boto3
import re
import csv
from concurrent.futures import ThreadPoolExecutor

"""
This script reads all the text files in a specific S3 bucket and prefix, and extracts the last page number from each file.
The output is written to a CSV file with the filename and the last page number.

This script is not used often and was needed for an adhoc analysis. It is not part of the regular data pipeline.

"""



BUCKET_NAME = "centene-national-contracting-files"
PREFIX = "batch_1_priority_files/text_files/"
OUTPUT_CSV = "batch1/batch1_s3_files_page_counts.csv"

s3_client = boto3.Session(profile_name='temp_cred').client('s3')

def process_file(s3_key):
    try:
        response = s3_client.get_object(Bucket=BUCKET_NAME, Key=s3_key)
        content = response['Body'].read().decode('utf-8')

        matches = re.findall(r"Start of Page No\. = (\d+)", content)

        last_page = matches[-1] if matches else None

        return s3_key.split('/')[-1], last_page
    except Exception as e:
        print(f"Error processing {s3_key}: {e}")
        return s3_key.split('/')[-1], None
    
def list_s3_objects(bucket, prefix):
    paginator = s3_client.get_paginator('list_objects_v2')
    for page in paginator.paginate(Bucket=bucket, Prefix=prefix):
        for obj in page.get('Contents', []):
            if obj['Key'].endswith('.txt'):
                yield obj['Key']

try:
    files = list(list_s3_objects(BUCKET_NAME, PREFIX))
except Exception as e:
    print(f"Error listing objects in S3: {e}")

results = []
with ThreadPoolExecutor(max_workers= 50) as executor:
    futures = [executor.submit(process_file, file) for file in files]
    for future in futures:
        results.append(future.result())

with open(OUTPUT_CSV, mode='w', newline='', encoding='utf-8') as csvfile:
    writer = csv.writer(csvfile)
    writer.writerow(["Filename", "Page Count"])
    writer.writerows(results)

print(f"Output written to {OUTPUT_CSV}")