#!/usr/bin/env python
"""

Before running this please run
CUDA_VISIBLE_DEVICES=0 th web.lua
CUDA_VISIBLE_DEVICES=0 /home/jerryli27/Programs/serving/bazel-bin/tensorflow_serving/model_servers/tensorflow_model_server --port=9000 --model_name=test --model_base_path=/bla/bla/bla
"""

from __future__ import absolute_import
from __future__ import unicode_literals
import CGIHTTPServer, SimpleHTTPServer, BaseHTTPServer
import SocketServer

import functools
import shutil
import os, sys
import base64
import json
import io
import time
from PIL import Image

import argparse

from cgi import parse_header, parse_multipart
from urlparse import parse_qs
from io import open
import scipy.misc
import numpy as np
import tensorflow as tf

from interface import face_detection_util
from interface import domain_translation_client
from interface import waifu2x_interface
from interface import labeler_client
from interface import sketch_refinement_client
from interface import interface_utils
import util_io
import util_misc

tf.app.flags.DEFINE_integer('gpu', -1,
                            'GPU ID (negative value indicates CPU)')
tf.app.flags.DEFINE_integer('port', 8000, 'Port of this server.')
tf.app.flags.DEFINE_string('host', '', 'Host of this server.')
tf.app.flags.DEFINE_integer('max_num_faces', 4, 'Maximum number of faces to do transfer.')


tf.app.flags.DEFINE_string('sketch_refinement_supervised_server', None, 'PredictionService host:port')
# tf.app.flags.mark_flag_as_required('sketch_refinement_supervised_server')
tf.app.flags.DEFINE_string('sketch_refinement_unsupervised_server', None, 'PredictionService host:port')
# tf.app.flags.mark_flag_as_required('sketch_refinement_unsupervised_server')

tf.app.flags.DEFINE_string('labeler_source_image_path', None,'')
tf.app.flags.DEFINE_string('labeler_output_image_path', None,'')
tf.app.flags.DEFINE_string('labeler_sketch_folder', None,'')
tf.app.flags.DEFINE_string('labeler_sketch_refinement_folder', None,'')
FLAGS = tf.app.flags.FLAGS

FACE_DETECTOR = face_detection_util.FaceDetector()
DT_CLIENT = domain_translation_client.DomainTranslationClient('', None)
WAIFU2X = waifu2x_interface.Waifu2x()
LABELER_CLIENT = labeler_client.LabelerClient()
SKETCH_REFINEMENT_CLIENT = sketch_refinement_client.SketchRefinementClient('', None, supervised=True)
SKETCH_REFINEMENT_CLIENT_UNSUPERVISED = sketch_refinement_client.SketchRefinementClient('', None, supervised=False)

class MyHandler(CGIHTTPServer.CGIHTTPRequestHandler):
  def __init__(self, req, client_addr, server):
    CGIHTTPServer.CGIHTTPRequestHandler.__init__(self, req, client_addr, server)


  def parse_POST(self):
    ctype, pdict = parse_header(self.headers['content-type'])
    pdict['boundary'] = str(pdict['boundary']).encode('utf-8')
    if ctype == 'multipart/form-data':
      postvars = parse_multipart(self.rfile, pdict)
    elif ctype == 'application/x-www-form-urlencoded':
      length = int(self.headers['content-length'])
      postvars = parse_qs(
        self.rfile.read(length),
        keep_blank_values=1)
    else:
      postvars = {}
    return postvars

  def do_POST(self):
    form = self.parse_POST()

    post_type = None
    if 'type' in form:
      post_type = form['type'][0]

    if post_type == 'labeler':
      if 'id' in form:
        id_str = form['id'][0]
        id_str = id_str.decode()

        skipped = (form['skip'][0] == 'true')
        if not skipped:
          bin1 = form['image'][0]
          bin1 = bin1.decode().split(',')[1]
          bin1 = base64.b64decode(bin1.encode())
          image_path = os.path.join(FLAGS.labeler_output_image_path,  os.path.splitext(id_str)[0] + '.png')  # id_str happens to be the file name.
          with open(image_path, 'wb') as fout1:
            fout1.write(bin1)
        LABELER_CLIENT.mark_current_as_done(skipped)

      image, sketch, image_id = LABELER_CLIENT.get_image_and_id()
      if image:
        self.post_success(image_id, {'image': image, 'sketch': sketch})
      else:
        self.post_success(image_id, {'error': 'Ran out of images.'})
      return
    elif post_type == 'sketch_refinement':
      id_str = form['id'][0]
      id_str = id_str.decode()
      if 'sketch_refinement' not in form or 'sketch' not in form:
        self.post_bad_request('must contain sketch_refinement and sketch', id_str)
        return
      x_pos = int(form['x'][0])
      y_pos = int(form['y'][0])
      # We only need the alpha channel. That will give us a grayscale (0/1 in this case) image.
      sketch_refinement_img = interface_utils.base64_to_numpy(form['sketch_refinement'][0], contains_format=True)[...,3]
      sketch_img = interface_utils.base64_to_numpy(form['sketch'][0], contains_format=True)
      sketch_img = util_misc.im2gray(sketch_img)
      # with open(image_path, 'wb') as fout1:
      #   fout1.write(sketch_refinement)
      out_path = os.path.join(FLAGS.labeler_sketch_refinement_folder, os.path.splitext(id_str)[0] + '_' + str(time.time()) + '.png')
      SKETCH_REFINEMENT_CLIENT.do_inference(out_path, center_point_xy=[x_pos, y_pos], sketch_image_np=sketch_img, image_np=sketch_refinement_img)
      sketch_refinement_exists = SKETCH_REFINEMENT_CLIENT.block_on_callback(out_path)
      if not sketch_refinement_exists:
        self.post_server_internal_error('sketch refinement timed out.', id_str)
        return
      LABELER_CLIENT.set_current_sketch_path(out_path)
      self.post_success(id_str, {'refined_sketch': interface_utils.get_image_encoding(out_path)})
      return


    if 'id' in form:
      id_str = form['id'][0]
      id_str = id_str.decode()
    else:
      id_str = 'test'

    if 'register_download' in form and form['register_download']:
      if 'subid' in form:
        subid_str = form['subid'][0]
        subid_str = subid_str.decode()
      else:
        subid_str = '0'
      print('TODO: do action for register download for id: %s and subid: %s' %(id_str, subid_str))
      self.post_success(id_str,)
      return

    elif 'image' in form:
      bin1 = form['image'][0]
      input_image_path = interface_utils.save_encoded_image(bin1, './static/images/inputs/' + id_str)

      cropped_image_pattern = './static/images/cropped_faces/' + id_str + '_%d.png'
      faces = FACE_DETECTOR.crop_face_and_save(input_image_path, cropped_image_pattern)
      num_faces = len(faces)
      if num_faces > FLAGS.max_num_faces:
        faces = faces[FLAGS.max_num_faces]
        num_faces = FLAGS.max_num_faces
      if num_faces == 0:
        shutil.copy(input_image_path, './static/images/cropped_faces/' + id_str + '_%d.png' %0)
        faces = [util_io.imread(input_image_path, (FLAGS.image_hw, FLAGS.image_hw))]
        num_faces = len(faces)

      transferred_image_file_format = './static/images/transferred_faces/' + id_str + '_%d.png'
      succeed, transferred_image_files = self.automatic_retry(functools.partial(
        self.domain_transfer, transferred_image_file_format=transferred_image_file_format, images=faces))
      if not succeed:
        self.post_server_internal_error('Domain transfer failed', id_str, {'num_faces': num_faces})
        return

      if 'do_waifu2x' in form:
        do_waifu2x = form['do_waifu2x'][0] == 'true'
      else:
        do_waifu2x = False
      if do_waifu2x:
        transferred_2x_image_file_format = './static/images/transferred_faces_2x/' + id_str + '_%d.png'
        succeed, transferred_2x_image_files = self.automatic_retry(functools.partial(
          self.call_waifu2x, transferred_image_file_format=transferred_image_file_format,
          transferred_2x_image_file_format=transferred_2x_image_file_format, num_images=num_faces))
        if not succeed:
          self.post_server_internal_error('Waifu2x failed', id_str, {'num_faces': num_faces})
          return
        transferred_image_to_be_combined_format = transferred_2x_image_file_format
      else:
        transferred_image_to_be_combined_format = transferred_image_file_format

      combined_image_pattern = './static/images/combined/' + id_str + '_%d.png'
      succeed, combined_images = self.automatic_retry(functools.partial(
        self.combine_original_and_transferred, images=faces, combined_image_pattern=combined_image_pattern,
        transferred_2x_image_file_format=transferred_image_to_be_combined_format))
      if not succeed:
        self.post_server_internal_error('Combine original and transferred failed.', id_str, {'num_faces': num_faces})
        return

      self.post_success(id_str, {'num_faces': num_faces})
    else:
      self.post_bad_request('Post request must contain image.', id_str)
    return

  def post_success(self, id_str, extra_kwargs=None):
    content = {
      'message': 'The command Completed Successfully',
      'Status': '200 OK',
      'success': True,
      'gpu': FLAGS.gpu,
      'id_str': id_str,
    }
    if extra_kwargs:
      content.update(extra_kwargs)
    content = json.dumps(content)
    self.send_response(200)
    self.send_header('Content-type', 'application/json')
    self.send_header('Content-Length', len(content))
    self.end_headers()
    self.wfile.write(content)

  def post_bad_request(self, message, id_str, extra_kwargs=None):
    content = {
      'message':message,
      'Status':'400 Bad Request',
      'success':False,
      'gpu': FLAGS.gpu,
      'id_str': id_str,
    }
    if extra_kwargs:
      content.update(extra_kwargs)
    content = json.dumps(content)
    self.send_response(400)
    self.send_header('Content-type', 'application/json')
    self.send_header('Content-Length', len(content))
    self.end_headers()
    self.wfile.write(content)

  def post_server_internal_error(self, error_message, id_str, extra_kwargs=None):
    content = {
      'message': 'Some internal server error occurred: Error is "%s". Please try again.' %error_message,
      'Status': '500 Internal Server Error',
      'success': False,
      'gpu': FLAGS.gpu,
      'id_str': id_str,
    }
    if extra_kwargs:
      content.update(extra_kwargs)
    content = json.dumps(content)
    self.send_response(500)
    self.send_header('Content-type', 'application/json')
    self.send_header('Content-Length', len(content))
    self.end_headers()
    self.wfile.write(content)

  def automatic_retry(self, func, num_tries=3):
    """Automatically reruns a function.

    Takes a function that returns a list of files as output. Check whether those files exists.
    If not, run the function again."""
    succeed = False
    while num_tries > 0 and not succeed:
      paths = func()
      succeed = True
      for path in paths:
        if not os.path.exists(path):
          succeed = False
          break
    return succeed, paths

  @staticmethod
  def domain_transfer(transferred_image_file_format, images):
    ret = []
    for i, image in enumerate(images):
      image_path = transferred_image_file_format % i
      ret.append(image_path)
      if os.path.exists(image_path):
        continue
      DT_CLIENT.do_inference(image_path, image_np=image)
      DT_CLIENT.block_on_callback(image_path)
    return ret

  @staticmethod
  def call_waifu2x(transferred_image_file_format, transferred_2x_image_file_format, num_images):
    ret = []
    for i in range(num_images):
      image_path = transferred_2x_image_file_format % i
      ret.append(image_path)
      if os.path.exists(image_path):
        continue
      WAIFU2X.post_request(transferred_image_file_format % i, image_path)
    return ret

  @staticmethod
  def combine_original_and_transferred(images, transferred_2x_image_file_format, combined_image_pattern):
    ret = []
    for i in range(len(images)):
      save_image_path = combined_image_pattern % i
      ret.append(save_image_path)
      if os.path.exists(save_image_path):
        continue
      # raise NotImplementedError('deal with error: IOError: cannot identify image file  u\'./static/images/transferred_faces/')
      transferred_image = util_io.imread(transferred_2x_image_file_format % i, )
      face_image = scipy.misc.imresize(images[i], (transferred_image.shape[0], transferred_image.shape[1]))
      combined_image = np.concatenate((face_image, transferred_image), axis=1)
      util_io.imsave(save_image_path, combined_image)
    return ret


def main(_):
  global DT_CLIENT
  DT_CLIENT = domain_translation_client.DomainTranslationClient(
    FLAGS.domain_translation_server, FLAGS.image_hw,
    # sample_images_path='/home/jerryli27/PycharmProjects/image2tag/data/inference_test_data/images_sample.npy',
    concurrency=FLAGS.concurrency)
  if FLAGS.labeler_source_image_path:
    LABELER_CLIENT.set_image_paths(FLAGS.labeler_source_image_path, os.path.join(FLAGS.labeler_output_image_path, 'finished_images.txt'), FLAGS.labeler_sketch_folder)
  else:
    tf.logging.warning('labeler client not set! empty flag `labeler_source_image_path`.')

  global SKETCH_REFINEMENT_CLIENT
  if FLAGS.sketch_refinement_supervised_server:
    SKETCH_REFINEMENT_CLIENT = sketch_refinement_client.SketchRefinementClient(
      FLAGS.sketch_refinement_supervised_server, FLAGS.image_hw, concurrency=FLAGS.concurrency, supervised=True
    )
    # SKETCH_REFINEMENT_CLIENT_UNSUPERVISED = sketch_refinement_client.SketchRefinementClient(
    #   FLAGS.sketch_refinement_unsupervised_server, FLAGS.image_hw, concurrency=FLAGS.concurrency, supervised=False
    # )

  if FLAGS.labeler_sketch_refinement_folder:
    util_io.touch_folder(FLAGS.labeler_sketch_refinement_folder)

  print 'GPU: {}'.format(FLAGS.gpu)
  httpd = BaseHTTPServer.HTTPServer((FLAGS.host, FLAGS.port), MyHandler)
  print 'serving at', FLAGS.host, ':', FLAGS.port
  httpd.serve_forever()


if __name__ == '__main__':
  tf.app.run()

# eg

# saved_model_cli show --dir '/home/xor/PycharmProjects/image2tag/checkpoints/pggan_sharew_sc/export/2'  --tag_set serve  --signature_def serving_default

# saved_model_cli run --dir '/home/xor/PycharmProjects/image2tag/checkpoints/pggan_sharew_sc/export/2'  --tag_set serve  --signature_def serving_default --input_exprs 'inputs=np.ones((1,4,4,3))'