import logging
import struct
import time
from pathlib import Path

import zmq

from flandre.config import C
from flandre.nodes.Node import Node
from flandre.utils.Msg import ImageArgMsg, KillMsg, SetSeqMetaMsg, SetPlayMode, SetDeviceConfigMsg, \
    ImagingConfigNameListMsg, BeamformerMsg, RobotRtsiMsg, SeqMetaMsg, DeviceEnabledMsg, RfFrameMsg
from flandre.utils.RfFrame import RfFrameMemory
from flandre.utils.RfMeta import RfFrameMeta, RfSequenceMeta

logger = logging.getLogger(__name__)


class Muxer(Node):
    topics = [SetSeqMetaMsg, SetPlayMode, SetDeviceConfigMsg, RfFrameMsg,
              ImageArgMsg, RobotRtsiMsg, SeqMetaMsg, DeviceEnabledMsg]

    def __init__(self, level=logging.INFO):
        super(Muxer, self).__init__(level=level)
        self.seq_meta = None
        self.seq_meta_live: RfSequenceMeta = None
        self.seq_meta_playback = None
        self.play_mode = None
        self.rep_socket: zmq.Socket = None
        self.req_driver_socket: zmq.Socket = None
        self.playback_rf_msg: RfFrameMsg | None = None
        self.device_enabled = False
        self.driver_data_raw = b''

    def custom_setup(self):
        self.rep_socket: zmq.Socket = self.c.ctx.socket(zmq.REP)
        self.rep_socket.bind(f'tcp://localhost:{C.muxer_rep_port}')
        self.req_driver_socket: zmq.Socket = self.c.ctx.socket(zmq.REQ)
        self.req_driver_socket.connect(C.driver_rep_socket)
        self.c.poller.register(self.rep_socket, zmq.POLLIN)

    def handle_rep_socket(self):
        self.rep_socket.recv()
        if self.play_mode is None:
            self.rep_socket.send(BeamformerMsg(b'init').encode_msg())
            return
        match self.play_mode:
            case 'playback':
                # logger.warning(f'test, {self.playback_rf_msg}')
                if self.playback_rf_msg is None:
                    self.rep_socket.send(BeamformerMsg(b'nop').encode_msg())
                    return
                data_msg = self.playback_rf_msg
            case 'live':
                if not self.device_enabled:
                    self.rep_socket.send(BeamformerMsg(b'init').encode_msg())
                    logger.warning('Device not enabled')
                    return
                self.req_driver_socket.send(b'')
                self.driver_data_raw = self.req_driver_socket.recv()
                if self.driver_data_raw == b'':
                    # todo fixit driver no empty
                    self.rep_socket.send(BeamformerMsg(b'nop').encode_msg())
                    return
                _, _, sequence_id, encoder = struct.unpack_from('=iiii', self.driver_data_raw)
                driver_data_body = self.driver_data_raw[4 + 8 + 4:]
                data_msg = RfFrameMsg(0, RfFrameMemory(
                    RfFrameMeta(
                        encoder=encoder,
                        sequence_id=sequence_id,
                        robot_x=self.rtsi.pos[0],
                        robot_y=self.rtsi.pos[1],
                        robot_z=self.rtsi.pos[2],
                        robot_roll=self.rtsi.pos[3],
                        robot_pitch=self.rtsi.pos[4],
                        robot_yal=self.rtsi.pos[5],
                        robot_force_x=self.rtsi.force[0],
                        robot_force_y=self.rtsi.force[1],
                        robot_force_z=self.rtsi.force[2],
                        robot_force_roll=self.rtsi.force[3],
                        robot_force_pitch=self.rtsi.force[4],
                        robot_force_yal=self.rtsi.force[5],
                    ),
                    self.seq_meta_live,
                    driver_data_body
                ))
            case _:
                raise NotImplementedError()
        # if (data_msg.data.__len__() // 2) != data_msg.rf_frame.prod():
        #     self.rep_socket.send(BeamformerMsg(b'nop').encode_msg())
        #     return
        self.rep_socket.send(BeamformerMsg(
            self.arg.encode_msg() +
            data_msg.encode_msg()
        ).encode_msg())

    def loop(self):
        self.rtsi = RobotRtsiMsg(
            pos=(0, 0, 0, 0, 0, 0),
            force=(0, 0, 0, 0, 0, 0),
        )
        device_socket = self.context.socket(zmq.PULL)

        self.arg = ImageArgMsg('', t_start=0, t_end=1499)
        self.c.poller.register(device_socket, zmq.POLLIN)
        self.send(ImagingConfigNameListMsg([path.stem for path in C.imaging_config_folder.glob('*.json')]))
        while True:
            socks = dict(self.c.poller.poll())
            for k in socks:
                if k == device_socket:
                    pass
                if k == self.rep_socket:
                    self.handle_rep_socket()
                if k == self.c.sub:
                    msg = self.recv()
                    if isinstance(msg, KillMsg):
                        if msg.name == '':
                            return
                    elif isinstance(msg, RfFrameMsg):
                        if msg.sender == 1:
                            self.playback_rf_msg = msg
                    elif isinstance(msg, ImageArgMsg):
                        self.arg = msg
                    elif isinstance(msg, SeqMetaMsg):
                        match msg.target:
                            case 'live':
                                self.seq_meta_live = RfSequenceMeta.from_name(msg.name)
                    elif isinstance(msg, SetPlayMode):
                        logger.info(f'set playmode {msg}')
                        self.play_mode = msg.value
                    elif isinstance(msg, RobotRtsiMsg):
                        self.rtsi = msg
                    elif isinstance(msg, DeviceEnabledMsg):
                        self.device_enabled = msg.value
