import logging

from peewee import IntegrityError

from data.database import Tag, Manifest, ManifestBlob, ManifestLegacyImage, db_transaction
from data.model.oci.tag import filter_to_alive_tags
from data.model.storage import lookup_repo_storages_by_content_checksum
from data.model.image import lookup_repository_images, get_image, synthesize_v1_image
from image.docker.schema1 import DockerSchema1Manifest, ManifestException

logger = logging.getLogger(__name__)

def lookup_manifest(repository_id, manifest_digest, allow_dead=False):
  """ Returns the manifest with the specified digest under the specified repository
      or None if none. If allow_dead is True, then manifests referenced by only
      dead tags will also be returned.
  """
  query = (Manifest
           .select()
           .where(Manifest.repository == repository_id)
           .where(Manifest.digest == manifest_digest))

  if not allow_dead:
    query = filter_to_alive_tags(query.join(Tag)).group_by(Manifest.id)

  try:
    return query.get()
  except Manifest.DoesNotExist:
    return None


def get_or_create_manifest(repository_id, manifest_interface_instance):
  """ Returns a tuple of the manifest in the specified repository with the matching digest
      (if it already exists) or, if not yet created, creates and returns the manifest, as well as
      if the manifest was created. Returns (None, None) if there was an error creating the manifest.
      Note that *all* blobs referenced by the manifest must exist already in the repository or this
      method will fail with a (None, None).
  """
  existing = lookup_manifest(repository_id, manifest_interface_instance.digest, allow_dead=True)
  if existing is not None:
    return existing, False

  assert len(list(manifest_interface_instance.layers)) > 0

  # TODO(jschorr): Switch this to supporting schema2 once we're ready.
  assert isinstance(manifest_interface_instance, DockerSchema1Manifest)

  # Ensure all the blobs in the manifest exist.
  digests = manifest_interface_instance.checksums
  query = lookup_repo_storages_by_content_checksum(repository_id, digests)
  blob_map = {s.content_checksum: s for s in query}
  for digest_str in manifest_interface_instance.blob_digests:
    if digest_str not in blob_map:
      logger.warning('Unknown blob `%s` under manifest `%s` for repository `%s`', digest_str,
                     manifest_interface_instance.digest, repository_id)
      return None, None

  # Determine and populate the legacy image if necessary.
  legacy_image_id = _populate_legacy_image(repository_id, manifest_interface_instance, blob_map)
  if legacy_image_id is None:
    return None, None

  legacy_image = get_image(repository_id, legacy_image_id)
  if legacy_image is None:
    return None, None

  # Create the manifest and its blobs.
  media_type = Manifest.media_type.get_id(manifest_interface_instance.content_type)
  storage_ids = {storage.id for storage in blob_map.values()}

  with db_transaction():
    # Create the manifest.
    try:
      manifest = Manifest.create(repository=repository_id,
                                 digest=manifest_interface_instance.digest,
                                 media_type=media_type,
                                 manifest_bytes=manifest_interface_instance.bytes)
    except IntegrityError:
      manifest = Manifest.get(repository=repository_id, digest=manifest_interface_instance.digest)
      return manifest, False

    # Insert the blobs.
    blobs_to_insert = [dict(manifest=manifest, repository=repository_id,
                            blob=storage_id) for storage_id in storage_ids]
    if blobs_to_insert:
      ManifestBlob.insert_many(blobs_to_insert).execute()

    # Set the legacy image (if applicable).
    ManifestLegacyImage.create(repository=repository_id, image=legacy_image, manifest=manifest)

    return manifest, True


def _populate_legacy_image(repository_id, manifest_interface_instance, blob_map):
  # Lookup all the images and their parent images (if any) inside the manifest.
  # This will let us know which v1 images we need to synthesize and which ones are invalid.
  docker_image_ids = list(manifest_interface_instance.legacy_image_ids)
  images_query = lookup_repository_images(repository_id, docker_image_ids)
  image_storage_map = {i.docker_image_id: i.storage for i in images_query}

  # Rewrite any v1 image IDs that do not match the checksum in the database.
  try:
    rewritten_images = manifest_interface_instance.rewrite_invalid_image_ids(image_storage_map)
    rewritten_images = list(rewritten_images)
    parent_image_map = {}

    for rewritten_image in rewritten_images:
      if not rewritten_image.image_id in image_storage_map:
        parent_image = None
        if rewritten_image.parent_image_id:
          parent_image = parent_image_map.get(rewritten_image.parent_image_id)
          if parent_image is None:
            parent_image = get_image(repository_id, rewritten_image.parent_image_id)
            if parent_image is None:
              return None

        synthesized = synthesize_v1_image(
          repository_id,
          blob_map[rewritten_image.content_checksum].id,
          blob_map[rewritten_image.content_checksum].image_size,
          rewritten_image.image_id,
          rewritten_image.created,
          rewritten_image.comment,
          rewritten_image.command,
          rewritten_image.compat_json,
          parent_image,
        )

        parent_image_map[rewritten_image.image_id] = synthesized
  except ManifestException:
    logger.exception("exception when rewriting v1 metadata")
    return None

  return rewritten_images[-1].image_id