import os, sys
from ctypes import sizeof, string_at
from socketserver import BaseRequestHandler, UDPServer
from .components import *
from .constants import *

# 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(0, src_path)

from protos import aerial_pb2 as pb



UDP_STATE_IDLE = 0
UDP_STATE_HAND_SHAKE = 1
UDP_STATE_DATA = 2
UDP_STATE_FOOTER  = 3

class MicroWaveReceiver:
    """
    micro wave data receiver
    use AddBuf method add udp packet to receiver
    when DataFrameReady == True, you can get data frame
    """
    def __init__(self):
        self.state = UDP_STATE_IDLE
        self.PointCount = 0
        self.TrackCount = 0
        self.DataFrameReady = False
        self.DataFrame = None 
        self.Header = None
        self.Footer = None
    
    def _receive_data(self, data) -> None:
        """
        receive data packet from udp
        """
        self.data_buf.extend(data)
        self.ReceivedDataLength += len(data)
        if self.ReceivedDataLength == self.DataLength:
            # if received data length equal expected data length
            # then we can get footer block
            self.state = UDP_STATE_FOOTER
        return

    def _receive_footer(self, data) -> None:
        # receive footer block
        self.footer_buf.extend(data)
        self.ReceivedFooterLength += len(data)
        if self.ReceivedFooterLength == sizeof(FooterBlock):
            self.Footer = FooterBlock.from_buffer_copy(self.footer_buf)
            # after received footer block, the data frame is ready
            self.DataFrameReady = True
            # reset state to idle, ready for next data frame
            self.state = UDP_STATE_IDLE
        return
    
    def AddBuf(self, buf: bytes) -> None:
        # add a udp packet to receiver
        buf_size = len(buf)
        if self.state == UDP_STATE_IDLE:
            if buf_size == sizeof(Handshake):
                handshake = Handshake.from_buffer_copy(buf)
                if HAND_SHAKE_MAGIC == string_at(handshake.Magic, 8):
                    self.DataFrameReady = False
                    self.state = UDP_STATE_HAND_SHAKE
                    self.DataFrame = None
                    self.Header = None
                    self.Footer = None
        elif self.state == UDP_STATE_HAND_SHAKE:
            # after received handshake packet, we can get header block 
            self.Header = HeaderBlock.from_buffer_copy(buf[:sizeof(HeaderBlock)])
            if HEADER_MAGIC == string_at(self.Header.Magic, 8):
                self.DataFrame = bytearray(buf)
                N = self.Header.NumberOfDetection
                Nt = self.Header.NumberOfTrack
                self.PointCount = 0
                self.TrackCount = 0
                self.DataLength = N * 8 + Nt * sizeof(TrackerBlock)
                self.RemainLength = self.DataLength + sizeof(FooterBlock)
                self.data_buf = bytearray()
                self.footer_buf = bytearray()
                self.Points = []
                self.ReceivedDataLength = 0
                self.ReceivedFooterLength = 0
                # the next packet should be data block
                self.state = UDP_STATE_DATA

                header_size = sizeof(HeaderBlock)
                # 处理第一个数据包中的数据
                data_size = min(buf_size - header_size, self.DataLength)
                self._receive_data(buf[header_size: header_size + data_size])
                
                # 判断第一个数据包中是否包含Footer Block
                if (header_size + data_size) < buf_size:
                    self._receive_footer(buf[header_size + data_size])
        elif self.state == UDP_STATE_DATA:
            # receive data block
            self.DataFrame.extend(buf)
            remain_data_size = self.DataLength - self.ReceivedDataLength
            data_size = min(buf_size, remain_data_size)
            self._receive_data(buf[: data_size])
            if data_size < buf_size:
                self._receive_footer(buf[data_size: ])
        elif self.state == UDP_STATE_FOOTER:
            # receive footer block
            self.DataFrame.extend(buf)
            self._receive_footer(buf)



class MicroWaveRequestHandler(BaseRequestHandler):
    """
    micro wave udp request handler
    """
    def handle(self):
        # Get message and client socket
        buf, _ = self.request
        receiver = self.server.Receiver
        # add buffer to receiver
        receiver.AddBuf(buf)
        if receiver.DataFrameReady:
            # dataframe ready, add to udp_queue
            msg = pb.RequestCommand()
            msg.sys_command = pb.SysCommand.MICRO_WAVE_COMMAND
            msg.sub_command = pb.SubCommand.MICRO_WAVE_SUB_COMMAND
            data1 = pb.MicroWaveMessage()
            data1.length = len(receiver.DataFrame)
            data1.data = bytes(receiver.DataFrame)
            data1_str = data1.SerializeToString()
            msg.data1.length = len(data1_str)
            msg.data1.data = data1_str
            self.server.msg_queue.put(msg)


class MicroWaveUDPServer(UDPServer):
    """
    micro wave udp server
    """
    def __init__(self, server_address, msg_queue):
        super().__init__(server_address, MicroWaveRequestHandler)
        # create receiver instance
        self.Receiver = MicroWaveReceiver()
        self.msg_queue = msg_queue