import os
import imageio
import boto3
from urllib.parse import urlparse
import logging
import cv2
import argparse
import common_utils as utils
from augmentation_options import Transform

s3 = boto3.resource('s3')

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# construct the argument parser and parse the arguments
ap = argparse.ArgumentParser()
ap.add_argument("-m", "--manifest_uri", required=True, help="S3 path to the manifest")
ap.add_argument("-b", "--output_s3_bucket", required=True, help="S3 bucket to upload result images")
ap.add_argument("-d", "--working_directory", required=False, help="the root directory to store video and frames",
                default=".")




def flip(image_local_path, working_dir, x_axis=True):
    img = imageio.imread(image_local_path)
    flipped_image = cv2.flip(img, 1 if x_axis else 0)
    fname = os.path.split(image_local_path)[1]

    # If original file is "test-1.jpg", a x-flipped file will be named "test-1-x-flip.jpg"
    flipped_fname_suffix = '-x-flip.' if x_axis else '-y-flip.'
    flipped_fname = fname.replace('.', flipped_fname_suffix)

    flipped_fpath = os.path.join(working_dir, flipped_fname)
    imageio.imwrite(flipped_fpath, flipped_image)
    logger.info("wrote image to {}".format(flipped_fpath))
    return flipped_fpath


def rotate(image_local_path, working_dir, cw=True):
    img = imageio.imread(image_local_path)
    rotated = cv2.transpose(img)
    rotated = cv2.flip(rotated, 1 if cw else 0)
    fname = os.path.split(image_local_path)[1]

    # If original file is "test-1.jpg", a cw rotated file will be named "test-1-cw-rotate.jpg"
    rotated_fname_suffix = "-cw-rotate." if cw else '-ccw-rotate.'
    rotated_fname = fname.replace('.', rotated_fname_suffix)

    rotated_fpath = os.path.join(working_dir, rotated_fname)
    imageio.imwrite(rotated_fpath, rotated)
    logger.info("wrote image to {}".format(rotated_fpath))
    return rotated_fpath


def transform_and_upload(transformation, image_path, working_directory, s3_bucket, s3_prefix, cleanup):
    
    if transformation is Transform.X_FLIP:
        transformed = flip(image_path, working_directory, x_axis=True)
    elif transformation is Transform.Y_FLIP:
        transformed = flip(image_path, working_directory, x_axis=False)
    elif transformation is Transform.CW_ROTATE:
        transformed = rotate(image_path, working_directory, cw=True)
    else:
        transformed = rotate(image_path, working_directory, cw=False)
    utils.upload_file(transformed, s3_bucket, s3_prefix)
    if cleanup:
        os.remove(transformed)


def transform_img(img_s3_uri, working_directory, output_s3_bucket, cleanup=True):
    o = urlparse(img_s3_uri)
    s3_bucket = o.netloc
    s3_key = o.path.lstrip('/')

    img_fname = os.path.split(s3_key)[1]
    img_id = os.path.splitext(img_fname)[0]
    logger.info("image: {}".format(img_id))

    img_ccw = 'frames/ccw/' + img_id + '-ccw-rotate.jpg'
    
    if utils.exists_in_s3(output_s3_bucket, img_ccw):
        logger.info("augmentation already exists: s3://{}/{}".format(output_s3_bucket, img_ccw))
    else:
        logger.info("augmentation does not exist: s3://{}/{}".format(output_s3_bucket, img_ccw))
        image_path = utils.download_file(img_s3_uri, working_directory)
        transform_and_upload(transformation=Transform.X_FLIP, image_path=image_path,
                             working_directory=working_directory, s3_bucket=output_s3_bucket, s3_prefix='frames/x-flipped',
                             cleanup=cleanup)
        transform_and_upload(transformation=Transform.Y_FLIP, image_path=image_path,
                             working_directory=working_directory, s3_bucket=output_s3_bucket, s3_prefix='frames/y-flipped',
                             cleanup=cleanup)
        transform_and_upload(transformation=Transform.CW_ROTATE, image_path=image_path,
                             working_directory=working_directory, s3_bucket=output_s3_bucket, s3_prefix='frames/cw',
                             cleanup=cleanup)
        transform_and_upload(transformation=Transform.CCW_ROTATE, image_path=image_path,
                             working_directory=working_directory, s3_bucket=output_s3_bucket, s3_prefix='frames/ccw',
                             cleanup=cleanup)
        if cleanup:
            os.remove(image_path)


def main():
    args = vars(ap.parse_args())
    manifest_s3_uri = args["manifest_uri"]

    output_s3_bucket = args['output_s3_bucket']
    
    working_directory = args["working_directory"]
    if not os.path.isdir(working_directory):
        utils.create_tmp_dir(working_directory)
    logger.info("working directory: {}".format(working_directory))

    manifest_file = utils.download_file(manifest_s3_uri, working_dir=working_directory)
    manifest_lines = utils.read_manifest_file(manifest_file)
    print(len(manifest_lines))
    for line in manifest_lines:
        img_s3_uri = line['source-ref']
        transform_img(img_s3_uri, working_directory, output_s3_bucket)


if __name__ == "__main__":
    main()