"""Modbus PDU (Protocol Data Unit) codec.

Encodes/decodes request and response PDUs for the function codes
supported by the WAGO 750-362 coupler:

    FC1  Read Coils             (max 2000 bits)
    FC2  Read Discrete Inputs   (max 2000 bits)
    FC3  Read Holding Registers (max 125 words)
    FC4  Read Input Registers   (max 125 words)
    FC5  Write Single Coil
    FC6  Write Single Register
    FC11 Get Comm Event Count
    FC15 Write Multiple Coils   (max 800 bits)
    FC16 Write Multiple Registers (max 100 words)
    FC22 Mask Write Register
    FC23 Read/Write Multiple Registers (125 read / 100 write)
"""

from __future__ import annotations

from dataclasses import dataclass, field

# Function codes supported by the WAGO 750-362 (manual §11.2.2)
FC_READ_COILS = 0x01
FC_READ_DISCRETE_INPUTS = 0x02
FC_READ_HOLDING_REGISTERS = 0x03
FC_READ_INPUT_REGISTERS = 0x04
FC_WRITE_SINGLE_COIL = 0x05
FC_WRITE_SINGLE_REGISTER = 0x06
FC_GET_COMM_EVENT_COUNT = 0x0B
FC_WRITE_MULTIPLE_COILS = 0x0F
FC_WRITE_MULTIPLE_REGISTERS = 0x10
FC_MASK_WRITE_REGISTER = 0x16
FC_READ_WRITE_MULTIPLE_REGISTERS = 0x17

SUPPORTED_FUNCTION_CODES = frozenset(
    {
        FC_READ_COILS,
        FC_READ_DISCRETE_INPUTS,
        FC_READ_HOLDING_REGISTERS,
        FC_READ_INPUT_REGISTERS,
        FC_WRITE_SINGLE_COIL,
        FC_WRITE_SINGLE_REGISTER,
        FC_GET_COMM_EVENT_COUNT,
        FC_WRITE_MULTIPLE_COILS,
        FC_WRITE_MULTIPLE_REGISTERS,
        FC_MASK_WRITE_REGISTER,
        FC_READ_WRITE_MULTIPLE_REGISTERS,
    }
)

# Access limits per 750-362 (statistics counters TooManyRegisters/TooManyBits)
MAX_READ_BITS = 2000
MAX_WRITE_BITS = 800
MAX_READ_REGISTERS = 125
MAX_WRITE_REGISTERS = 100

# Exception codes (manual §11.2.3, Table 83)
EXC_ILLEGAL_FUNCTION = 0x01
EXC_ILLEGAL_DATA_ADDRESS = 0x02
EXC_ILLEGAL_DATA_VALUE = 0x03
EXC_SLAVE_DEVICE_FAILURE = 0x04
EXC_ACKNOWLEDGE = 0x05
EXC_SERVER_BUSY = 0x06

COIL_ON = 0xFF00
COIL_OFF = 0x0000


class PduError(ValueError):
    """Raised for structurally invalid PDUs."""


@dataclass(slots=True, frozen=True)
class ReadRequest:
    """FC1/2/3/4 common request shape."""

    function_code: int
    address: int
    count: int


@dataclass(slots=True, frozen=True)
class WriteSingleCoilRequest:
    function_code: int = FC_WRITE_SINGLE_COIL
    address: int = 0
    value: bool = False


@dataclass(slots=True, frozen=True)
class WriteSingleRegisterRequest:
    function_code: int = FC_WRITE_SINGLE_REGISTER
    address: int = 0
    value: int = 0


@dataclass(slots=True, frozen=True)
class WriteMultipleCoilsRequest:
    function_code: int = FC_WRITE_MULTIPLE_COILS
    address: int = 0
    values: tuple[bool, ...] = field(default=())


@dataclass(slots=True, frozen=True)
class WriteMultipleRegistersRequest:
    function_code: int = FC_WRITE_MULTIPLE_REGISTERS
    address: int = 0
    values: tuple[int, ...] = field(default=())


@dataclass(slots=True, frozen=True)
class MaskWriteRegisterRequest:
    function_code: int = FC_MASK_WRITE_REGISTER
    address: int = 0
    and_mask: int = 0xFFFF
    or_mask: int = 0x0000


@dataclass(slots=True, frozen=True)
class ReadWriteMultipleRegistersRequest:
    function_code: int = FC_READ_WRITE_MULTIPLE_REGISTERS
    read_address: int = 0
    read_count: int = 0
    write_address: int = 0
    values: tuple[int, ...] = field(default=())


@dataclass(slots=True, frozen=True)
class GetCommEventCountRequest:
    function_code: int = FC_GET_COMM_EVENT_COUNT


Request = (
    ReadRequest
    | WriteSingleCoilRequest
    | WriteSingleRegisterRequest
    | WriteMultipleCoilsRequest
    | WriteMultipleRegistersRequest
    | MaskWriteRegisterRequest
    | ReadWriteMultipleRegistersRequest
    | GetCommEventCountRequest
)


class ReadResponse:
    """Response for FC1/2 (bits) or FC3/4 (words)."""

    __slots__ = ("function_code", "bits", "registers")

    def __init__(self, function_code: int, bits: list[bool] | None = None,
                 registers: list[int] | None = None) -> None:
        self.function_code = function_code
        self.bits = bits if bits is not None else []
        self.registers = registers if registers is not None else []

    def encode(self) -> bytes:
        if self.function_code in (FC_READ_COILS, FC_READ_DISCRETE_INPUTS):
            payload = self._pack_bits(self.bits)
            return bytes([self.function_code, len(payload)]) + payload
        data = b"".join(v.to_bytes(2, "big") for v in self.registers)
        byte_count = len(data)
        if byte_count != 2 * len(self.registers):
            raise PduError("register count mismatch")
        return bytes([self.function_code, byte_count]) + data

    @staticmethod
    def _pack_bits(bits: list[bool]) -> bytes:
        out = bytearray((len(bits) + 7) // 8)
        for i, bit in enumerate(bits):
            if bit:
                out[i // 8] |= 1 << (i % 8)
        return bytes(out)


class WriteSingleCoilResponse:
    __slots__ = ("function_code", "address", "value")

    def __init__(self, address: int, value: bool) -> None:
        self.function_code = FC_WRITE_SINGLE_COIL
        self.address = address
        self.value = value

    def encode(self) -> bytes:
        return (
            bytes([self.function_code])
            + self.address.to_bytes(2, "big")
            + (COIL_ON if self.value else COIL_OFF).to_bytes(2, "big")
        )


class WriteSingleRegisterResponse:
    __slots__ = ("function_code", "address", "value")

    def __init__(self, address: int, value: int) -> None:
        self.function_code = FC_WRITE_SINGLE_REGISTER
        self.address = address
        self.value = value

    def encode(self) -> bytes:
        return (
            bytes([self.function_code])
            + self.address.to_bytes(2, "big")
            + self.value.to_bytes(2, "big")
        )


class WriteMultipleResponse:
    __slots__ = ("function_code", "address", "count")

    def __init__(self, function_code: int, address: int, count: int) -> None:
        self.function_code = function_code
        self.address = address
        self.count = count

    def encode(self) -> bytes:
        return (
            bytes([self.function_code])
            + self.address.to_bytes(2, "big")
            + self.count.to_bytes(2, "big")
        )


class MaskWriteRegisterResponse:
    __slots__ = ("function_code", "address", "and_mask", "or_mask")

    def __init__(self, address: int, and_mask: int, or_mask: int) -> None:
        self.function_code = FC_MASK_WRITE_REGISTER
        self.address = address
        self.and_mask = and_mask
        self.or_mask = or_mask

    def encode(self) -> bytes:
        return (
            bytes([self.function_code])
            + self.address.to_bytes(2, "big")
            + self.and_mask.to_bytes(2, "big")
            + self.or_mask.to_bytes(2, "big")
        )


class GetCommEventCountResponse:
    __slots__ = ("function_code", "status", "event_count")

    def __init__(self, status: int = 0, event_count: int = 0) -> None:
        self.function_code = FC_GET_COMM_EVENT_COUNT
        self.status = status
        self.event_count = event_count

    def encode(self) -> bytes:
        return (
            bytes([self.function_code])
            + self.status.to_bytes(2, "big")
            + self.event_count.to_bytes(2, "big")
        )


class ExceptionResponse:
    __slots__ = ("function_code", "exception_code")

    def __init__(self, function_code: int, exception_code: int) -> None:
        self.function_code = function_code
        self.exception_code = exception_code

    def encode(self) -> bytes:
        return bytes([self.function_code | 0x80, self.exception_code])


Response = (
    ReadResponse
    | WriteSingleCoilResponse
    | WriteSingleRegisterResponse
    | WriteMultipleResponse
    | MaskWriteRegisterResponse
    | GetCommEventCountResponse
    | ExceptionResponse
)


def decode_request(pdu: bytes) -> Request:
    """Decode a request PDU.

    Raises :class:`PduError` for structurally invalid PDUs
    (the caller then answers with exception 0x03).
    Unsupported function codes raise :class:`UnsupportedFunction`,
    a subclass, so the server can reply 0x01 while counting BadFunction.
    """
    if not pdu:
        raise PduError("empty PDU")
    fc = pdu[0]
    if fc not in SUPPORTED_FUNCTION_CODES:
        raise UnsupportedFunction(fc)
    body = pdu[1:]
    if fc in (FC_READ_COILS, FC_READ_DISCRETE_INPUTS, FC_READ_HOLDING_REGISTERS,
              FC_READ_INPUT_REGISTERS):
        if len(body) != 4:
            raise PduError(f"FC{fc}: expected 4 data bytes, got {len(body)}")
        address = int.from_bytes(body[0:2], "big")
        count = int.from_bytes(body[2:4], "big")
        max_count = (
            MAX_READ_BITS if fc in (FC_READ_COILS, FC_READ_DISCRETE_INPUTS)
            else MAX_READ_REGISTERS
        )
        if count < 1 or count > max_count:
            raise IllegalDataValue(f"FC{fc}: count out of range: {count}")
        return ReadRequest(fc, address, count)
    if fc == FC_WRITE_SINGLE_COIL:
        if len(body) != 4:
            raise PduError(f"FC{fc}: expected 4 data bytes, got {len(body)}")
        address = int.from_bytes(body[0:2], "big")
        raw = int.from_bytes(body[2:4], "big")
        if raw not in (COIL_ON, COIL_OFF):
            raise IllegalDataValue(f"FC{fc}: invalid coil value: {raw:#06x}")
        return WriteSingleCoilRequest(fc, address, raw == COIL_ON)
    if fc == FC_WRITE_SINGLE_REGISTER:
        if len(body) != 4:
            raise PduError(f"FC{fc}: expected 4 data bytes, got {len(body)}")
        address = int.from_bytes(body[0:2], "big")
        value = int.from_bytes(body[2:4], "big")
        return WriteSingleRegisterRequest(fc, address, value)
    if fc == FC_GET_COMM_EVENT_COUNT:
        if body:
            raise PduError(f"FC{fc}: no data bytes expected, got {len(body)}")
        return GetCommEventCountRequest(fc)
    if fc == FC_WRITE_MULTIPLE_COILS:
        if len(body) < 6:
            raise PduError(f"FC{fc}: too short: {len(body)}")
        address = int.from_bytes(body[0:2], "big")
        count = int.from_bytes(body[2:4], "big")
        byte_count = body[4]
        if count < 1 or count > MAX_WRITE_BITS:
            raise IllegalDataValue(f"FC{fc}: count out of range: {count}")
        if byte_count != (count + 7) // 8 or len(body) != 5 + byte_count:
            raise PduError(f"FC{fc}: byte count mismatch")
        values = [
            bool(body[5 + i // 8] & (1 << (i % 8))) for i in range(count)
        ]
        return WriteMultipleCoilsRequest(fc, address, tuple(values))
    if fc == FC_WRITE_MULTIPLE_REGISTERS:
        if len(body) < 6:
            raise PduError(f"FC{fc}: too short: {len(body)}")
        address = int.from_bytes(body[0:2], "big")
        count = int.from_bytes(body[2:4], "big")
        byte_count = body[4]
        if count < 1 or count > MAX_WRITE_REGISTERS:
            raise IllegalDataValue(f"FC{fc}: count out of range: {count}")
        if byte_count != 2 * count or len(body) != 5 + byte_count:
            raise PduError(f"FC{fc}: byte count mismatch")
        values = tuple(
            int.from_bytes(body[5 + 2 * i:7 + 2 * i], "big") for i in range(count)
        )
        return WriteMultipleRegistersRequest(fc, address, values)
    if fc == FC_MASK_WRITE_REGISTER:
        if len(body) != 6:
            raise PduError(f"FC{fc}: expected 6 data bytes, got {len(body)}")
        address = int.from_bytes(body[0:2], "big")
        and_mask = int.from_bytes(body[2:4], "big")
        or_mask = int.from_bytes(body[4:6], "big")
        return MaskWriteRegisterRequest(fc, address, and_mask, or_mask)
    if fc == FC_READ_WRITE_MULTIPLE_REGISTERS:
        if len(body) < 10:
            raise PduError(f"FC{fc}: too short: {len(body)}")
        read_address = int.from_bytes(body[0:2], "big")
        read_count = int.from_bytes(body[2:4], "big")
        write_address = int.from_bytes(body[4:6], "big")
        write_count = int.from_bytes(body[6:8], "big")
        byte_count = body[8]  # 1-byte field per Modbus spec
        if read_count < 1 or read_count > MAX_READ_REGISTERS:
            raise IllegalDataValue(f"FC{fc}: read count out of range: {read_count}")
        if write_count < 1 or write_count > MAX_WRITE_REGISTERS:
            raise IllegalDataValue(f"FC{fc}: write count out of range: {write_count}")
        if byte_count != 2 * write_count or len(body) != 9 + byte_count:
            raise PduError(f"FC{fc}: byte count mismatch")
        values = tuple(
            int.from_bytes(body[9 + 2 * i:11 + 2 * i], "big")
            for i in range(write_count)
        )
        return ReadWriteMultipleRegistersRequest(
            fc, read_address, read_count, write_address, values
        )
    raise UnsupportedFunction(fc)  # pragma: no cover


class UnsupportedFunction(PduError):
    """Function code not implemented by the 750-362."""

    def __init__(self, function_code: int) -> None:
        self.function_code = function_code
        super().__init__(f"unsupported function code: {function_code:#04x}")


class IllegalDataValue(PduError):
    """Field value outside allowed range (exception 0x03 path)."""