import datetime
from socketserver import BaseRequestHandler, UDPServer
import time
import os, sys
import open3d as o3d
from multiprocessing import  Pool
import numpy as np

from ..utils import file_fliter, write_ply
from .servers import MicroWaveReceiver
from .parsers import MicroWaveParser

# # add python path of src to sys.path
# src_path = os.path.join(__file__, *(['..'] * 2))
# src_path = os.path.abspath(src_path)
# sys.path.insert(1, src_path)

# udp_get
class PointCloudUDPRequestHandler(BaseRequestHandler):

    def handle(self):
        packet_threshold = self.server.packet_threshold 
        packed_count = self.server.packed_count
        path = self.server.path

        print(time.asctime(),' Got connection from {}, {}'.format(self.client_address, packed_count))
        # Get message and client socket
        msg, sock = self.request
        time_tag = datetime.datetime.now().strftime('%Y-%m-%d-%H-%M-%S')
        f = open('{}/{}_packet{}.udp'.format(path,time_tag,packed_count), 'wb')
        f.write(msg)
        f.close()
        self.server.packed_count += 1
        self.server.udp_list.append(msg)
        
        if packed_count >= packet_threshold:
            UDPServer.server_close(self.server)

def udp_get_tool(microwaveHost,packetThreshold,deviceName,saveFolder,subFolder):

    # host = ''
    # host = '192.168.2.65' # use this host to get the udp data from microwave device
    port = 9911
    server =UDPServer((microwaveHost, port), PointCloudUDPRequestHandler)


    if subFolder is not None:
        path = os.path.join(saveFolder,subFolder)
    else:
        start = datetime.datetime.now()
        path = os.path.join(saveFolder,f"{deviceName}-{start.strftime('%Y-%m-%d-%H-%M-%S')}")
    if not os.path.exists(path):
        os.makedirs(path)

    server.path = os.path.abspath(path)
    server.packet_threshold = packetThreshold
    server.packed_count = 0
    server.udp_list = []

    print(time.asctime(),'Server started on port', port)
    print('....')
    print('ctrl-c to quit server.')
    try:
        server.serve_forever()
    except KeyboardInterrupt:
        server.server_close()
        print(time.asctime(),"Server Stopped")
    except:
        server.server_close()
        print(time.asctime(),"Server Stopped")
    return server.udp_list

# udp2bin
def udp2bin_tool(udp_folder_Or_list):
    Save = True
    bin_list = []

    if isinstance(udp_folder_Or_list, list):
        udp_files = udp_folder_Or_list
        Save = False
    else:
        UDPfolder = udp_folder_Or_list
        # example pod1-2022-02-21-13-25-03_packet125.udp --> 125
        sort_lambda = lambda x:int(x.split('packet')[1][:-4])
        udp_files = file_fliter(UDPfolder,'udp', sort_lambda)
        Save = True

    receiver = MicroWaveReceiver()


    for file in udp_files:
        data = None
        if isinstance(file,str):
            filename = os.path.join(UDPfolder, file)
            with open(filename, "rb") as f:
                data = f.read()
        else:
            data = file
        
        receiver.AddBuf(data)
        if receiver.DataFrameReady:
            frame_number = receiver.Header.FrameNumber
            if Save:
                bin_name = os.path.join(UDPfolder, "{}.bin".format(frame_number))
                print(time.asctime(),"Frame {} ready".format(frame_number))
                with open(bin_name, "wb") as bin_file:
                    bin_file.write(receiver.DataFrame)
                    bin_file.close()
            
            bin_list.append((frame_number,receiver.DataFrame))

    return bin_list

def bin2ply_f(folder,file):

    with open(os.path.join(folder,file), "rb") as bin_file:
        print(f"Extracting the data in {file} ...")
        buf = bin_file.read()
        buf_data = MicroWaveParser.ParserFrame(buf)
        for frame in buf_data:
            
            Frame_id = frame[0]
            FrameNumber = frame[1]
            point_cloud = np.array(frame[2])
            if len(point_cloud) == 0 :
                print(f"Notice: Frame {FrameNumber} is empty!")
            else:
                write_ply(point_cloud,folder,str(FrameNumber))

# bin2ply
def bin2ply_tool(folder,nums_worker = 4):
    pool = Pool(processes = nums_worker)

    bin_files = file_fliter(folder,'bin')

    for file in bin_files:
        pool.apply_async(func=bin2ply_f,args=(folder,file,))

    pool.close()
    pool.join()


def plot_ply(path):
    if isinstance(path,str):
        abspath = os.path.abspath(path)
    else:
        raise ValueError("Not a path!")
    
    if os.path.isdir(abspath):
        # open folder to load all ply files
        ply_files = file_fliter(abspath,'ply')
        show = []
        for ply in ply_files:
            pcd = o3d.io.read_point_cloud(os.path.join(abspath,ply))
            show.append(pcd)
    else:
        show = [o3d.io.read_point_cloud(abspath)]

    o3d.visualization.draw_geometries(show)
    return