import json

from cStringIO import StringIO
from enum import Enum, unique

from digest.checksums import compute_simple
from test.registry.protocols import (RegistryProtocol, Failures, ProtocolOptions, PushResult,
                                     PullResult)

@unique
class V1ProtocolSteps(Enum):
  """ Defines the various steps of the protocol, for matching failures. """
  PUT_IMAGES = 'put-images'
  GET_IMAGES = 'get-images'


class V1Protocol(RegistryProtocol):
  FAILURE_CODES = {
    V1ProtocolSteps.PUT_IMAGES: {
      Failures.UNAUTHENTICATED: 401,
      Failures.UNAUTHORIZED: 403,
    },
    V1ProtocolSteps.GET_IMAGES: {
      Failures.UNAUTHENTICATED: 403,
      Failures.UNAUTHORIZED: 403,
    },
  }

  def __init__(self, jwk):
    pass

  def _auth_for_credentials(self, credentials):
    if credentials is None:
      return None

    return credentials

  def ping(self, session):
    assert session.get('/v1/_ping').status_code == 200

  def pull(self, session, namespace, repo_name, tag_names, images, credentials=None,
           expected_failure=None, options=None):
    options = options or ProtocolOptions()
    auth = self._auth_for_credentials(credentials)
    tag_names = [tag_names] if isinstance(tag_names, str) else tag_names
    prefix = '/v1/repositories/%s/%s/' % (namespace, repo_name)

    # Ping!
    self.ping(session)

    # GET /v1/repositories/{namespace}/{repository}/images
    result = self.conduct(session, 'GET', prefix + 'images', auth=auth,
                          expected_status=(200, expected_failure, V1ProtocolSteps.GET_IMAGES))
    if expected_failure is not None:
      return

    # GET /v1/repositories/{namespace}/{repository}/tags
    tags_result = self.conduct(session, 'GET', prefix + 'tags', auth=auth).json()
    assert len(tags_result.values()) == len(tag_names)

    tag_image_id = tags_result['latest']
    if not options.munge_shas:
      # Ensure we have a matching image ID.
      known_ids = {image.id for image in images}
      assert tag_image_id in known_ids

    # Retrieve the ancestry of the tagged image.
    image_prefix = '/v1/images/%s/' % tag_image_id
    ancestors = self.conduct(session, 'GET', image_prefix + 'ancestry', auth=auth).json()

    assert len(ancestors) == len(images)
    for index, image_id in enumerate(reversed(ancestors)):
      # /v1/images/{imageID}/{ancestry, json, layer}
      image_prefix = '/v1/images/%s/' % image_id
      self.conduct(session, 'GET', image_prefix + 'ancestry', auth=auth)

      result = self.conduct(session, 'GET', image_prefix + 'json', auth=auth)
      assert result.json()['id'] == image_id

      # Ensure we can HEAD the image layer.
      self.conduct(session, 'HEAD', image_prefix + 'layer', auth=auth)

      # And retrieve the layer data.
      result = self.conduct(session, 'GET', image_prefix + 'layer', auth=auth)
      assert result.content == images[index].bytes

    return PullResult(manifests=None)

  def push(self, session, namespace, repo_name, tag_names, images, credentials=None,
           expected_failure=None, options=None):
    auth = self._auth_for_credentials(credentials)
    tag_names = [tag_names] if isinstance(tag_names, str) else tag_names

    # Ping!
    self.ping(session)

    # PUT /v1/repositories/{namespace}/{repository}/
    result = self.conduct(session, 'PUT', '/v1/repositories/%s/%s/' % (namespace, repo_name),
                          expected_status=(201, expected_failure, V1ProtocolSteps.PUT_IMAGES),
                          json_data={},
                          auth=auth)

    if expected_failure is not None:
      return

    headers = {}
    headers['Authorization'] = 'token ' + result.headers['www-authenticate']

    for image in images:
      # PUT /v1/images/{imageID}/json
      image_json_data = {'id': image.id}
      if image.size is not None:
        image_json_data['Size'] = image.size

      if image.parent_id is not None:
        image_json_data['parent'] = image.parent_id

      self.conduct(session, 'PUT', '/v1/images/%s/json' % image.id,
                   json_data=image_json_data, headers=headers)

      # PUT /v1/images/{imageID}/layer
      self.conduct(session, 'PUT', '/v1/images/%s/layer' % image.id,
                   data=StringIO(image.bytes), headers=headers)

      # PUT /v1/images/{imageID}/checksum
      checksum = compute_simple(StringIO(image.bytes), json.dumps(image_json_data))
      checksum_headers = {'X-Docker-Checksum-Payload': checksum}
      checksum_headers.update(headers)

      self.conduct(session, 'PUT', '/v1/images/%s/checksum' % image.id,
                   headers=checksum_headers)

    # PUT /v1/repositories/{namespace}/{repository}/tags/latest
    for tag_name in tag_names:
      self.conduct(session, 'PUT',
                   '/v1/repositories/%s/%s/tags/%s' % (namespace, repo_name, tag_name),
                   data='"%s"' % images[-1].id,
                   headers=headers)

    # PUT /v1/repositories/{namespace}/{repository}/images
    self.conduct(session, 'PUT', '/v1/repositories/%s/%s/images' % (namespace, repo_name),
                 expected_status=204, headers=headers)

    return PushResult(checksums=None, manifests=None)