from asyncio import Future, StreamReader, StreamWriter, Task, shield
import asyncio
import functools
import operator
import struct
from typing import Awaitable, Dict, Optional, Self
from dataclasses import dataclass

import serial

import async_serial
from logger import spr

MSP_MAGIC_OUT = b"$R<"
MSP_MAGIC_IN = b"$R>"
MSP_CODE_POLL = 3

MSP_TTY_DEF_BAUDRATE = 250000

def make_msp(code: int, data: bytes = b'') -> bytes:
    ret = bytearray(MSP_MAGIC_OUT)
    if len(data) > 255:
        raise ValueError()
    ret.append(len(data))
    ret.append(code)
    ret += data
    checksum = functools.reduce(operator.xor, map(int, ret[len(MSP_MAGIC_OUT):]))
    ret.append(checksum)
    return bytes(ret)

async def read_msp(r: StreamReader) -> tuple[int, bytes]:
    magic_pos = 0
    while magic_pos < len(MSP_MAGIC_IN):
        b, = await r.readexactly(1)
        if b == MSP_MAGIC_IN[magic_pos]:
            magic_pos += 1
        else:
            magic_pos = 0
            print('Unexpected byte from RCBenchmark serial')
    size, = await r.readexactly(1)
    code, = await r.readexactly(1)
    data = await r.readexactly(size)
    checksum, = await r.readexactly(1)
    _ = checksum
    return code, data


def make_poll(esc_pwm: int) -> bytes: 
    data = struct.pack('<HHHH', esc_pwm, 0, 0, 0)
    # data = struct.pack('>HHHH', esc_pwm, 0, 0, 0)
    # return make_msp(MSP_CODE_POLL, data)
    return data


#        esc_voltage, esc_current, esc_power, load_thrust, load_left, rot_e, rot_o, temp0, temp1, temp2, basic_data_flag, acc_x, acc_y, acc_z, vibration, raw_pressure_p, raw_pressure_t, load_right, pro_data_flag = struct.unpack_from('<ffffffffffchhhhhhfc', data)

@dataclass
class PollResponse:
    esc_voltage: float
    esc_current: float
    esc_power: float
    load_thrust: float
    load_left: float
    rot_e: float
    rot_o: float
    temp0: float
    temp1: float
    temp2: float
    basic_data_flag: int
    acc_x: int
    acc_y: int
    acc_z: int
    vibration: int
    raw_pressure_p: int
    raw_pressure_t: int
    load_right: float
    pro_data_flag: int

    @classmethod
    def from_bytes(cls, data: bytes) -> Self:
        unpacked = struct.unpack_from('<ffffffffffchhhhhhfc', data)
        # print(unpacked)
        return cls(*unpacked)




class MSPSlave:
    w: StreamWriter
    r: StreamReader

    pending_requests: Dict[int, Future[bytes]]
    _reader_task: Optional[Task]


    def __init__(self, reader: StreamReader, writer: StreamWriter):
        self.r = reader
        self.w = writer
        self._reader_task = None
        self.pending_requests = {}

    async def ensure_reader(self):
        if self._reader_task is None:
            self._reader_task = asyncio.create_task(self._reader())

    async def _reader(self):
        try:
            while True:
                code, data = await read_msp(self.r)
                self.handle_packet(code, data)
        except Exception as e:
            for req in self.pending_requests.values():
                req.set_exception(e)
            

    def handle_packet(self, code: int, data: bytes):
        if code in self.pending_requests:
            self.pending_requests[code].set_result(data)
        else:
            raise RuntimeError(f'Unexpected packet: {code}')

    def do_request(self, code: int, data: bytes) -> Awaitable[bytes]: 
        if code in self.pending_requests:
            raise RuntimeError('Duplicate request')
        f = Future()
        def future_done(_):
            del self.pending_requests[code]
        self.pending_requests[code] = f
        f.add_done_callback(future_done)
        self.w.write(make_msp(code, data))
        return shield(f)

    async def do_poll(self, esc_pwm: int) -> PollResponse:
        resp = await self.do_request(MSP_CODE_POLL, make_poll(esc_pwm))
        return PollResponse.from_bytes(resp)

    # TODO
    # ASYNC rozjebion
    # TODO
    @staticmethod
    async def open_connection(tty: str, baudrate = MSP_TTY_DEF_BAUDRATE) -> 'MSPSlave':
        s = serial.Serial(port=tty, baudrate = baudrate)

        while True:
            l = s.readline()
            print(l)
            if l == b'Ready\r\n':
                break

        spr('MSP: Ready')

        r, w = async_serial.wrap_serial(s)
        m = MSPSlave(r, w)
        await m.ensure_reader()
        return m