import re
import base64
import datetime
import sys
import pickle
import gzip
import json
import time
import gradio as gr

from ai import *
from tricks import *
from decompositioner import *
from rendering import *


import os
import glob
import shutil

splash = glob.glob('ui/web-mobile/splash*')[0]
os.remove(splash)
shutil.copy('res/splash.png', splash)
with open('ui/web-mobile/index.html', 'r', encoding='utf-8') as f:
    page = f.read()
with open('ui/web-mobile/index.html', 'w', encoding='utf-8') as f:
    f.write(page.replace('Cocos Creator | ', ''))


def cv2_encode(image: np.ndarray, name):
    if image is None:
        return 'null'
    if '.jpg' in name:
        _, data = cv2.imencode('.jpeg', image)
        return 'data:image/jpeg;base64,' + base64.b64encode(data).decode('utf8')
    else:
        _, data = cv2.imencode('.png', image)
        return 'data:image/png;base64,' + base64.b64encode(data).decode('utf8')


def get_request_image(request, name):
    img = request.get(name)
    img = re.sub('^data:image/.+;base64,', '', img)
    img = base64.b64decode(img)
    img = np.fromstring(img, dtype=np.uint8)
    img = cv2.imdecode(img, -1)
    return img


def npcache(history, room_id, name, data):
    rooms = list(filter(lambda _room: _room["id"] == room_id, history))
    if len(rooms) == 0:
        room = { "id": room_id }
        history.append(room)
    else:
        room = rooms[0]
    room[name] = data


def npread(history, room_id, name):
    rooms = list(filter(lambda _room: _room["id"] == room_id, history))
    if len(rooms) == 0:
        return None
    else:
        room = rooms[0]
    return room.get(name, None)


def upload_sketch(json_str, history):
    request = json.loads(json_str)
    timenow = time.time()
    ID = datetime.datetime.now().strftime('H%HM%MS%S')
    room = datetime.datetime.now().strftime('%b%dH%HM%MS%S') + 'R' + str(np.random.randint(100, 999))
    sketch = from_png_to_jpg(get_request_image(request, 'sketch'))
    npcache(history, room, 'sketch.original.jpg', sketch)
    sketch = go_tail(cli_norm(min_resize(sketch, 512)))
    print('original_sketch saved')
    s256 = go_vector(go_cal(mk_resize(sketch, 8)))[:, :, 0]
    print('s256')
    s512 = go_vector(go_cal(d_resize(sketch, s256.shape, 2.0)))[:, :, 0]
    print('s512')
    s1024 = go_vector(go_cal(d_resize(sketch, s256.shape, 4.0)))[:, :, 0]
    print('s1024')
    npcache(history, room, 'sketch.s1024.png', s1024)
    npcache(history, room, 'sketch.s512.png', s512)
    npcache(history, room, 'sketch.s256.png', s256)
    print('edge processed')
    fill = double_fill(s1024, s512, s256)
    npcache(history, room, 'sketch.fill', fill)
    print('filled')
    npcache(history, room, 'sketch.colorization.png', np.min(sketch, axis=2))
    print('sketch processed')
    print(time.time() - timenow)
    if len(history) > 5:
        history = history[-5:]
    return room + '_' + ID, history


def request_result(json_str, history):
    request = json.loads(json_str)
    timenow = time.time()
    room = request.get("room")
    if len(list(filter(lambda _room: _room["id"] == room, history))) == 0:
        return None, history
    skipper = str(request.get("skipper"))
    light_r = float(request.get("r"))
    light_g = float(request.get("g"))
    light_b = float(request.get("b"))
    light_h = float(request.get("h"))
    light_max = max([light_r, light_g, light_b, light_h])
    inv4 = int(request.get("inv4"))
    print('inv4=' + str(inv4))
    light_r = (light_r + 1e-5) / (light_max + 1e-5)
    light_g = (light_g + 1e-5) / (light_max + 1e-5)
    light_b = (light_b + 1e-5) / (light_max + 1e-5)
    light_h *= 600.0
    light_d = request.get("d")
    need_render = int(request.get("need_render"))
    print([light_r, light_g, light_b, light_d])
    ID = datetime.datetime.now().strftime('H%HM%MS%S')
    points = request.get("points")
    points = json.loads(points)
    npcache(history, room, 'points.' + ID + '.txt', points)
    if len(points) > 500:
        return None, history
    for _ in range(len(points)):
        points[_][1] = 1 - points[_][1]
    fill = npread(history, room, 'sketch.fill')
    s1024 = npread(history, room, 'sketch.s1024.png')
    sketch = npread(history, room, 'sketch.colorization.png')
    print(skipper)
    if npread(history, room, 'albedo.' + skipper + '.png') is not None:
        albedo = npread(history, room, 'albedo.' + skipper + '.png')
        npcache(history, room, 'albedo.' + ID + '.png', albedo)
        print('albedo readed')
    else:
        faceID = int(request.get("faceID")) - 65535
        print(faceID)
        if faceID > -1:
            print('fake ref')
            face = from_png_to_jpg(cv2.imread("refs/" + str(faceID + 1) + ".png", cv2.IMREAD_UNCHANGED))
        else:
            print('get ref')
            face = from_png_to_jpg(get_request_image(request, 'face'))
        npcache(history, room, 'face.' + ID + '.jpg', face)
        face = s_enhance(face)
        print('request result room = ' + str(room) + ', ID = ' + str(ID))
        print('processing painting in ' + room)
        if inv4 > 0:
            sketch_1024 = k_resize(sketch, 64)
        else:
            sketch_1024 = k_resize(sketch, 48)
        hints_1024 = ini_hint(sketch_1024)
        hints_1024 = opreate_normal_hint(hints_1024, points, length=2, skip_sp=True)
        baby = go_head(
            sketch=sketch_1024,
            global_hint=k_resize(face, 14),
            local_hint=hints_1024
        )
        npcache(history, room, 'baby.' + ID + '.jpg', baby)
        print('baby born')
        composition = d_resize(re_deatlize(deatlize(balance_fill(baby, fill, opreate_normal_hint(ini_hint(s1024), points, length=2, skip_sp=True), s1024)), s1024), sketch.shape)
        npcache(history, room, 'composition.' + ID + '.jpg', composition)
        gird = process_overlay(composition, sketch)
        npcache(history, room, 'gird.' + ID + '.jpg', gird)
        print('composition saved')
        if inv4 > 0:
            albedo = go_render(sketch_1024, d_resize(composition, sketch_1024.shape, 0.5), hints_1024)
            albedo = go_tail(albedo)
            albedo = d_resize(re_deatlize(d_resize(albedo, s1024.shape), s1024), sketch.shape)
            albedo = cv2.cvtColor(albedo, cv2.COLOR_RGB2YUV)
            albedo[:, :, 0] = go_vgg7(albedo[:, :, 0])
            albedo = cv2.cvtColor(albedo, cv2.COLOR_YUV2RGB)
        else:
            albedo = re_deatlize(d_resize(baby, s1024.shape), s1024)
            albedo = d_resize(albedo, sketch.shape, 0.25)
            albedo = go_tail(albedo)
            albedo = go_tail(albedo)
            albedo = d_resize(albedo, sketch.shape)
            boundary = sketch.astype(np.float32)
            boundary = cv2.GaussianBlur(boundary, (0, 0), 1.618) - boundary
            boundary = boundary.clip(0, 255)
            albedo = cv2.cvtColor(albedo, cv2.COLOR_RGB2HSV).astype(np.float32)
            albedo[:, :, 1] += albedo[:, :, 1] * boundary / 48.0
            albedo[:, :, 2] -= boundary
            albedo = cv2.cvtColor(albedo.clip(0, 255).astype(np.uint8), cv2.COLOR_HSV2RGB)
        npcache(history, room, 'albedo.' + ID + '.png', albedo)
        print('albedo saved')
        if need_render == 0:
            npcache(history, room, 'result.' + ID + '.jpg', albedo)
            # cv2.imwrite('results/' + room + '.' + ID + '.jpg', albedo)
            print(time.time() - timenow)
            return room + '_' + ID, history
        HSV, YUV, DEL = process_albedo(albedo, composition, sketch)
        npcache(history, room, 'HSV.' + ID + '.jpg', HSV)
        npcache(history, room, 'YUV.' + ID + '.jpg', YUV)
        npcache(history, room, 'DEL.' + ID + '.jpg', DEL)
        print('HSV YUV DEL')
    albedo_s1024 = d_resize(albedo, s1024.shape)
    matting = go_mat(albedo_s1024)
    matting = np.tile(matting[:, :, None], [1, 1, 3])
    matting = shade_fill(matting, fill, opreate_normal_hint(ini_hint(s1024), points, length=2, skip_sp=False), s1024)
    matting = matting[:, :, 0]
    depth = np.zeros_like(matting, dtype=np.uint8) + 255
    depth[matting < 127] = 127
    depth[s1024 < 250] = 0
    npcache(history, room, 'depth.' + ID + '.jpg', depth)
    print('depth saved')
    normal = go_norm(depth).astype(np.float32)
    normal = ((normal + 1e-4) / (np.max(normal, axis=2, keepdims=True) + 1e-4) * 255.0).clip(0, 255).astype(np.uint8)
    normal[matting < 127] = 255
    normal = re_deatlize(normal, s1024)
    normal = d_resize(normal, sketch.shape)
    npcache(history, room, 'normal.' + ID + '.jpg', normal)
    print('norm saved')
    mask = np.zeros_like(matting, dtype=np.uint8) + 255
    mask[matting < 127] = 0
    mask = d_resize(mask, sketch.shape)
    mask[mask < 127] = 0
    mask[mask > 0] = 255
    if int(light_d) == 0:
        result = small_render(normal, mask, albedo, s1024, r=light_r, g=light_g, b=light_b, h=light_h, left=True, top=True)
    elif int(light_d) == 1:
        result = small_render(normal, mask, albedo, s1024, r=light_r, g=light_g, b=light_b, h=light_h, left=False, top=True)
    elif int(light_d) == 2:
        result = small_render(normal, mask, albedo, s1024, r=light_r, g=light_g, b=light_b, h=light_h, left=True, top=False)
    else:
        result = small_render(normal, mask, albedo, s1024, r=light_r, g=light_g, b=light_b, h=light_h, left=False, top=False)
    if need_render == 2:
        npcache(history, room, 'result.' + ID + '.jpg', result)
        # cv2.imwrite('results/' + room + '.' + ID + '.jpg', result)
        print(time.time() - timenow)
        return room + '_' + ID, history
    print('result saved')
    preview = np.concatenate([np.tile(sketch[:, :, None], [1, 1, 3]), albedo, result], axis=1)
    npcache(history, room, 'preview.' + ID + '.jpg', preview)
    print('preview saved')
    npcache(history, room, 'result.' + ID + '.jpg', result)
    # cv2.imwrite('results/' + room + '.' + ID + '.jpg', preview)
    print(time.time() - timenow)
    return room + '_' + ID, history


def download_result(name, history):
    room_id, name = name.split('/')
    rooms = list(filter(lambda _room: _room["id"] == room_id, history))
    if len(rooms) == 0:
        return None
    else:
        room = rooms[0]

    real_name = None
    for k in room.keys():
        if name in k:
            real_name = k
            break
    if real_name is None:
        return None
    name = real_name
    
    result = room.get(name, None)

    if 'points' in name:
        return json.dumps(result)
    
    return cv2_encode(result, name)
    

with gr.Blocks() as demo:
    history = gr.State(value=[])
    with gr.Row():
        with gr.Column():
            btn_show = gr.Button("Open Style2Paints V4.2")
    btn_show.click(None, _js="(_) => open('file/ui/web-mobile/index.html')")

    with gr.Row():
        with gr.Box():
            with gr.Row():
                upload_sketch_json = gr.Textbox(label="upload_sketch(json string)")
            with gr.Row():
                upload_sketch_btn = gr.Button(label="Submit sketch json")
            with gr.Row():
                upload_sketch_result = gr.Textbox(label="Result", interactive=False)
            upload_sketch_btn.click(upload_sketch, [upload_sketch_json, history], [upload_sketch_result, history], api_name="upload_sketch")
        
        with gr.Box():
            with gr.Row():
                request_result_json = gr.Textbox(label="request_result(json string)")
            with gr.Row():
                request_result_btn = gr.Button(label="Submit json of request for result")
            with gr.Row():
                request_result_result = gr.Textbox(label="Result", interactive=False)
            upload_sketch_btn.click(request_result, [request_result_json, history], [request_result_result, history], api_name="request_result")
        
        with gr.Box():
            with gr.Row():
                download_result_json = gr.Textbox(label="download_result(json string)")
            with gr.Row():
                download_result_btn = gr.Button(label="Submit json of download for result")
            with gr.Row():
                download_result_result = gr.Textbox(label="Result", interactive=False)
            upload_sketch_btn.click(download_result, [download_result_json, history], [download_result_result], api_name="download_result")

demo.launch()