import json
import logging
import struct
from pathlib import Path

import numpy as np
import zmq

from flandre import C
from flandre.nodes.Node import Node
from flandre.utils.Msg import ImageArgMsg, KillMsg, SetSeqMetaMsg, SetPlayMode, SetDeviceConfigMsg, SetRecordMsg, \
    RequestRfFrameMsg, RecordFrameMsg, RobotRtsiMsg, SeqMetaMsg
from flandre.utils.RfFrame import b2t
from flandre.utils.RfMeta import RfFrameMeta, RfSequenceMeta

logger = logging.getLogger(__name__)


class Recorder(Node):
    topics = [SetSeqMetaMsg, SeqMetaMsg, SetPlayMode, SetDeviceConfigMsg, SetRecordMsg, 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
        self.arr: list[RobotRtsiMsg] = []
        self.isrecord = False

    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)
        last_ts = 0
        while True:
            socks = dict(self.c.poller.poll())
            for k in socks:
                if k == device_socket:
                    buffer = device_socket.recv()
                    ts, sequence_id, encoder, s = b2t(buffer)
                    if last_ts == ts:
                        logger.warning(f"Record {ts}")
                    last_ts = ts
                    if self.seq_meta_live is not None and s.__len__() // 2 == np.prod(self.seq_meta_live.shape):
                        (self.record_path / f'{ts}.bin').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, SeqMetaMsg):
                        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:
                            self.isrecord = True
                            self.arr.clear()
                            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:
                            self.isrecord = False
                            device_socket.disconnect(C.live_push_socket)
                            arrj = []
                            for msg in self.arr:
                                arrj.append(msg.dict)
                            (self.record_path / f'rtsi.{self.arr[0].ts}.json').write_text(json.dumps(arrj))

                    elif isinstance(msg, RobotRtsiMsg):
                        self.rtsi = msg
                        if self.isrecord:
                            self.arr.append(self.rtsi)
