import os
import cv2
import bpy
import json
import socket
import numpy as np
from math import radians
from base64 import b64encode
from bpy.types import Operator
from mathutils import Vector, Quaternion


def encodeImage(image):
    # encode the image
    image_encode = cv2.imencode('.jpg', image, [cv2.IMWRITE_JPEG_QUALITY, bpy.context.scene.quality])[1]
    image_bytes = np.array(image_encode).tobytes()
    image_str = b64encode(image_bytes).decode('utf-8')
    return image_str


def packDataFlow(data):
    length_flow = str(len(data))
    length_flow = length_flow.zfill(8)
    return length_flow.encode('utf-8') + data


def drivenCharacter(rots, trans):
    try:
        scene = bpy.context.scene
        character = bpy.data.objects[scene.character_name]
        bones = character.pose.bones
        bone_names = ['Pelvis', 'L_Hip', 'R_Hip', 'Spine1', 'L_Knee', 'R_Knee', 'Spine2', 'L_Ankle', 'R_Ankle',
                      'Spine3', 'L_Foot', 'R_Foot', 'Neck', 'L_Collar', 'R_Collar', 'Head', 'L_Shoulder', 'R_Shoulder',
                      'L_Elbow', 'R_Elbow', 'L_Wrist', 'R_Wrist']

        rots = np.array(rots)
        trans = np.array(trans)

        scene.frame_current += scene.insert_interval
        if scene.translation:
            bones['Pelvis'].location = Vector((100 * trans[0], -100 * trans[1], -100 * trans[2]))
        else:
            bones['Pelvis'].location = Vector((0, 0, 0))

        if scene.insert_keyframe:
            bones['Pelvis'].keyframe_insert('location')

        for idx in range(len(bone_names)):
            bone = bones[bone_names[idx]]

            bone_rot = Quaternion(
                Vector((rots[4 * idx], rots[4 * idx + 1], rots[4 * idx + 2])),
                rots[4 * idx + 3])
            quat = Quaternion((1.0, 0.0, 0.0), radians(-180))

            if bone_names[idx] == 'Pelvis':
                bone.rotation_quaternion = (quat) @ bone_rot
            else:
                bone.rotation_quaternion = bone_rot
            if scene.insert_keyframe:
                bone.keyframe_insert('rotation_quaternion', frame=scene.frame_current)

        scene.frame_end = scene.frame_current
        return True
    except:
        return False


class PlayAnimation(Operator):
    bl_idname = 'ops.webcam_animation'
    bl_label = 'PLAY'

    def execute(self, ctx):
        self.client = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.client.connect((ctx.scene.ip, int(ctx.scene.port)))
        ctx.scene.frame_start = ctx.scene.frame_current
        # self.webcam = cv2.VideoCapture(0)
        # self.webcam.set(3, ctx.scene.width)
        # self.webcam.set(4, ctx.scene.height)

        # init = json.dumps(
        #     {
        #         'type': 'init',
        #         'gpu': str(ctx.scene.gpu),
        #         'mode': 'webcam',  ## modify
        #     }
        # ).encode('utf-8')
        # self.client.send(packDataFlow(init))

        self._timer = ctx.window_manager.event_timer_add(1 / ctx.scene.fps, window=ctx.window)
        ctx.window_manager.modal_handler_add(self)

        return {'RUNNING_MODAL'}

    def modal(self, ctx, evt):
        if evt.type == 'TIMER':
            # ret, frame = self.webcam.read()
            # while not ret:
            #     ret, frame = self.webcam.read()
            # ret, frame = True, cv2.imread(r"C:\Users\jiangchengwei\AppData\Roaming\Blender Foundation\Blender\3.4\scripts\addons\addon_CK230\ops\bus.jpg")
            # ret, frame = True, cv2.imread(r"C:\blender_imgs\test.jpg")
            # ret, frame =True, np.ones((480, 640, 3), dtype="uint8")
            # print(frame.shape)
            # cv2.imshow('frame', frame)
            # cv2.waitKey(1)
            # img_str = encodeImage(frame)
            # image = json.dumps(
            #     {
            #         'type': 'image',
            #         'content': 'none',
            #     }
            # ).encode('utf-8')
            # img_pack = packDataFlow(image)
            # self.client.send(img_pack)
            data = self.client.recv(4096).decode('utf-8')
            if data != 'none':
                try:
                    data = json.loads(data)
                except:
                    return {'RUNNING_MODAL'}
                if not drivenCharacter(data['poses'], data['trans']):
                    self.report({"ERROR"}, "The armature/character's name is not found or the bones is not fixed")
                    self.close(ctx)
                    return {'FINISHED'}

        if evt.type == 'ESC':
            self.close(ctx)
            return {'FINISHED'}

        return {'RUNNING_MODAL'}

    # def modal(self, ctx, evt):
    #
    #     return {'RUNNING_MODAL'}

    def close(self, ctx):
        # done = json.dumps(
        #     {
        #         'type': 'done'
        #     }
        # ).encode('utf-8')
        # self.client.send(packDataFlow(done))
        self.client.close()
        # cv2.destroyAllWindows()
        # self.webcam.release()
        ctx.window_manager.event_timer_remove(self._timer)
