import logging
import struct
from pathlib import Path

import numpy as np
import zmq

from flandre.config import LIVE_SOCKET,C
from flandre.nodes.Node import Node
from flandre.utils.Msg import ImageArgMsg, KillMsg, SetSeqMetaMsg, SetPlayMode, SetDeviceConfigMsg, SetRecordMsg, \
    RfFrameWithMetaMsg, RequestRfFrameMsg, RecordFrameMsg, RobotRtsiMsg
from flandre.utils.RfFile import RfSequenceMeta
from flandre.utils.RfMeta import RfFrameMeta

logger = logging.getLogger(__name__)


class Recorder(Node):
    topics = [SetSeqMetaMsg, SetPlayMode, SetDeviceConfigMsg, SetRecordMsg, RfFrameWithMetaMsg, RequestRfFrameMsg,
              ImageArgMsg, RobotRtsiMsg]

    def __init__(self, level=logging.INFO):
        super(Recorder, self).__init__(level=level)
        self.seq_meta_live = None
        self.record_enable = False
        self.record_commit = None
        self.record_path: Path | None = None
        self.play_mode = None

    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.c.poller.register(device_socket, zmq.POLLIN)
        while True:
            socks = dict(self.c.poller.poll())
            for k in socks:
                if k == device_socket:
                    buffer = device_socket.recv()
                    _, sequence_id, encoder = struct.unpack_from('=iqi', buffer)
                    s = buffer[4 + 8 + 4:]
                    if self.seq_meta_live is not None and s.__len__() // 2 == np.prod(self.seq_meta_live.shape):
                        (self.record_path / 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],
                        ).filename).write_bytes(s)
                        self.send(RecordFrameMsg(s.__len__(), sequence_id))
                if k == self.c.sub:
                    msg = self.recv()
                    if isinstance(msg, KillMsg):
                        if msg.name == '':
                            return
                    elif isinstance(msg, SetSeqMetaMsg):
                        match msg.target:
                            case 'live':
                                self.seq_meta_live = RfSequenceMeta.from_name(msg.name)
                    elif isinstance(msg, SetRecordMsg):
                        self.record_enable = msg.enable
                        if msg.enable:
                            seq_meta = self.seq_meta_live
                            self.record_commit = msg.commit
                            seq_meta.commit = msg.commit
                            p = Path(msg.base) / seq_meta.name
                            p.mkdir(parents=True, exist_ok=True)
                            self.record_path = p
                            device_socket.connect(C.live_push_socket)
                        else:
                            device_socket.disconnect(C.live_push_socket)
                    elif isinstance(msg, RobotRtsiMsg):
                        self.rtsi = msg
