import os, sys
from ctypes import sizeof, string_at
from socketserver import BaseRequestHandler, UDPServer
import warnings
from .components import *
from .constants import *
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 _is_handshake(self,data):
        handshake = Handshake.from_buffer_copy(data)
        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
            return True
        else:
            return False

    def _get_header(self, buf):
        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
            return True
        else:
            False
    
    def _receive_data(self, data) -> None:
        """
        receive data packet from udp
        """
        self.data_buf.extend(bytes(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 True
        elif self.ReceivedDataLength > self.DataLength:
            self.state = UDP_STATE_IDLE
            warnings.warn(f"encounter an error when getting the detection data in udp packets ...")
            return False
        else:
            return True

    def _receive_footer(self, data) -> None:
        # receive footer block
        self.footer_buf.extend(bytes(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 buf_size == sizeof(Handshake):
            self._is_handshake(buf)

        if self.state == UDP_STATE_HAND_SHAKE:
            # after received handshake packet, we can get header block 
            try:
                # make sure to get a complete header
                result = self._get_header(buf)
            except Exception as e:
                warnings.warn(f"encounter an unknow error {e} when getting the header data in udp packets ...")
                result = False

            if not result:
                # waiting for the next handshake udp packet
                self.state = UDP_STATE_IDLE
            else:
                # process the detection data in the first data packets
                header_size = sizeof(HeaderBlock)
                data_size = min(buf_size - header_size, self.DataLength)

                if 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(bytes(buf))
            remain_data_size = self.DataLength - self.ReceivedDataLength
            data_size = min(buf_size, remain_data_size)
            if 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(bytes(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