import docker
import logging
import shutil
import os
import re
import requests
import json

from flask import Flask, jsonify, abort, make_response
from zipfile import ZipFile
from tempfile import TemporaryFile, mkdtemp
from multiprocessing.pool import ThreadPool
from base64 import b64encode


BUFFER_SIZE = 8 * 1024
LOG_FORMAT = '%(asctime)-15s - %(levelname)s - %(pathname)s - ' + \
             '%(funcName)s - %(message)s'

app = Flask(__name__)
logger = logging.getLogger(__name__)


def count_steps(dockerfile_path):
  with open(dockerfile_path, 'r') as dockerfileobj:
    steps = 0
    for line in dockerfileobj.readlines():
      stripped = line.strip()
      if stripped and stripped[0] is not '#':
        steps += 1
    return steps


def prepare_zip(request_file):
  build_dir = mkdtemp(prefix='docker-build-')

  # Save the zip file to temp somewhere
  with TemporaryFile() as zip_file:
    zip_file.write(request_file.content)
    to_extract = ZipFile(zip_file)
    to_extract.extractall(build_dir)

  return build_dir


def prepare_dockerfile(request_file):
  build_dir = mkdtemp(prefix='docker-build-')
  dockerfile_path = os.path.join(build_dir, "Dockerfile")
  with open(dockerfile_path, 'w') as dockerfile:
    dockerfile.write(request_file.content)

  return build_dir


def total_completion(statuses, total_images):
  percentage_with_sizes = float(len(statuses.values()))/total_images
  sent_bytes = sum([status[u'current'] for status in statuses.values()])
  total_bytes = sum([status[u'total'] for status in statuses.values()])
  return float(sent_bytes)/total_bytes*percentage_with_sizes


def build_image(build_dir, tag_name, num_steps, result_object):
  try:
    logger.debug('Starting build.')
    docker_cl = docker.Client(timeout=1200)
    result_object['status'] = 'building'
    build_status = docker_cl.build(path=build_dir, tag=tag_name, stream=True)

    current_step = 0
    built_image = None
    for status in build_status:
      # logger.debug('Status: %s', str(status))
      step_increment = re.search(r'Step ([0-9]+) :', status)
      if step_increment:
        current_step = int(step_increment.group(1))
        logger.debug('Step now: %s/%s' % (current_step, num_steps))
        result_object['current_command'] = current_step
        continue

      complete = re.match(r'Successfully built ([a-z0-9]+)$', status)
      if complete:
        built_image = complete.group(1)
        logger.debug('Final image ID is: %s' % built_image)
        continue

    shutil.rmtree(build_dir)

    # Get the image count
    if not built_image:
      result_object['status'] = 'error'
      result_object['message'] = 'Unable to build dockerfile.'
      return

    history = json.loads(docker_cl.history(built_image))
    num_images = len(history)
    result_object['total_images'] = num_images

    result_object['status'] = 'pushing'
    logger.debug('Pushing to tag name: %s' % tag_name)
    resp = docker_cl.push(tag_name, stream=True)

    for status_str in resp:
      status = json.loads(status_str)
      logger.debug('Status: %s', status_str)
      if u'status' in status:
        status_msg = status[u'status']

        if status_msg == 'Pushing':
          if u'progressDetail' in status and u'id' in status:
            image_id = status[u'id']
            detail = status[u'progressDetail']

            if u'current' in detail and 'total' in detail:
              images = result_object['image_completion']

              images[image_id] = detail
              result_object['push_completion'] = total_completion(images,
                                                                  num_images)

      elif u'errorDetail' in status:
        result_object['status'] = 'error'
        if u'message' in status[u'errorDetail']:
          result_object['message'] = str(status[u'errorDetail'][u'message'])
        return

    result_object['status'] = 'complete'
  except Exception as e:
    logger.exception('Exception when processing request.')
    result_object['status'] = 'error'
    result_object['message'] = str(e.message)


MIME_PROCESSORS = {
  'application/zip': prepare_zip,
  'text/plain': prepare_dockerfile,
  'application/octet-stream': prepare_dockerfile,
}

# If this format it should also be changed in the api method get_repo_builds
build = {
  'total_commands': None,
  'current_command': None,
  'push_completion': 0.0,
  'status': 'waiting',
  'message': None,
  'image_completion': {},
}
pool = ThreadPool(1)


@app.before_first_request
def start_build():
  resource_url = os.environ['RESOURCE_URL']
  tag_name = os.environ['TAG']
  acccess_token = os.environ['TOKEN']

  logger.debug('Starting job with resource url: %s tag: %s and token: %s' %
               (resource_url, tag_name, acccess_token))

  # Save the token
  host = re.match(r'([a-z0-9.:]+)/.+/.+$', tag_name)
  if host:
    docker_endpoint = 'http://%s/v1/' % host.group(1)
    dockercfg_path = os.path.join(os.environ.get('HOME', '.'), '.dockercfg')
    token = b64encode('$token:%s' % acccess_token)
    with open(dockercfg_path, 'w') as dockercfg:
      payload = {
        docker_endpoint: {
          'auth': token,
          'email': '',
        }
      }
      dockercfg.write(json.dumps(payload))

  else:
    raise Exception('Invalid tag name: %s' % tag_name)

  docker_resource = requests.get(resource_url)
  c_type = docker_resource.headers['content-type']

  logger.info('Request to build file of type: %s with tag: %s' %
              (c_type, tag_name))

  if c_type not in MIME_PROCESSORS:
    raise Exception('Invalid dockerfile content type: %s' % c_type)

  build_dir = MIME_PROCESSORS[c_type](docker_resource)

  dockerfile_path = os.path.join(build_dir, "Dockerfile")
  num_steps = count_steps(dockerfile_path)
  logger.debug('Dockerfile had %s steps' % num_steps)

  logger.info('Sending job to builder pool.')
  build['total_commands'] = num_steps

  pool.apply_async(build_image, [build_dir, tag_name, num_steps,
                                 build])


@app.route('/build/', methods=['GET'])
def get_status():
  if build:
    return jsonify(build)
  abort(404)


@app.route('/status/', methods=['GET'])
def health_check():
  return make_response('Running')


if __name__ == '__main__':
  logging.basicConfig(level=logging.DEBUG, format=LOG_FORMAT)
  app.run(host='0.0.0.0', port=5002, threaded=True)