import argparse
import logging

from data import model
from data.database import ImageStoragePlacement, ImageStorageLocation
from app import storage


logger = logging.getLogger(__name__)


PATHSPECS = [
  (storage.image_json_path, True),
  (storage.image_layer_path, True),
  (storage.image_ancestry_path, True),
  (storage.image_file_trie_path, False),
  (storage.image_file_diffs_path, False),
]


def migrate_image(image, destination_location):
  logger.debug('Migrating image: %s -> %s', image.docker_image_id, destination_location.name)
  destination_location_set = {destination_location.name}

  for path_func, required in PATHSPECS:
    path = path_func(image.storage.uuid)

    if storage.exists(image.storage.locations, path):
      if not storage.exists(destination_location_set, path):
        logger.debug('Migrating path: %s', path)

        with storage.stream_read_file(image.storage.locations, path) as file_to_migrate:
          storage.stream_write(destination_location_set, path, file_to_migrate)
      else:
        logger.debug('File already present in destination: %s', path)
    elif required:
      raise RuntimeError('Required file not present in image to migrate: %s', path)

  # Successfully migrated, now write the placement
  ImageStoragePlacement.create(location=destination_location, storage=image.storage)

parser = argparse.ArgumentParser(description='Replicate an image storage.')
parser.add_argument('--namespace', type=str, required=True,
                    help='Namespace for the repository containing the image to be replicated')
parser.add_argument('--repository', type=str, required=True,
                    help='Name for the repository containing the image to be replicated')
parser.add_argument('--imageid', type=str, default=None,
                    help='Specific image to migrate, entire repo will be migrated if omitted')
parser.add_argument('--to', type=str, required=True,
                    help='Storage region to which the data should be replicated')

if __name__ == "__main__":
  logging.basicConfig(level=logging.DEBUG)
  logging.getLogger('boto').setLevel(logging.CRITICAL)

  args = parser.parse_args()

  location = ImageStorageLocation.get(name=args.to)

  images = []
  if args.imageid is not None:
    images = [model.image.get_image_by_id(args.namespace, args.repository, args.imageid)]
  else:
    images = model.image.get_repository_images(args.namespace, args.repository)

  for img in images:
    migrate_image(img, location)