import json
import logging
import boto3
import os
from configparser import ConfigParser
from botocore.exceptions import ClientError

# 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 load configuration 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 construct JSON path from S3 object path
def generate_output_json_path(src_folder,dest_folder, s3_object):
    
    # split string to remove source folder path
    txt = s3_object.split(src_folder)
    if len(txt) == 2:
        txt = txt[1]
    else:
        txt = txt[0]
    # change file extension - remove .pdf and add .json 
    if txt.lower().endswith(".pdf"):
        return dest_folder+txt[0:-4]+".json"

# Function to get Textract document analysis
def get_textract_document_detection(job_id, textract_client):
    # Initialize an empty list to store blocks
    all_blocks = []
    next_token = None
    response = {}
    flag = True
    logger.info("Get document detection")
    try:
        while True:
            if flag:
                # Call the Textract API to get document analysis
                response = textract_client.get_document_text_detection(
                    JobId=job_id
                )
                flag = False
            else:
                # Call the Textract API to get document analysis
                response = textract_client.get_document_text_detection(
                    JobId=job_id,
                    NextToken=next_token
                )
            
            job_status = response["JobStatus"]
            logger.info("Job %s status is %s.", job_id, job_status)
            
            # Merge the blocks from the current response
            all_blocks.extend(response.get('Blocks', []))
            
            # Check if there are more blocks to retrieve
            next_token = response.get('NextToken')
            if not next_token:
                logger.info("No more Textract response to retrieve")
                break
            
    except ClientError:
        logger.exception("Couldn't get data for job %s.", job_id)
        raise
    else:
        # Remove unnecessary keys from the last response
        last_response = response.copy()
        last_response.pop('Blocks', None)
        last_response.pop('ResponseMetadata', None)
        logger.info("Removed 'ResponseMetadata' key")

        # Merge with {'Blocks': all_blocks}
        final_response = {'Blocks': all_blocks}
        final_response.update(last_response)
        logger.info("Final Textract response is contructed")
        return final_response
    
# Function to get Textract document analysis
def get_textract_document_analysis(job_id, textract_client):
    # Initialize an empty list to store blocks
    all_blocks = []
    next_token = None
    response = {}
    flag = True
    logger.info("Get document analysis")
    try:
        while True:
            if flag:
                # Call the Textract API to get document analysis
                response = textract_client.get_document_analysis(
                    JobId=job_id
                )
                flag = False
            else:
                # Call the Textract API to get document analysis
                response = textract_client.get_document_analysis(
                    JobId=job_id,
                    NextToken=next_token
                )
            
            job_status = response["JobStatus"]
            logger.info("Job %s status is %s.", job_id, job_status)
            
            # Merge the blocks from the current response
            all_blocks.extend(response.get('Blocks', []))
            
            # Check if there are more blocks to retrieve
            next_token = response.get('NextToken')
            if not next_token:
                logger.info("No more Textract response to retrieve")
                break
            
    except ClientError:
        logger.exception("Couldn't get data for job %s.", job_id)
        raise
    else:
        # Remove unnecessary keys from the last response
        last_response = response.copy()
        last_response.pop('Blocks', None)
        last_response.pop('ResponseMetadata', None)
        logger.info("Removed 'ResponseMetadata' key")

        # Merge with {'Blocks': all_blocks}
        final_response = {'Blocks': all_blocks}
        final_response.update(last_response)
        logger.info("Final Textract response is contructed")
        return final_response

# Function to save Textract response to S3
def upload_response_to_s3(response, bucket_name, object_key, s3_client):
    # Convert the response to JSON
    response_json = json.dumps(response)
    
    try:
        tags = "env=dev"
        # Upload the JSON response to S3
        s3_client.put_object(
            Bucket=bucket_name,
            Key=object_key,
            Body=response_json,
            ContentType='application/json'
        )
        logger.info(f"Saved analysis response to S3: {object_key}")
    except Exception as e:
        logger.error(f"Error saving response to S3 - {object_key}: {str(e)}")        

# Function to move file within S3
def move_file_within_s3(source_bucket, source_path, destination_path):
    try:
        # Copy the file to the destination folder
        s3_client.copy_object(Bucket=source_bucket, CopySource={'Bucket': source_bucket, 'Key': source_path}, Key=destination_path)

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

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

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

    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', '')

    # Validate environment variables
    file_path_array = property_file_path.split("/")
    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))

        STAGING_LOCATION = config_dict['FOLDER_LOCATIONS']['STAGING_LOCATION'].format(batch_id)
        OUTPUT_LOCATION = config_dict['FOLDER_LOCATIONS']['OUTPUT_LOCATION'].format(batch_id)
        PROCESSED_LOCATION = config_dict['FOLDER_LOCATIONS']['PROCESSED_LOCATION'].format(batch_id)
        UNPROCESSED_LOCATION = config_dict['FOLDER_LOCATIONS']['UNPROCESSED_LOCATION'].format(batch_id)
        PROCESS_TYPE = str(config_dict['OTHERS']['PROCESS_TYPE']).upper()

        logger.info('STAGGING_LOCATION: ' + STAGING_LOCATION)
        logger.info('OUTPUT_LOCATION: ' + OUTPUT_LOCATION)
        logger.info('PROCESSED_LOCATION: ' + PROCESSED_LOCATION)
        logger.info('UNPROCESSED_LOCATION: ' + UNPROCESSED_LOCATION)
        logger.info('PROCESS_TYPE: ' + str(PROCESS_TYPE))

        # 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'])
            message_body = json.loads(record_body['Message'])
            logger.info('MESSAGE_BODY: ' + str(message_body))
            try:

                # Extract relevant information from the message body
                job_id = message_body.get('JobId')
                document_location = message_body.get('DocumentLocation')
                s3_object_name = document_location.get('S3ObjectName')
                
                # Check if the status is "SUCCEEDED"
                if message_body.get('Status') == 'SUCCEEDED':

                    document = {}
                    if PROCESS_TYPE == "ANALYSIS":

                        # Call the function to get document analysis using Textract
                        document = get_textract_document_analysis(job_id, textract_client)

                    elif PROCESS_TYPE == "DETECTION":

                        # Call the function to get document analysis using Textract
                        document = get_textract_document_detection(job_id, textract_client)
                    
                    # Save the document analysis response to S3
                    s3_object_key = generate_output_json_path(STAGING_LOCATION,OUTPUT_LOCATION, s3_object_name)
                    upload_response_to_s3(document, S3_BUCKET_NAME, s3_object_key, s3_client)

                    # Construct the destination paths
                    destination_path = PROCESSED_LOCATION + s3_object_name.replace(STAGING_LOCATION,"")

                    # Move file to processed folder
                    move_file_within_s3(S3_BUCKET_NAME, s3_object_name, destination_path)

                    success_message = 'Processed file '+ str(s3_object_key)
                    logger.info(success_message)
                    return {
                        'statusCode': 200,
                        'body': success_message
                    }
                else:
                    error_message = f"Skipping message with JobId {message_body.get('JobId')} as Status is not 'SUCCEEDED'"
                    logger.info(error_message)
                    
                    # Construct the destination paths
                    destination_path = UNPROCESSED_LOCATION + s3_object_name.replace(STAGING_LOCATION,"")

                    # Move file to unprocessed folder
                    move_file_within_s3(S3_BUCKET_NAME, s3_object_name, destination_path)
                    return {
                        'statusCode': 500,
                        'body': error_message
                    }
        
            except Exception as e:
                error_message = f"Error processing message with JobId {message_body.get('JobId')}: {str(e)}"
                logger.error(error_message)  
                return {
                        'statusCode': 500,
                        'body': error_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
            }