import unittest
import json as py_json

from flask import url_for
from endpoints.api import api
from app import app, storage
from initdb import setup_database_for_testing, finished_database_for_testing
from data import model, database

ADMIN_ACCESS_USER = 'devtable'
PUBLIC_USER = 'public'

REPO = 'somerepo'

class TestGarbageColection(unittest.TestCase):
  def setUp(self):
    setup_database_for_testing(self)
    self.app = app.test_client()
    self.ctx = app.test_request_context()
    self.ctx.__enter__()

  def tearDown(self):
    finished_database_for_testing(self)
    self.ctx.__exit__(True, None, None)

  def createImage(self, docker_image_id, repository_obj, username):
    preferred = storage.preferred_locations[0]
    image = model.find_create_or_link_image(docker_image_id, repository_obj, username, {},
                                            preferred)
    image.storage.uploading = False
    image.storage.save()

    # Create derived images as well.
    for i in range(0, 2):
      model.find_or_create_derived_storage(image.storage, 'squash', preferred)

    # Add some additional placements to the image.
    for location_name in ['local_eu']:
      location = database.ImageStorageLocation.get(name=location_name)

      try:
        database.ImageStoragePlacement.get(location=location, storage=image.storage)
      except:
        continue

      database.ImageStoragePlacement.create(location=location, storage=image.storage)

    return image.storage

  def createRepository(self, namespace=ADMIN_ACCESS_USER, name=REPO, **kwargs):
    user = model.get_user(namespace)
    repo = model.create_repository(namespace, name, user)

    # Populate the repository with the tags.
    image_map = {}
    for tag_name in kwargs:
      image_ids = kwargs[tag_name]
      parent = None

      for image_id in image_ids:
        if not image_id in image_map:
          image_map[image_id] = self.createImage(image_id, repo, namespace)

        # Set the ancestors for the image.
        parent = model.set_image_metadata(image_id, namespace, name, '', '', '', parent=parent)

      # Set the tag for the image.
      model.create_or_update_tag(namespace, name, tag_name, image_ids[-1])

    return repo

  def gcNow(self, repository):
    model.garbage_collect_repository(repository.namespace_user.username, repository.name)

  def deleteTag(self, repository, tag):
    model.delete_tag(repository.namespace_user.username, repository.name, tag)
    model.garbage_collect_repository(repository.namespace_user.username, repository.name)

  def moveTag(self, repository, tag, docker_image_id):
    model.create_or_update_tag(repository.namespace_user.username, repository.name, tag, docker_image_id)
    model.garbage_collect_repository(repository.namespace_user.username, repository.name)

  def assertNotDeleted(self, repository, *args):
    for docker_image_id in args:
      self.assertTrue(bool(model.get_image_by_id(repository.namespace_user.username, repository.name, docker_image_id)))

  def assertDeleted(self, repository, *args):
    for docker_image_id in args:
      try:
        # Verify the image is missing when accessed by the repository.
        model.get_image_by_id(repository.namespace_user.username, repository.name, docker_image_id)
      except model.DataModelException as ex:
        return

      self.fail('Expected image %s to be deleted' % docker_image_id)


  def test_one_tag(self):
    """ Create a repository with a single tag, then remove that tag and verify that the repository
        is now empty. """
    repository = self.createRepository(latest = ['i1', 'i2', 'i3'])
    self.deleteTag(repository, 'latest')
    self.assertDeleted(repository, 'i1', 'i2', 'i3')

  def test_two_tags_unshared_images(self):
    """ Repository has two tags with no shared images between them. """
    repository = self.createRepository(latest = ['i1', 'i2', 'i3'], other = ['f1', 'f2'])
    self.deleteTag(repository, 'latest')
    self.assertDeleted(repository, 'i1', 'i2', 'i3')
    self.assertNotDeleted(repository, 'f1', 'f2')

  def test_two_tags_shared_images(self):
    """ Repository has two tags with shared images. Deleting the tag should only remove the
        unshared images.
    """
    repository = self.createRepository(latest = ['i1', 'i2', 'i3'], other = ['i1', 'f1'])
    self.deleteTag(repository, 'latest')
    self.assertDeleted(repository, 'i2', 'i3')
    self.assertNotDeleted(repository, 'i1', 'f1')

  def test_unrelated_repositories(self):
    """ Two repositories with different images. Removing the tag from one leaves the other's
        images intact.
    """
    repository1 = self.createRepository(latest = ['i1', 'i2', 'i3'], name = 'repo1')
    repository2 = self.createRepository(latest = ['j1', 'j2', 'j3'], name = 'repo2')

    self.deleteTag(repository1, 'latest')

    self.assertDeleted(repository1, 'i1', 'i2', 'i3')
    self.assertNotDeleted(repository2, 'j1', 'j2', 'j3')

  def test_related_repositories(self):
    """ Two repositories with shared images. Removing the tag from one leaves the other's
        images intact.
    """
    repository1 = self.createRepository(latest = ['i1', 'i2', 'i3'], name = 'repo1')
    repository2 = self.createRepository(latest = ['i1', 'i2', 'j1'], name = 'repo2')

    self.deleteTag(repository1, 'latest')

    self.assertDeleted(repository1, 'i3')
    self.assertNotDeleted(repository2, 'i1', 'i2', 'j1')

  def test_inaccessible_repositories(self):
    """ Two repositories under different namespaces should result in the images being deleted
        but not completely removed from the database.
    """
    repository1 = self.createRepository(namespace = ADMIN_ACCESS_USER, latest = ['i1', 'i2', 'i3'])
    repository2 = self.createRepository(namespace = PUBLIC_USER, latest = ['i1', 'i2', 'i3'])

    self.deleteTag(repository1, 'latest')
    self.assertDeleted(repository1, 'i1', 'i2', 'i3')
    self.assertNotDeleted(repository2, 'i1', 'i2', 'i3')


  def test_multiple_shared_images(self):
    """ Repository has multiple tags with shared images. Selectively deleting the tags, and
        verifying at each step.
    """
    repository = self.createRepository(
      latest = ['i1', 'i2', 'i3'],
      other = ['i1', 'f1', 'f2'],
      third = ['t1', 't2', 't3'],
      fourth = ['i1', 'f1'])

    # Delete tag other. Should delete f2, since it is not shared.
    self.deleteTag(repository, 'other')
    self.assertDeleted(repository, 'f2')
    self.assertNotDeleted(repository,'i1', 'i2', 'i3', 't1', 't2', 't3', 'f1')

    # Move tag fourth to i3. This should remove f1 since it is no longer referenced.
    self.moveTag(repository, 'fourth', 'i3')
    self.assertDeleted(repository, 'f1')
    self.assertNotDeleted(repository,'i1', 'i2', 'i3', 't1', 't2', 't3')

    # Delete tag 'latest'. This should do nothing since fourth is on the same branch.
    self.deleteTag(repository, 'latest')
    self.assertNotDeleted(repository,'i1', 'i2', 'i3', 't1', 't2', 't3')

    # Delete tag 'third'. This should remove t1->t3.
    self.deleteTag(repository, 'third')
    self.assertDeleted(repository, 't1', 't2', 't3')
    self.assertNotDeleted(repository,'i1', 'i2', 'i3')

    # Add tag to i1.
    self.moveTag(repository, 'newtag', 'i1')
    self.assertNotDeleted(repository,'i1', 'i2', 'i3')

    # Delete tag 'fourth'. This should remove i2 and i3.
    self.deleteTag(repository, 'fourth')
    self.assertDeleted(repository, 'i2', 'i3')
    self.assertNotDeleted(repository, 'i1')

    # Delete tag 'newtag'. This should remove the remaining image.
    self.deleteTag(repository, 'newtag')
    self.assertDeleted(repository, 'i1')

  def test_empty_gc(self):
    repository = self.createRepository(
      latest = ['i1', 'i2', 'i3'],
      other = ['i1', 'f1', 'f2'],
      third = ['t1', 't2', 't3'],
      fourth = ['i1', 'f1'])

    self.gcNow(repository)
    self.assertNotDeleted(repository,'i1', 'i2', 'i3', 't1', 't2', 't3', 'f1', 'f2')

  def test_gc_storage_empty(self):
    model.garbage_collect_storage(set())

if __name__ == '__main__':
  unittest.main()
