import boto3
import pandas as pd
from io import StringIO
import os
import json
import requests


"""
This script is used to process files from a source S3 bucket and trigger the pipeline for processing.

The script does the following:
    1. Load files from the source S3 bucket based on the directory prefix
    2. Create a batch in the client bucket
    3. Copy files to the client bucket and add tags
    4. Move the original files to a processed folder in the source bucket
    5. Create a contract list and send it to the pipeline API
    6. Update the CSV with the processed batch details
    7. Trigger the pipeline API

The script uses the following variables:
    - source_bucket: The source S3 bucket name 
        (Client bucket where all source files are located. 
        NOTE: This function will move the files to processed folder in the source bucket, so the source prefix will be empty after the execution)
    
    - api_trigger_url: The URL of the pipeline trigger API ( This explicit url can be reaplced by using AWS SDK to get the API URL)
    
    - create_batch_url: The URL of the create batch API
    
    - csv_file: The name of the CSV file to store processed batch details
        (Usually we add relevant bucket of files to track this, e.g. Facility_files.csv )
    
    - processed_csv_s3_key: The S3 key to store the processed CSV file
    
    - client_bucket: The client S3 bucket name
    
    - directory_prefix_list: A list of directory prefixes to search for files in the source bucket

    - client_name: The name of the client (for logs)

    - username: The username of the user triggering the pipeline (for logs)

The script can be scheduled to run periodically to process files from the source bucket.

Add this code to a lambda function, ensure there is a lambda layer for pandas
AWS provides a layer for pandas, you can use that or create your own layer (AWSSDKPandas-Python310)

Once the function is created, we can create an event trigger for the lambda function to run at a specific time or a a specific interval.

Note: The script assumes that the pipeline trigger API and create batch API are already set up and working correctly.

"""



source_bucket = 'centene-texas-files'
api_trigger_url = 'https://<APISTAGE>.execute-api.us-east-2.amazonaws.com/dev/trigger-pipeline'
create_batch_url = "https://<APISTAGE>.execute-api.us-east-2.amazonaws.com/dev/create-batch"
csv_file = 'ancillary_files.csv' #file count need to as prefix
processed_csv_s3_key = f'processed_batch_ids/{csv_file}'
client_bucket = 'doczyai-use2-u-cn1-s3-textract-processing-001'
client_name = 'Centene-texas'
username = '<Username>'

directory_prefix_list = [
    'for-textract/'
    # Can add more prefixes here if needed


    # 'TX-Legacy-agreements_2/batch2/for-processing/', 
    # 'TX-Legacy-agreements_2/batch2/for-processing/',
    # 'TX-Legacy-agreements_3/batch3/for-processing/',
    # 'TX-Legacy-agreements_4/batch4/for-processing/',
    # 'TX-Legacy-Amendments/TX-Legacy-Amendments/for-processing/'
]


def lambda_handler(event, context):    
    s3_resource = boto3.resource('s3')
    
    for directory_prefix in directory_prefix_list:
        bucket_obj = s3_resource.Bucket(source_bucket)
        files = [obj.key for obj in bucket_obj.objects.filter(Prefix=directory_prefix) if obj.key.endswith('.Pdf') or obj.key.endswith('.pdf')]
        
        if files:
            print(f"Processing files for directory prefix: {directory_prefix}")
            process_files(source_bucket, directory_prefix, files[:60], create_batch_url, api_trigger_url)
            break
        else:
            print(f"No files found for directory prefix: {directory_prefix}. Moving to the next prefix...")

def process_files(source_bucket, directory_prefix, files, create_batch_url, api_trigger_url):
    s3_client = boto3.client('s3')
    

    # Load existing CSV from S3 if it exists
    existing_df = None
    try:
        csv_obj = s3_client.get_object(Bucket=source_bucket, Key=processed_csv_s3_key)
        existing_df = pd.read_csv(csv_obj['Body'])
        print(f"Loaded existing CSV with {len(existing_df)} records")
    except s3_client.exceptions.NoSuchKey:
        print("No existing CSV found. A new one will be created.")

    batch_id, landing_zone = create_batch(client_bucket, create_batch_url)
    if batch_id == 'failed_cases':
        print("Batch creation failed. Skipping this batch...")
        return
    
    contract_list = []
    for s3_key in files:

        filename = os.path.basename(s3_key).replace('.Pdf', '.pdf')
        new_s3_key = f'contracts-landing-zone/{batch_id}/{filename}'
        
        # Copy the file to the new key with .pdf extension
        copy_source = {'Bucket': source_bucket, 'Key': s3_key}
        s3_client.copy_object(CopySource=copy_source, Bucket=client_bucket, Key=new_s3_key)
        
        # Add tags to the copied file
        s3_client.put_object_tagging(
            Bucket=client_bucket,
            Key=new_s3_key,
            Tagging={
                'TagSet': [
                    {
                        'Key': 'BatchId',
                        'Value': batch_id
                    }
                ]
            }
        )
        
        # Move the original file to the processed subfolder
        processed_s3_key = f'processed/{os.path.basename(s3_key)}'
        s3_client.copy_object(CopySource=copy_source, Bucket=source_bucket, Key=processed_s3_key)
        s3_client.delete_object(Bucket=source_bucket, Key=s3_key)
        
        # Add file details to contract list
        contract_list.append({
            "contract_name": filename,
            "groups": ["A", "C"],  # This should be dynamic based on your requirements
            "contract_source_path": new_s3_key
        })
    print(f"Uploaded {len(files)} files to S3 bucket for batch id {batch_id}")
    data = {
        "s3_bucket": client_bucket, 
        "batch_id": batch_id,
        "client_name": client_name,
        "username": username,
        "contract_list": contract_list
    }
    
    # Convert contract list to DataFrame and append to existing DataFrame
    new_df = pd.DataFrame(contract_list)
    new_df['batch_id'] = batch_id
    if existing_df is not None:
        existing_df = pd.concat([existing_df, new_df])
    else:
        existing_df = new_df
    
    # Upload updated CSV back to S3
    csv_buf = StringIO()
    existing_df.to_csv(csv_buf, header=True, index=False)
    csv_buf.seek(0)
    
    s3_client.put_object(Bucket=source_bucket, Body=csv_buf.getvalue(), Key=processed_csv_s3_key)
    
    print(f"Saved updated batch details to {processed_csv_s3_key} and sending to API...")
    response = requests.post(api_trigger_url, json=data)
    print(response.text)
    print(response.status_code)

    return {
        "statusCode": 200,
        "body": json.dumps("Processed files successfully")
    }

def create_batch(client_bucket, create_batch_url):
    myobj = { "client-bucket-name": client_bucket }
    response = requests.post(create_batch_url, json=myobj)
    if response.status_code >= 200 and response.status_code < 300:
        try:
            response_body = json.loads(json.loads(response.text)['body'])
            new_batch_id = response_body['batch_id']
            landing_zone = response_body['landing_zone']
        except:
            print(myobj)
            print(response.text)
            new_batch_id = 'failed_cases'
            landing_zone = 'contracts_landing_zone'
    else:
        print(response.text)
        new_batch_id = 'failed_cases'
        landing_zone = 'contracts_landing_zone'
    
    return new_batch_id, landing_zone