import struct

from enum import Enum
from typing import Union, Type

import forcen_public_utils.checked as ch


AllowablePytypes = Union[int, float, bool, str, None]
AllowablePyObjects = Union[Type[int], Type[float], Type[bool], Type[str]]


class DataTypes(Enum):
    INT8 = "int8"
    INT16 = "int16"
    INT32 = "int32"
    INT64 = "int64"

    UINT8 = "uint8"
    UINT16 = "uint16"
    UINT32 = "uint32"
    UINT64 = "uint64"

    FLOAT32 = "float32"
    FLOAT64 = "float64"

    BOOL = "bool"

    STR = "str"


def data_type_text_to_enum(text: str) -> ch.MaybeVal[DataTypes]:
    for typ in DataTypes:
        if text == typ.value:
            return ch.ok(typ)
    return ch.bad(message="Cannot find type")


class DataTypeFormats(Enum):
    INT8_FMT = "b"
    INT16_FMT = "h"
    INT32_FMT = "i"
    INT64_FMT = "q"

    UINT8_FMT = "B"
    UINT16_FMT = "H"
    UINT32_FMT = "I"
    UINT64_FMT = "Q"

    FLOAT32_FMT = "f"
    FLOAT64_FMT = "d"

    BOOL_FMT = "?"

    STR_FMT = "s"


class ByteConverter:
    ENDIAN_FMT = "<"

    DATATYPE_TO_FMT = {
        DataTypes.INT8.value: DataTypeFormats.INT8_FMT.value,
        DataTypes.INT16.value: DataTypeFormats.INT16_FMT.value,
        DataTypes.INT32.value: DataTypeFormats.INT32_FMT.value,
        DataTypes.INT64.value: DataTypeFormats.INT64_FMT.value,
        DataTypes.UINT8.value: DataTypeFormats.UINT8_FMT.value,
        DataTypes.UINT16.value: DataTypeFormats.UINT16_FMT.value,
        DataTypes.UINT32.value: DataTypeFormats.UINT32_FMT.value,
        DataTypes.UINT64.value: DataTypeFormats.UINT64_FMT.value,
        DataTypes.FLOAT32.value: DataTypeFormats.FLOAT32_FMT.value,
        DataTypes.FLOAT64.value: DataTypeFormats.FLOAT64_FMT.value,
        DataTypes.BOOL.value: DataTypeFormats.BOOL_FMT.value,
        DataTypes.STR.value: DataTypeFormats.STR_FMT.value,
    }

    STR_ENCODING = "utf-8"

    NUMERIC_TYPES = [
        DataTypes.INT8.value,
        DataTypes.INT16.value,
        DataTypes.INT32.value,
        DataTypes.INT64.value,
        DataTypes.UINT8.value,
        DataTypes.UINT16.value,
        DataTypes.UINT32.value,
        DataTypes.UINT64.value,
        DataTypes.FLOAT32.value,
        DataTypes.FLOAT64.value,
    ]

    INT_TYPES = [
        DataTypes.INT8.value,
        DataTypes.INT16.value,
        DataTypes.INT32.value,
        DataTypes.INT64.value,
        DataTypes.UINT8.value,
        DataTypes.UINT16.value,
        DataTypes.UINT32.value,
        DataTypes.UINT64.value,
    ]

    FLOAT_TYPES = [DataTypes.FLOAT32.value, DataTypes.FLOAT64.value]

    def _size_of_msg_type(self, msgtype: str) -> ch.MaybeVal[int]:
        return ch.ok(struct.calcsize(self.DATATYPE_TO_FMT[msgtype]))

    def bytes_to_msg(
        self, msg_in_bytes: bytearray, msgtype: str
    ) -> ch.MaybeVal[Union[AllowablePytypes, None]]:
        if len(msg_in_bytes) == 0:
            return ch.ok("")

        if msgtype not in self.DATATYPE_TO_FMT.keys():
            return ch.bad(
                message=f"Message type {msgtype} is not supported. Check the jig configuration files",
            )

        if msgtype == DataTypes.STR.value:
            try:
                return ch.ok(msg_in_bytes.decode(self.STR_ENCODING))
            except ValueError as e:
                return ch.bad(message=f"Count not unpack bytes due to error: {e}")

        else:
            fmt = f"{self.ENDIAN_FMT}{self.DATATYPE_TO_FMT[msgtype]}"

            # struct.unpack always returns a tuple so grab first element to get raw datatype
            try:
                return ch.ok(struct.unpack(fmt, msg_in_bytes)[0])
            except (struct.error, IndexError) as e:
                return ch.bad(message=f"Count not unpack bytes due to error {e}")

    def msg_to_bytes(
        self, msg: AllowablePytypes, msgtype: str
    ) -> ch.MaybeVal[bytearray]:
        if msg is None:
            return ch.ok(bytearray())

        if msgtype not in self.DATATYPE_TO_FMT.keys():
            return ch.bad(
                message=f"Message type {msgtype} is not supported. Check the jig configuration files",
            )

        if msgtype == DataTypes.STR.value and type(msg) is str:
            try:
                return ch.ok(bytearray(msg.encode(self.STR_ENCODING)))
            except ValueError as e:
                return ch.bad(message=f"Could not pack value due to error {e}")

        else:
            fmt = f"{self.ENDIAN_FMT}{self.DATATYPE_TO_FMT[msgtype]}"
            try:
                return ch.ok(bytearray(struct.pack(fmt, msg)))
            except struct.error as e:
                return ch.bad(message=f"Could not pack value due to error {e}")
