import json

from playhouse.test_utils import assert_query_count

from app import docker_v2_signing_key, storage

from digest.digest_tools import sha256_digest
from data.database import (Tag, ManifestBlob, ImageStorageLocation, ManifestChild,
                           ImageStorage, Image, RepositoryTag, get_epoch_timestamp_ms)
from data.model.oci.manifest import lookup_manifest, get_or_create_manifest
from data.model.oci.tag import filter_to_alive_tags, get_tag
from data.model.oci.shared import get_legacy_image_for_manifest
from data.model.oci.label import list_manifest_labels
from data.model.oci.retriever import RepositoryContentRetriever
from data.model.repository import get_repository, create_repository
from data.model.image import find_create_or_link_image
from data.model.blob import store_blob_record_and_temp_link
from data.model.storage import get_layer_path
from image.docker.schema1 import DockerSchema1ManifestBuilder, DockerSchema1Manifest
from image.docker.schema2.manifest import DockerSchema2ManifestBuilder
from image.docker.schema2.list import DockerSchema2ManifestListBuilder
from util.bytes import Bytes

from test.fixtures import *

def test_lookup_manifest(initialized_db):
  found = False
  for tag in filter_to_alive_tags(Tag.select()):
    found = True
    repo = tag.repository
    digest = tag.manifest.digest
    with assert_query_count(1):
      assert lookup_manifest(repo, digest) == tag.manifest

  assert found

  for tag in Tag.select():
    repo = tag.repository
    digest = tag.manifest.digest
    with assert_query_count(1):
      assert lookup_manifest(repo, digest, allow_dead=True) == tag.manifest


def test_lookup_manifest_dead_tag(initialized_db):
  dead_tag = Tag.select().where(Tag.lifetime_end_ms <= get_epoch_timestamp_ms()).get()
  assert dead_tag.lifetime_end_ms <= get_epoch_timestamp_ms()

  assert lookup_manifest(dead_tag.repository, dead_tag.manifest.digest) is None
  assert (lookup_manifest(dead_tag.repository, dead_tag.manifest.digest, allow_dead=True) ==
          dead_tag.manifest)


def create_manifest_for_testing(repository, differentiation_field='1'):
  # Populate a manifest.
  layer_json = json.dumps({
    'config': {},
    "rootfs": {
      "type": "layers",
      "diff_ids": []
    },
    "history": [],
  })

  # Add a blob containing the config.
  _, config_digest = _populate_blob(layer_json)

  remote_digest = sha256_digest('something')
  builder = DockerSchema2ManifestBuilder()
  builder.set_config_digest(config_digest, len(layer_json))
  builder.add_layer(remote_digest, 1234, urls=['http://hello/world' + differentiation_field])
  manifest = builder.build()

  created = get_or_create_manifest(repository, manifest, storage)
  assert created
  return created.manifest, manifest


def test_lookup_manifest_child_tag(initialized_db):
  repository = create_repository('devtable', 'newrepo', None)
  manifest, manifest_impl = create_manifest_for_testing(repository)

  # Mark the hidden tag as dead.
  hidden_tag = Tag.get(manifest=manifest, hidden=True)
  hidden_tag.lifetime_end_ms = hidden_tag.lifetime_start_ms
  hidden_tag.save()

  # Ensure the manifest cannot currently be looked up, as it is not pointed to by an alive tag.
  assert lookup_manifest(repository, manifest.digest) is None
  assert lookup_manifest(repository, manifest.digest, allow_dead=True) is not None

  # Populate a manifest list.
  list_builder = DockerSchema2ManifestListBuilder()
  list_builder.add_manifest(manifest_impl, 'amd64', 'linux')
  manifest_list = list_builder.build()

  # Write the manifest list, which should also write the manifests themselves.
  created_tuple = get_or_create_manifest(repository, manifest_list, storage)
  assert created_tuple is not None

  # Since the manifests are not yet referenced by a tag, they cannot be found.
  assert lookup_manifest(repository, manifest.digest) is None
  assert lookup_manifest(repository, manifest_list.digest) is None

  # Unless we ask for "dead" manifests.
  assert lookup_manifest(repository, manifest.digest, allow_dead=True) is not None
  assert lookup_manifest(repository, manifest_list.digest, allow_dead=True) is not None


def _populate_blob(content):
  digest = str(sha256_digest(content))
  location = ImageStorageLocation.get(name='local_us')
  blob = store_blob_record_and_temp_link('devtable', 'newrepo', digest, location,
                                         len(content), 120)
  storage.put_content(['local_us'], get_layer_path(blob), content)
  return blob, digest


@pytest.mark.parametrize('schema_version', [
  1,
  2,
])
def test_get_or_create_manifest(schema_version, initialized_db):
  repository = create_repository('devtable', 'newrepo', None)

  expected_labels = {
    'Foo': 'Bar',
    'Baz': 'Meh',
  }

  layer_json = json.dumps({
    'id': 'somelegacyid',
    'config': {
      'Labels': expected_labels,
    },
    "rootfs": {
      "type": "layers",
      "diff_ids": []
    },
    "history": [
      {
        "created": "2018-04-03T18:37:09.284840891Z",
        "created_by": "do something",
      },
    ],
  })

  # Create a legacy image.
  find_create_or_link_image('somelegacyid', repository, 'devtable', {}, 'local_us')

  # Add a blob containing the config.
  _, config_digest = _populate_blob(layer_json)

  # Add a blob of random data.
  random_data = 'hello world'
  _, random_digest = _populate_blob(random_data)

  # Build the manifest.
  if schema_version == 1:
    builder = DockerSchema1ManifestBuilder('devtable', 'simple', 'anothertag')
    builder.add_layer(random_digest, layer_json)
    sample_manifest_instance = builder.build(docker_v2_signing_key)
  elif schema_version == 2:
    builder = DockerSchema2ManifestBuilder()
    builder.set_config_digest(config_digest, len(layer_json))
    builder.add_layer(random_digest, len(random_data))
    sample_manifest_instance = builder.build()

  # Create a new manifest.
  created_manifest = get_or_create_manifest(repository, sample_manifest_instance, storage)
  created = created_manifest.manifest
  newly_created = created_manifest.newly_created

  assert newly_created
  assert created is not None
  assert created.media_type.name == sample_manifest_instance.media_type
  assert created.digest == sample_manifest_instance.digest
  assert created.manifest_bytes == sample_manifest_instance.bytes.as_encoded_str()
  assert created_manifest.labels_to_apply == expected_labels

  # Verify it has a temporary tag pointing to it.
  assert Tag.get(manifest=created, hidden=True).lifetime_end_ms

  # Verify the legacy image.
  legacy_image = get_legacy_image_for_manifest(created)
  assert legacy_image is not None
  assert legacy_image.storage.content_checksum == random_digest

  # Verify the linked blobs.
  blob_digests = [mb.blob.content_checksum for mb
                  in ManifestBlob.select().where(ManifestBlob.manifest == created)]

  assert random_digest in blob_digests
  if schema_version == 2:
    assert config_digest in blob_digests

  # Retrieve it again and ensure it is the same manifest.
  created_manifest2 = get_or_create_manifest(repository, sample_manifest_instance, storage)
  created2 = created_manifest2.manifest
  newly_created2 = created_manifest2.newly_created

  assert not newly_created2
  assert created2 == created

  # Ensure it again has a temporary tag.
  assert Tag.get(manifest=created2, hidden=True).lifetime_end_ms

  # Ensure the labels were added.
  labels = list(list_manifest_labels(created))
  assert len(labels) == 2

  labels_dict = {label.key: label.value for label in labels}
  assert labels_dict == expected_labels


def test_get_or_create_manifest_invalid_image(initialized_db):
  repository = get_repository('devtable', 'simple')

  latest_tag = get_tag(repository, 'latest')
  parsed = DockerSchema1Manifest(Bytes.for_string_or_unicode(latest_tag.manifest.manifest_bytes),
                                 validate=False)

  builder = DockerSchema1ManifestBuilder('devtable', 'simple', 'anothertag')
  builder.add_layer(parsed.blob_digests[0], '{"id": "foo", "parent": "someinvalidimageid"}')
  sample_manifest_instance = builder.build(docker_v2_signing_key)

  created_manifest = get_or_create_manifest(repository, sample_manifest_instance, storage)
  assert created_manifest is None


def test_get_or_create_manifest_list(initialized_db):
  repository = create_repository('devtable', 'newrepo', None)

  expected_labels = {
    'Foo': 'Bar',
    'Baz': 'Meh',
  }

  layer_json = json.dumps({
    'id': 'somelegacyid',
    'config': {
      'Labels': expected_labels,
    },
    "rootfs": {
      "type": "layers",
      "diff_ids": []
    },
    "history": [
      {
        "created": "2018-04-03T18:37:09.284840891Z",
        "created_by": "do something",
      },
    ],
  })

  # Create a legacy image.
  find_create_or_link_image('somelegacyid', repository, 'devtable', {}, 'local_us')

  # Add a blob containing the config.
  _, config_digest = _populate_blob(layer_json)

  # Add a blob of random data.
  random_data = 'hello world'
  _, random_digest = _populate_blob(random_data)

  # Build the manifests.
  v1_builder = DockerSchema1ManifestBuilder('devtable', 'simple', 'anothertag')
  v1_builder.add_layer(random_digest, layer_json)
  v1_manifest = v1_builder.build(docker_v2_signing_key).unsigned()

  v2_builder = DockerSchema2ManifestBuilder()
  v2_builder.set_config_digest(config_digest, len(layer_json))
  v2_builder.add_layer(random_digest, len(random_data))
  v2_manifest = v2_builder.build()

  # Write the manifests.
  v1_created = get_or_create_manifest(repository, v1_manifest, storage)
  assert v1_created
  assert v1_created.manifest.digest == v1_manifest.digest

  v2_created = get_or_create_manifest(repository, v2_manifest, storage)
  assert v2_created
  assert v2_created.manifest.digest == v2_manifest.digest

  # Build the manifest list.
  list_builder = DockerSchema2ManifestListBuilder()
  list_builder.add_manifest(v1_manifest, 'amd64', 'linux')
  list_builder.add_manifest(v2_manifest, 'amd32', 'linux')
  manifest_list = list_builder.build()

  # Write the manifest list, which should also write the manifests themselves.
  created_tuple = get_or_create_manifest(repository, manifest_list, storage)
  assert created_tuple is not None

  created_list = created_tuple.manifest
  assert created_list
  assert created_list.media_type.name == manifest_list.media_type
  assert created_list.digest == manifest_list.digest

  # Ensure the child manifest links exist.
  child_manifests = {cm.child_manifest.digest: cm.child_manifest
                     for cm in ManifestChild.select().where(ManifestChild.manifest == created_list)}
  assert len(child_manifests) == 2
  assert v1_manifest.digest in child_manifests
  assert v2_manifest.digest in child_manifests

  assert child_manifests[v1_manifest.digest].media_type.name == v1_manifest.media_type
  assert child_manifests[v2_manifest.digest].media_type.name == v2_manifest.media_type


def test_get_or_create_manifest_list_duplicate_child_manifest(initialized_db):
  repository = create_repository('devtable', 'newrepo', None)

  expected_labels = {
    'Foo': 'Bar',
    'Baz': 'Meh',
  }

  layer_json = json.dumps({
    'id': 'somelegacyid',
    'config': {
      'Labels': expected_labels,
    },
    "rootfs": {
      "type": "layers",
      "diff_ids": []
    },
    "history": [
      {
        "created": "2018-04-03T18:37:09.284840891Z",
        "created_by": "do something",
      },
    ],
  })

  # Create a legacy image.
  find_create_or_link_image('somelegacyid', repository, 'devtable', {}, 'local_us')

  # Add a blob containing the config.
  _, config_digest = _populate_blob(layer_json)

  # Add a blob of random data.
  random_data = 'hello world'
  _, random_digest = _populate_blob(random_data)

  # Build the manifest.
  v2_builder = DockerSchema2ManifestBuilder()
  v2_builder.set_config_digest(config_digest, len(layer_json))
  v2_builder.add_layer(random_digest, len(random_data))
  v2_manifest = v2_builder.build()

  # Write the manifest.
  v2_created = get_or_create_manifest(repository, v2_manifest, storage)
  assert v2_created
  assert v2_created.manifest.digest == v2_manifest.digest

  # Build the manifest list, with the child manifest repeated.
  list_builder = DockerSchema2ManifestListBuilder()
  list_builder.add_manifest(v2_manifest, 'amd64', 'linux')
  list_builder.add_manifest(v2_manifest, 'amd32', 'linux')
  manifest_list = list_builder.build()

  # Write the manifest list, which should also write the manifests themselves.
  created_tuple = get_or_create_manifest(repository, manifest_list, storage)
  assert created_tuple is not None

  created_list = created_tuple.manifest
  assert created_list
  assert created_list.media_type.name == manifest_list.media_type
  assert created_list.digest == manifest_list.digest

  # Ensure the child manifest links exist.
  child_manifests = {cm.child_manifest.digest: cm.child_manifest
                     for cm in ManifestChild.select().where(ManifestChild.manifest == created_list)}
  assert len(child_manifests) == 1
  assert v2_manifest.digest in child_manifests
  assert child_manifests[v2_manifest.digest].media_type.name == v2_manifest.media_type

  # Try to create again and ensure we get back the same manifest list.
  created2_tuple = get_or_create_manifest(repository, manifest_list, storage)
  assert created2_tuple is not None
  assert created2_tuple.manifest == created_list


def test_get_or_create_manifest_with_remote_layers(initialized_db):
  repository = create_repository('devtable', 'newrepo', None)

  layer_json = json.dumps({
    'config': {},
    "rootfs": {
      "type": "layers",
      "diff_ids": []
    },
    "history": [
      {
        "created": "2018-04-03T18:37:09.284840891Z",
        "created_by": "do something",
      },
      {
        "created": "2018-04-03T18:37:09.284840891Z",
        "created_by": "do something",
      },
    ],
  })

  # Add a blob containing the config.
  _, config_digest = _populate_blob(layer_json)

  # Add a blob of random data.
  random_data = 'hello world'
  _, random_digest = _populate_blob(random_data)

  remote_digest = sha256_digest('something')

  builder = DockerSchema2ManifestBuilder()
  builder.set_config_digest(config_digest, len(layer_json))
  builder.add_layer(remote_digest, 1234, urls=['http://hello/world'])
  builder.add_layer(random_digest, len(random_data))
  manifest = builder.build()

  assert remote_digest in manifest.blob_digests
  assert remote_digest not in manifest.local_blob_digests

  assert manifest.has_remote_layer
  assert not manifest.has_legacy_image
  assert manifest.get_schema1_manifest('foo', 'bar', 'baz', None) is None

  # Write the manifest.
  created_tuple = get_or_create_manifest(repository, manifest, storage)
  assert created_tuple is not None

  created_manifest = created_tuple.manifest
  assert created_manifest
  assert created_manifest.media_type.name == manifest.media_type
  assert created_manifest.digest == manifest.digest

  # Verify the legacy image.
  legacy_image = get_legacy_image_for_manifest(created_manifest)
  assert legacy_image is None

  # Verify the linked blobs.
  blob_digests = {mb.blob.content_checksum for mb
                  in ManifestBlob.select().where(ManifestBlob.manifest == created_manifest)}

  assert random_digest in blob_digests
  assert config_digest in blob_digests
  assert remote_digest not in blob_digests


def create_manifest_for_testing(repository, differentiation_field='1', include_shared_blob=False):
  # Populate a manifest.
  layer_json = json.dumps({
    'config': {},
    "rootfs": {
      "type": "layers",
      "diff_ids": []
    },
    "history": [],
  })

  # Add a blob containing the config.
  _, config_digest = _populate_blob(layer_json)

  remote_digest = sha256_digest('something')
  builder = DockerSchema2ManifestBuilder()
  builder.set_config_digest(config_digest, len(layer_json))
  builder.add_layer(remote_digest, 1234, urls=['http://hello/world' + differentiation_field])

  if include_shared_blob:
    _, blob_digest = _populate_blob('some data here')
    builder.add_layer(blob_digest, 4567)

  manifest = builder.build()

  created = get_or_create_manifest(repository, manifest, storage)
  assert created
  return created.manifest, manifest


def test_retriever(initialized_db):
  repository = create_repository('devtable', 'newrepo', None)

  layer_json = json.dumps({
    'config': {},
    "rootfs": {
      "type": "layers",
      "diff_ids": []
    },
    "history": [
      {
        "created": "2018-04-03T18:37:09.284840891Z",
        "created_by": "do something",
      },
      {
        "created": "2018-04-03T18:37:09.284840891Z",
        "created_by": "do something",
      },
    ],
  })

  # Add a blob containing the config.
  _, config_digest = _populate_blob(layer_json)

  # Add a blob of random data.
  random_data = 'hello world'
  _, random_digest = _populate_blob(random_data)

  # Add another blob of random data.
  other_random_data = 'hi place'
  _, other_random_digest = _populate_blob(other_random_data)

  remote_digest = sha256_digest('something')

  builder = DockerSchema2ManifestBuilder()
  builder.set_config_digest(config_digest, len(layer_json))
  builder.add_layer(other_random_digest, len(other_random_data))
  builder.add_layer(random_digest, len(random_data))
  manifest = builder.build()

  assert config_digest in manifest.blob_digests
  assert random_digest in manifest.blob_digests
  assert other_random_digest in manifest.blob_digests

  assert config_digest in manifest.local_blob_digests
  assert random_digest in manifest.local_blob_digests
  assert other_random_digest in manifest.local_blob_digests

  # Write the manifest.
  created_tuple = get_or_create_manifest(repository, manifest, storage)
  assert created_tuple is not None

  created_manifest = created_tuple.manifest
  assert created_manifest
  assert created_manifest.media_type.name == manifest.media_type
  assert created_manifest.digest == manifest.digest

  # Verify the linked blobs.
  blob_digests = {mb.blob.content_checksum for mb
                  in ManifestBlob.select().where(ManifestBlob.manifest == created_manifest)}

  assert random_digest in blob_digests
  assert other_random_digest in blob_digests
  assert config_digest in blob_digests

  # Delete any Image rows linking to the blobs from temp tags.
  for blob_digest in blob_digests:
    storage_row = ImageStorage.get(content_checksum=blob_digest)
    for image in list(Image.select().where(Image.storage == storage_row)):
      all_temp = all([rt.hidden for rt
                      in RepositoryTag.select().where(RepositoryTag.image == image)])
      if all_temp:
        RepositoryTag.delete().where(RepositoryTag.image == image).execute()
        image.delete_instance(recursive=True)

  # Verify the blobs in the retriever.
  retriever = RepositoryContentRetriever(repository, storage)
  assert (retriever.get_manifest_bytes_with_digest(created_manifest.digest) ==
          manifest.bytes.as_encoded_str())

  for blob_digest in blob_digests:
    assert retriever.get_blob_bytes_with_digest(blob_digest) is not None