import os
import argparse

file_stats = []

parser = argparse.ArgumentParser(description = 'Automatically sort media from a directory into folders by media type in another directory.')
parser.add_argument('--source', required = True, help = 'The absolute path of the source directory containing files to be sorted.')
parser.add_argument('--destination', required = True, help = 'An empty destination directory.')
args = parser.parse_args()

if not os.path.isdir(args.source):
  exit('Source is not a directory')

if not os.path.isdir(args.destination):
  exit('Destination is not a directory')

def file_category(extension: str):
  extension = extension.lower()
  images = ['.jpg', '.png', '.gif', '.bmp', '.webp']
  videos = ['.mp4', '.3gp', '.webm']
  if extension in images:
    return "images"
  elif extension in videos:
    return "videos"
  else:
    return False

for root, dirs, files in os.walk(args.source):

  for name in files:
    filepath = os.path.join(root, name)
    existing_file = next((item for item in file_stats if item['name'] == name), False)
    size = os.path.getsize(filepath)
    filename, file_extension = os.path.split(name)
    file_category = file_category(file_extension)

    if not existing_file:
      # Store some file stats so we can compare duplicates later.
      file_stats.append({'root': root, 'name': name, 'size': size})

      if not file_category:
        try:
          # Put the file at the base of the destination.
          os.rename(filepath, os.path.join(args.destination, name))
        except:
          print('Unable to move file %s' % (name))

      else:
        try:
          # Put the file in the appropriate dir under the destination.
          os.rename(filepath, os.path.join(args.destination, file_category, name))
        except:
          print('Unable to move file %s' % (name))

    else:
      if size != existing_file['size']:

        if not file_category:
          try:
            # Put the file at the base of the destination.
            os.rename(filepath, os.path.join(args.destination, name))
          except:
            print('Unable to move file %s' % (name))

        else:
          try:
          # Put the file in the appropriate dir under the destination.
            os.rename(filepath, os.path.join(args.destination, file_category, name))
          except:
            print('Unable to move file %s' % (name))