# Import necessary libraries
import io
from os import makedirs
from os.path import splitext
import time

import boto3
import openpyxl as op

from validations.validate import (
    validate_worksheet
)

import config

logger = config.get_logger(__name__)


def expand_column_width(ws):
    """Expand column width for legibility."""
    for col in ws.columns:
        max_length = 0
        column = col[0].column
        for cell in col:
            try:
                if len(str(cell.value)) > max_length:
                    max_length = len(cell.value)
            except:  # noqa
                pass
        adjusted_width = (max_length + 2)
        ws.column_dimensions[
            ws.cell(row=1, column=column).column_letter].width = adjusted_width


def main():
    """Main function."""
    start_time = time.time()
    # Create an S3 resource object using the session
    s3 = boto3.resource('s3')

    # Specify the bucket and file name
    bucket_name = config.BUCKET
    file_name = config.KEY

    # Create a streaming object to pull data from the file
    streaming_obj = s3.Bucket(bucket_name).Object(file_name).get()['Body']

    # Load the streaming object into a BytesIO object
    stream = io.BytesIO(streaming_obj.read())
    # Load the workbook from the BytesIO object
    wb = op.load_workbook(stream)

    # Get the active worksheet
    ws = wb.active

    # Get the first worksheet
    # ws = wb[wb.sheetnames[0]]

    # Get the sheet named 'Sheet1'
    # ws = wb['Sheet1']

    # Iterate over the rows and columns where we want to add comments
    validate_worksheet(ws)

    # Add '_validated' to the file name
    file_name = splitext(file_name)[0] + '_validated.xlsx'

    # Expand the column width for legibility
    expand_column_width(ws)

    # Save the workbook locally if in debug mode
    if config.DEBUG:
        makedirs('data/output', exist_ok=True)
        wb.save('data/output/'+file_name.split('/')[-1])

    # Create a byte stream from the wb and stream write it to S3
    stream = io.BytesIO()
    wb.save(stream)
    stream.seek(0)
    s3.Bucket(bucket_name).upload_fileobj(stream, file_name)

    end_time = time.time()
    elapsed_time = end_time - start_time
    logger.info(f'Done! Elapsed time: {elapsed_time} seconds')


if __name__ == '__main__':
    main()