import dataclasses
import json
import struct
from enum import auto, Enum
from pathlib import Path

import numpy as np

from flandre.utils.RfFrame import RfFrame, RfFrameMemory, RfFrameFile
from flandre.utils.RfMeta import RfFrameMeta, RfSequenceMeta
from flandre.utils.RfSequence import RfSequence


class BG(Enum):
    Msg1 = auto()
    Msg2 = auto()
    RefreshDeviceMsg = auto()
    BMMsg = auto()
    TickMsg = auto()
    KillMsg = auto()
    StrMsg = auto()
    MoveAxisMsg = auto()
    ImageArgMsg = auto()
    SetSeqMetaMsg = auto()
    SeqMetaMsg = auto()
    SeqIdMinMax = auto()
    SetBaseMsg = auto()
    PlaybackSeqListMsg = auto()
    SetPlayMode = auto()
    SetDeviceEnabledMsg = auto()
    SetDeviceConnectedMsg = auto()
    DeviceEnabledMsg = auto()
    DeviceConnectedMsg = auto()
    SetDeviceConfigMsg = auto()
    DeviceOnlineMsg = auto()
    DeviceConfigListMsg = auto()
    SetRecordMsg = auto()
    RobotRtsiMsg = auto()
    RecordFrameMsg = auto()
    SeqIdList = auto()
    SetWindowVisibleMsg = auto()
    SetSidMsg = auto()
    ImagingConfigNameListMsg = auto()

    BytesMsg = auto()
    BeamformerMsg = auto()
    RequestRfFrameMsg = auto()
    MidiMsg = auto()
    JoystickMsg = auto()
    DeviceZero = auto()
    SetDeviceSwitchMsg = auto()
    DeviceSwitchMsg = auto()
    HeaderByteMsg = auto()

    RfFrameMsg = auto()
    # RfFrameWithMetaMsg = auto()
    RfMatMsg = auto()

    KeyPressMsg = auto()
    RGB888Msg = auto()


class Msg:
    MAGIC = 1919810

    @classmethod
    def decode_base(cls, data: bytes) -> 'Msg':
        return cls(**json.loads(data.decode()))

    def encode(self) -> bytes:
        return json.dumps(self.__dict__).encode()

    @classmethod
    def decode(cls, data: bytes) -> 'Msg':
        return cls.decode_base(data)

    @classmethod
    def decode_msg(cls, msg: bytes):
        magic, eid = struct.unpack('II', msg[:8])
        assert magic == cls.MAGIC
        class_: 'Msg' = globals()[BG(eid).name]
        return class_.decode(msg[8:])

    @classmethod
    def eid(cls):
        return struct.pack('I', BG[cls.__name__].value)

    @classmethod
    def magic(cls):
        return struct.pack('I', cls.MAGIC)

    def encode_msg(self):
        return self.magic() + self.eid() + self.encode()


class HeaderByteMsg(Msg):
    header: dict
    data: bytes

    def __init__(self, header: dict, data: bytes):
        self.header = header
        self.data = data
        if self.data is None:
            self.data = b''

    def encode(self) -> bytes:
        e = json.dumps(self.header).encode()
        return struct.pack('I', e.__len__()) + e + self.data

    @classmethod
    def decode(cls, data: bytes) -> 'HeaderByteMsg':
        header_len = struct.unpack('I', data[:4])[0]
        header = json.loads(data[4:4 + header_len])
        data = data[4 + header_len:]
        return HeaderByteMsg(header, data)


@dataclasses.dataclass
class Msg1(Msg):
    a: int = 0
    b: int = 1


class RefreshDeviceMsg(Msg):
    pass


@dataclasses.dataclass
class KillMsg(Msg):
    name: str = ''


@dataclasses.dataclass
class Msg2(Msg):
    a: int = 2
    b: int = 2


@dataclasses.dataclass
class TickMsg(Msg):
    time: float = 0


@dataclasses.dataclass
class StrMsg(Msg):
    value: str = ''


@dataclasses.dataclass
class KeyPressMsg(StrMsg):
    pass


@dataclasses.dataclass
class BoolMsg(Msg):
    value: bool


@dataclasses.dataclass
class SetWindowVisibleMsg(Msg):
    sender: str
    name: str
    value: bool


class SetDeviceEnabledMsg(BoolMsg):
    pass


class SetDeviceConnectedMsg(BoolMsg):
    pass


class DeviceEnabledMsg(BoolMsg):
    pass


class DeviceConnectedMsg(BoolMsg):
    pass


class DeviceOnlineMsg(BoolMsg):
    pass


@dataclasses.dataclass
class SetRecordMsg(Msg):
    enable: bool
    commit: str = ''
    base: str = ''


@dataclasses.dataclass
class SetDeviceConfigMsg(Msg):
    name: str
    txt: str


class SetBaseMsg(StrMsg):
    pass


@dataclasses.dataclass
class PlaybackSeqListMsg(Msg):
    value: list[str]


@dataclasses.dataclass
class ImagingConfigNameListMsg(Msg):
    value: list[str]


class SetPlayMode(StrMsg):
    pass


@dataclasses.dataclass
class SetSeqMetaMsg(Msg):
    target: str
    name: str


@dataclasses.dataclass
class SeqMetaMsg(Msg):
    target: str
    name: str


@dataclasses.dataclass
class SeqIdMinMax(Msg):
    min: int
    max: int


@dataclasses.dataclass
class SeqIdList(Msg):
    li: list[int]


@dataclasses.dataclass
class MoveAxisMsg(Msg):
    sender: str
    axis: str
    value: int


@dataclasses.dataclass
class ImageArgMsg(Msg):
    sender: str
    t_end: int
    t_start: int = 0
    v2: int = 5900
    dct_center: int = 0
    dct_bandwidth: int = 0
    f_rows: int = 0
    beta: int = 10

    @staticmethod
    def from_path(p: Path):
        return ImageArgMsg(
            sender='load',
            **json.loads(p.read_text())
        )

    def json(self):
        arg_d = dict()
        for field in dataclasses.fields(ImageArgMsg):
            match field.name:
                case 'sender':
                    pass
                case _:
                    arg_d[field.name] = self.__getattribute__(field.name)
        return json.dumps(arg_d)


@dataclasses.dataclass
class DeviceConfigListMsg(Msg):
    arr: list[tuple[str, str]]


@dataclasses.dataclass
class RecordFrameMsg(Msg):
    size: int
    current_sid: int


@dataclasses.dataclass
class SetSidMsg(Msg):
    value: int


class BMMsg(Msg):
    def __init__(self, t: int, data: bytes):
        self.data = data
        self.t = t

    def encode(self) -> bytes:
        return struct.pack('I', self.t) + self.data

    @classmethod
    def decode(cls, data: bytes) -> 'Msg':
        return cls(
            struct.unpack('I', data[:4])[0],
            data[4:]
        )


class RfFrameMsg(HeaderByteMsg):
    def __init__(self, sender: int, rf_frame: RfFrame):
        self.sender = sender
        self.rf_frame = rf_frame
        if isinstance(rf_frame, RfFrameFile):
            super().__init__(dict(
                sender=sender,
                type='RfFrameFile',
                frame_meta=rf_frame.meta.name,
                frame_filename=rf_frame.filename,
                seq_path=rf_frame.seq.path.__str__(),
            ), rf_frame.data)
        elif isinstance(rf_frame, RfFrameMemory):
            super().__init__(dict(
                sender=sender,
                type='RfFrameMemory',
                frame_meta=rf_frame.meta.name,
                seq_meta=rf_frame.seq_meta,
            ), rf_frame.data)

    @classmethod
    def decode(cls, data: bytes) -> 'RfFrameMsg':
        msg = super(RfFrameMsg, cls).decode(data)
        if msg.header['type'] == 'RfFrameFile':
            return RfFrameMsg(
                msg.header['sender'],
                RfFrameFile(
                    RfFrameMeta.from_name(msg.header['frame_meta']),
                    RfSequence(msg.header['seq_path']),
                    msg.data if msg.data.__len__() > 0 else None,
                    msg.header['frame_filename']
                ))
        elif msg.header['type'] == 'RfFrameMemory':
            return RfFrameMsg(
                msg.header['sender'],
                RfFrameMemory(
                    RfFrameMeta.from_name(msg.header['frame_meta']),
                    RfSequenceMeta.from_name(msg.header['seq_meta']),
                    msg.data,
                ))
        raise NotImplementedError()


class RGB888Msg(HeaderByteMsg):

    def __init__(self, b: bytes, w: int, h: int):
        self.w = w
        self.h = h
        super().__init__(dict(h=h, w=w), b)

    @classmethod
    def decode(cls, data) -> 'RGB888Msg':
        msg = super(RGB888Msg, cls).decode(data)
        return RGB888Msg(msg.data, msg.header['w'], msg.header['h'])


class RfMatMsg(HeaderByteMsg):

    def __init__(self, rfmat: 'RfMat'):
        self.rfmat = rfmat

        super().__init__(dict(
            frame_meta=rfmat.frame_meta.name,
            seq_meta=rfmat.seq_meta.name,
            data_shape=rfmat.m.shape,
            dtype=str(rfmat.m.dtype),
        ), rfmat.m.tobytes())

    @classmethod
    def decode(cls, data) -> 'RfMatMsg':
        from flandre.utils.RfMat import RfMat
        msg = super(RfMatMsg, cls).decode(data)
        dt = np.dtype(msg.header['dtype'])
        mat = np.frombuffer(msg.data, dtype=dt).reshape(msg.header['data_shape'])
        rfmat = RfMat(
            mat,
            RfFrameMeta.from_name(msg.header['frame_meta']),
            RfSequenceMeta.from_name(msg.header['seq_meta']),
        )
        return RfMatMsg(rfmat)


# class RfFrameWithMetaMsg(HeaderByteMsg):
#     @dataclasses.dataclass
#     class Header:
#         sender: int
#         meta: str
#         is_zip: bool
# 
#     def __init__(self, sender: int, meta: RfFrameMeta, data: bytes, is_zip: bool = False):
#         self._header = self.Header(sender, meta.name, is_zip)
#         super().__init__(self._header.__dict__, data)
# 
#     @property
#     def sender(self) -> int:
#         return self._header.sender
# 
#     @property
#     def meta(self) -> RfFrameMeta:
#         return RfFrameMeta.from_name(self._header.meta)
# 
#     @classmethod
#     def decode(cls, data) -> 'RfFrameWithMetaMsg':
#         msg = super(RfFrameWithMetaMsg, cls).decode(data)
#         header = cls.Header(**msg.header)
#         return RfFrameWithMetaMsg(header.sender, RfFrameMeta.from_name(header.meta), msg.data, header.is_zip)


@dataclasses.dataclass
class BytesMsg(Msg):
    value: bytes

    def encode(self) -> bytes:
        return self.value

    @classmethod
    def decode(cls, data: bytes) -> 'Msg':
        return cls(data)

    def split(self):
        pass
        # id2 = r.index(Msg.magic(), 1)
        # id3 = r.index(Msg.magic(), id2 + 1)
        # print(id2, id3)
        # seq_msg: SetSeqMetaMsg = Msg.decode_msg(r[0:id2])
        # arg_msg: ImageArgMsg = Msg.decode_msg(r[id2:id3])
        # b_msg: RfFrameMsg = Msg.decode_msg(r[id3:-1])


class BeamformerMsg(BytesMsg):
    pass


class DeviceZero(Msg):
    pass


@dataclasses.dataclass
class MidiMsg(Msg):
    type: str
    channel: int = None
    pitch: int = None
    control: int = None
    value: int = None
    velocity: int = None
    note: int = None


@dataclasses.dataclass
class JoystickMsg(Msg):
    key: str
    value: float


class DeviceSwitchMsg(StrMsg):
    pass


class SetDeviceSwitchMsg(BoolMsg):
    pass


@dataclasses.dataclass
class RobotRtsiMsg(Msg):
    pos: tuple[int, int, int, int, int, int]
    force: tuple[int, int, int, int, int, int]


class RequestRfFrameMsg(Msg):
    pass


def test():
    values = set(item.name for item in BG)
    for k in globals().keys():
        if k.endswith('Msg') and k not in ['Msg', 'BoolMsg']:
            if k not in values:
                raise RuntimeError(f"Unknown msg type: {k}")


test()

if __name__ == '__main__':
    # c = HeaderByteMsg(dict(a=1, b='s'), b'asdasd')
    # c2 = c.decode(c.encode())
    # print(c2.header, c2.data)

    # c = RfFrameWithMetaMsg(1, RfFrameMeta(1, 11, 111), b'asdasdasdads', True)
    # print(c.sender, c.meta, c.data, c._header)
    # c2 = c.decode(c.encode())
    # print(c2.sender, c2.meta, c2.data, c2._header)

    rfmat = RfMat(np.array([1, 2, 3]), RfFrameMeta(), RfSequenceMeta())
    e = RfMatMsg(rfmat).encode()
    d = RfMatMsg.decode(e)
