from util.gzipwrap import GzipWrap, GZIP_BUFFER_SIZE
from util.streamlayerformat import StreamLayerMerger
from app import app

import copy
import json
import tarfile

class FileEstimationException(Exception):
  """ Exception raised by build_docker_load_stream if the estimated size of the layer TAR
      was lower than the actual size. This means the sent TAR header is wrong, and we have
      to fail.
  """
  pass


def build_docker_load_stream(namespace, repository, tag, synthetic_image_id,
                             layer_json, get_image_iterator, get_layer_iterator):
  """ Builds and streams a synthetic .tar.gz that represents a squashed version
      of the given layers, in `docker load` V1 format.
  """
  return GzipWrap(_import_format_generator(namespace, repository, tag,
                                           synthetic_image_id, layer_json,
                                           get_image_iterator, get_layer_iterator))


def _import_format_generator(namespace, repository, tag, synthetic_image_id,
                             layer_json, get_image_iterator, get_layer_iterator):
  # Docker import V1 Format (.tar):
  #  repositories - JSON file containing a repo -> tag -> image map
  #  {image ID folder}:
  #     json - The layer JSON
  #     layer.tar - The TARed contents of the layer
  #     VERSION - The docker import version: '1.0'
  layer_merger = StreamLayerMerger(get_layer_iterator)

  # Yield the repositories file:
  synthetic_layer_info = {}
  synthetic_layer_info[tag + '.squash'] = synthetic_image_id

  hostname = app.config['SERVER_HOSTNAME']
  repositories = {}
  repositories[hostname + '/' + namespace + '/' + repository] = synthetic_layer_info

  yield _tar_file('repositories', json.dumps(repositories))

  # Yield the image ID folder.
  yield _tar_folder(synthetic_image_id)

  # Yield the JSON layer data.
  layer_json = _build_layer_json(layer_json, synthetic_image_id)
  yield _tar_file(synthetic_image_id + '/json', json.dumps(layer_json))

  # Yield the VERSION file.
  yield _tar_file(synthetic_image_id + '/VERSION', '1.0')

  # Yield the merged layer data's header.
  estimated_file_size = 0
  for image in get_image_iterator():
    estimated_file_size += image.storage.uncompressed_size

  yield _tar_file_header(synthetic_image_id + '/layer.tar', estimated_file_size)

  # Yield the contents of the merged layer.
  yielded_size = 0
  for entry in layer_merger.get_generator():
    yield entry
    yielded_size += len(entry)

  # If the yielded size is more than the estimated size (which is unlikely but possible), then
  # raise an exception since the tar header will be wrong.
  if yielded_size > estimated_file_size:
    raise FileEstimationException()

  # If the yielded size is less than the estimated size (which is likely), fill the rest with
  # zeros.
  if yielded_size < estimated_file_size:
    to_yield = estimated_file_size - yielded_size
    while to_yield > 0:
      yielded = min(to_yield, GZIP_BUFFER_SIZE)
      yield '\0' * yielded
      to_yield -= yielded

  # Yield any file padding to 512 bytes that is necessary.
  yield _tar_file_padding(estimated_file_size)

  # Last two records are empty in TAR spec.
  yield '\0' * 512
  yield '\0' * 512


def _build_layer_json(layer_json, synthetic_image_id):
  updated_json = copy.deepcopy(layer_json)
  updated_json['id'] = synthetic_image_id

  if 'parent' in updated_json:
    del updated_json['parent']

  if 'config' in updated_json and 'Image' in updated_json['config']:
    updated_json['config']['Image'] = synthetic_image_id

  if 'container_config' in updated_json and 'Image' in updated_json['container_config']:
    updated_json['container_config']['Image'] = synthetic_image_id

  return updated_json


def _tar_file(name, contents):
  length = len(contents)
  tar_data = _tar_file_header(name, length)
  tar_data += contents
  tar_data += _tar_file_padding(length)
  return tar_data


def _tar_file_padding(length):
  if length % 512 != 0:
    return '\0' * (512 - (length % 512))

  return ''

def _tar_file_header(name, file_size):
  info = tarfile.TarInfo(name=name)
  info.type = tarfile.REGTYPE
  info.size = file_size
  return info.tobuf()


def _tar_folder(name):
  info = tarfile.TarInfo(name=name)
  info.type = tarfile.DIRTYPE
  return info.tobuf()