import boto3
import time
from configparser import ConfigParser
import logging
import os
from botocore.exceptions import ClientError
import json
from urllib.parse import unquote_plus

# Initialize logger
logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)

# Initialize S3 & Textract clients
s3_client = boto3.client('s3')
textract_client = boto3.client('textract')

# Function to generate a Unix timestamp
def generate_unix_timestamp():
    # Get the current time in seconds since the epoch
    unix_timestamp = int(time.time())
    return unix_timestamp

# Function to retrieve configuration values from S3
def load_config_from_s3(bucket_name, file_key):
    
    # Download the config file from S3
    response = s3_client.get_object(Bucket=bucket_name, Key=file_key)
    config_content = response['Body'].read().decode('utf-8')

    # Parse the config file
    config_parser = ConfigParser()
    config_parser.read_string(config_content)

    # Convert the configuration to a dictionary
    config_dict = {}
    for section in config_parser.sections():
        config_dict[section] = {key.upper(): value for key, value in config_parser.items(section)}

    return config_dict

# Function to move a file from source to destination in S3
def move_file_within_s3(source_bucket, source_key, destination_key):
    try:
        tags = "env=dev"
        # Copy the file to the destination folder
        s3_client.copy_object(Bucket=source_bucket, CopySource={'Bucket': source_bucket, 'Key': source_key}, Key=destination_key, Tagging=f'{tags}')

        # Delete the file from the source folder
        s3_client.delete_object(Bucket=source_bucket, Key=source_key)

        logger.info(f"File moved from {source_key} to {destination_key}")
    except ClientError as e:
        logger.error(f"Error moving file: {e}")
    except Exception as e:
        logger.error(f"Error moving file: {e}")

# Function to get a list of PDF files in a given S3 folder
def get_pdf_files_list_from_s3(source_bucket, source_folder):
    file_list = []
    # List S3 Object & iterate (as per max files allowed)
    s3_list_response = s3_client.list_objects_v2(Bucket=source_bucket, Prefix=source_folder)

    if s3_list_response and s3_list_response['ResponseMetadata']['HTTPStatusCode'] == 200 and s3_list_response['KeyCount'] != 0:
        objects = s3_list_response['Contents']

        for s3_object in objects:
            # Skip non-PDF files
            if not s3_object['Key'].lower().endswith('.pdf'):
                continue

            file_list.append(s3_object['Key'])

    return file_list

def start_textract_detection_job( bucket_name,
        document_file_name,
        sns_topic_arn,
        sns_role_arn,
        job_tag,):
    try:
        # Define the parameters for the start_document_analysis API
        start_document_detection_params = {
            'DocumentLocation': {
                'S3Object': {
                    'Bucket': bucket_name,
                    'Name': document_file_name
                }
            },
            'ClientRequestToken': 'unique-token-'+str(generate_unix_timestamp()),  # Use a unique token for each request
            'JobTag': job_tag,  # Use a tag to identify your job
            'NotificationChannel': {
                'SNSTopicArn': sns_topic_arn,
                'RoleArn': sns_role_arn  # Role to allow Textract service to notify SNS topic when response is ready
            }
        }
        
        logger.info('start_document_detection_params ' + str(start_document_detection_params))
        
        # Send the request to start document detection
        textract_response = textract_client.start_document_text_detection(**start_document_detection_params)
        
        job_id = textract_response["JobId"]
        logger.info(
            "Started text detection job %s on %s.", job_id, document_file_name
        )
    except ClientError:
        logger.exception("Couldn't detect text in %s.", document_file_name)
        raise
    else:
        return job_id

def start_textract_analysis_job(
        bucket_name,
        document_file_name,
        analysis_feature_type,
        sns_topic_arn,
        sns_role_arn,
        job_tag,
    ):

    try:
        # Define the parameters for the start_document_analysis API
        start_document_analysis_params = {
            'DocumentLocation': {
                'S3Object': {
                    'Bucket': bucket_name,
                    'Name': document_file_name
                }
            },
            'FeatureTypes': analysis_feature_type,  # Customize based on requirements
            'ClientRequestToken': 'unique-token-'+str(generate_unix_timestamp()),  # Use a unique token for each request
            'JobTag': job_tag,  # Use a tag to identify your job
            'NotificationChannel': {
                'SNSTopicArn': sns_topic_arn,
                'RoleArn': sns_role_arn  # Role to allow Textract service to notify SNS topic when response is ready
            }
        }
        
        logger.info('start_document_analysis_params ' + str(start_document_analysis_params))
        
        # Send the request to start document analysis
        textract_response = textract_client.start_document_analysis(**start_document_analysis_params)
        
        job_id = textract_response["JobId"]
        logger.info(
            "Started text analysis job %s on %s.", job_id, document_file_name
        )
    except ClientError:
        logger.exception("Couldn't analyze text in %s.", document_file_name)
        raise
    else:
        return job_id

# AWS Lambda handler function
def lambda_handler(event, context):

    try:
        # Extract AWS account ID and region from the Lambda ARN
        aws_account_id = context.invoked_function_arn.split(":")[4]
        aws_region = context.invoked_function_arn.split(":")[3]

        logger.info('## ENVIRONMENT VARIABLES\r' + str(os.environ))

        # Read environment variables
        property_file_path = os.environ.get('PROPERTY_FILE_S3_PATH', '')
        batch_id = os.environ.get('BATCH_ID', '')

        # Read config.properties
        file_path_array = property_file_path.split("/")
        
        # Valid if file_path_array has more than 2 elements
        if len(file_path_array) > 1:
            
            # Extract BUCKET_NAME and config_file_path
            S3_BUCKET_NAME = file_path_array[0]
            CONFIG_FILE_PATH = "/".join(file_path_array[1:])
            
            logger.info(f'S3_BUCKET_NAME: {S3_BUCKET_NAME}')
            logger.info(f'CONFIG_FILE_PATH: {CONFIG_FILE_PATH}')

            # Load config file
            config_dict = load_config_from_s3(S3_BUCKET_NAME, CONFIG_FILE_PATH)

            logger.info('## CONFIG DICTIONARY\r' + str(config_dict))
           
            # Extract configuration values
            SOURCE_LOCATION = config_dict['FOLDER_LOCATIONS']['SOURCE_LOCATION'].format(batch_id)  # SOURCE_LOCATION
            STAGING_LOCATION = config_dict['FOLDER_LOCATIONS']['STAGING_LOCATION'].format(batch_id)
            ANALYSIS_FEATURE_TYPE = config_dict['OTHERS']['ANALYSIS_FEATURE_TYPE'].split(",") # Analysis FeatureType
            SENDER_MAX_FILES =  int(config_dict['OTHERS']['SENDER_MAX_FILES'])
            SNS_TOPIC_ARN = config_dict['RESOURCES']['SNS_TOPIC_ARN'].replace("{aws_region}",aws_region).replace("{aws_account_id}",aws_account_id)
            TEXTRACT_ROLE_ARN = config_dict['RESOURCES']['TEXTRACT_ROLE_ARN'].replace("{aws_account_id}",aws_account_id) # Textract IAM Role ARN to publish to SNS
            JOB_TAG = config_dict['OTHERS']['JOB_TAG']
            PROCESS_TYPE = str(config_dict['OTHERS']['PROCESS_TYPE']).upper()

            logger.info('SOURCE_LOCATION: ' + SOURCE_LOCATION)
            logger.info('STAGING_LOCATION: ' + STAGING_LOCATION)
            logger.info('ANALYSIS_FEATURE_TYPE: ' + str(ANALYSIS_FEATURE_TYPE))
            logger.info('SNS_TOPIC_ARN: ' + SNS_TOPIC_ARN)
            logger.info('TEXTRACT_ROLE_ARN: ' + TEXTRACT_ROLE_ARN)
            logger.info('SENDER_MAX_FILES: ' + str(SENDER_MAX_FILES))
            logger.info('JOB_TAG: ' + str(JOB_TAG))
            logger.info('PROCESS_TYPE: ' + str(PROCESS_TYPE))

            # File count
            file_count = 0

            # Process each message from the SQS event
            for record in event['Records']:
                
                # Extract the message body from the record
                record_body = json.loads(record['body'])

                #logger.info('Message Count: ', str(len(record_body['Records'])) )
                
                for sqs_record in record_body['Records']:

                    # Construct the source and destination paths
                    source_path = unquote_plus(sqs_record['s3']['object']['key'])
                    destination_path = STAGING_LOCATION + source_path.replace(SOURCE_LOCATION,"")

                    # Move file to stagging
                    move_file_within_s3(S3_BUCKET_NAME, source_path, destination_path)

                    job_id = ""

                    if PROCESS_TYPE == "ANALYSIS":
                        # Start Textract analysis job
                        job_id = start_textract_analysis_job (
                                S3_BUCKET_NAME,
                                destination_path,
                                ANALYSIS_FEATURE_TYPE,
                                SNS_TOPIC_ARN,
                                TEXTRACT_ROLE_ARN,
                                JOB_TAG,
                            )       
                    elif PROCESS_TYPE == "DETECTION":
                        # Start Textract detection job
                        job_id = start_textract_detection_job (
                                S3_BUCKET_NAME,
                                destination_path,
                                SNS_TOPIC_ARN,
                                TEXTRACT_ROLE_ARN,
                                JOB_TAG,
                            )      

                    file_count = file_count + 1
                    logger.info(str(file_count) + '. ' + str(source_path) + " Job Id: " + str(job_id))
                    time.sleep(1)

            success_message = 'Total files sent to textract : '+ str(file_count)
            logger.info(success_message)
            return {
                'statusCode': 200,
                'body': success_message
            }
                

        else:
            error_message = 'Incorrect value for ENVIRONMENT VARIABLES: PROPERTY_FILE_S3_PATH\r' + str(property_file_path)
            logger.error(error_message)
            
            return {
                    'statusCode': 500,
                    'body': error_message
                }

    except ClientError as e:
        # Handle specific Textract client errors
        error_message = f"Error in Textract operation: {e}"
        logger.error(error_message)
        return {
            'statusCode': 500,
            'body': error_message
        }

    except Exception as e:
        # Handle other exceptions
        error_message = f"Unexpected error: {e}"
        logger.error(error_message)
        return {
            'statusCode': 500,
            'body': error_message
        }