rw_variable.py

import struct
from dataclasses import dataclass
from typing import AsyncGenerator, ClassVar, Generator

from ..enums import (
    Area,
    DataSizeByte,
    DataTransportSize,
    DataType,
    DataTypeTransportSize,
    FunctionCode,
    ItemReturnCode,
    MessageType,
    ParameterTransportSize,
)
from .data_item import DataItem
from .exceptions import ReadVariableException, WriteVariableException
from .headers import S7AckDataHeader, S7Header
from .packet import S7Packet
from .variable_address import VariableAddress


class RequestParameterItem:
    """_summary_

    ITEM_HEAD:
        variable specification: 0x12
        Length of following address specification: 0x0a
        Syntax Id: S7ANY (0x10)

    Raises:
        ValueError: _description_
        ValueError: _description_
        ValueError: _description_

    Returns:
        _type_: _description_
    """

    ITEM_HEAD: ClassVar[bytes] = b"\x12\x0a\x10"
    LENGTH = 12

    def __init__(self, address: VariableAddress, transport_size: ParameterTransportSize | None = None):
        # word_size is required by the request_generator to calculate maximum elements per packet
        if transport_size is None:
            self.transport_size = ParameterTransportSize[address.data_type.value]
        else:
            self.transport_size = transport_size
        self.word_size = DataSizeByte[address.data_type.value]
        self.length: int = address.amount
        self.address = address

    def serialize(self) -> bytes:
        if self.address.area != Area.DB:
            self.address.db_number = 0
        packet = self.ITEM_HEAD + struct.pack(
            "!BHHB",
            self.transport_size,
            self.length,
            self.address.db_number,
            self.address.area,
        )
        if self.transport_size in [ParameterTransportSize.COUNTER, ParameterTransportSize.TIMER]:
            address = self.address.start
        else:
            address = (self.address.start << 3) + self.address.start_bit
        return bytes(packet + address.to_bytes(length=3, byteorder="big"))

    def validate(self) -> None:
        if self.address.area != Area.DB:
            self.address.db_number = 0

        if self.transport_size != ParameterTransportSize.BIT and self.address.start_bit != 0:
            raise ValueError(f"{self.transport_size.name} Invalid start_bit: {self.address.start_bit}. should be 0")

        if self.address.db_number < 0 or self.address.db_number > 65535 or self.address.start < 0 or self.length < 1:
            raise ValueError("Invalid parameters")

        if self.transport_size == ParameterTransportSize.BIT and self.length > 1:
            raise ValueError("Invalid transport size")

    @classmethod
    def parse(cls, packet: bytes) -> "RequestParameterItem":
        _, transport_size, amount, db_number, area, address = struct.unpack_from("!3sBHHB3s", packet)
        transport_size = ParameterTransportSize(transport_size)
        start_bit = 0
        start_address = int.from_bytes(address, byteorder="big")
        if transport_size in [ParameterTransportSize.COUNTER, ParameterTransportSize.TIMER]:
            start = start_address
        else:
            start = start_address >> 3
            start_bit = start_address & 0x000007
        data_type = DataType[transport_size.name]
        variable_address = VariableAddress(
            area=area,
            db_number=db_number,
            start=start,
            data_type=data_type,
            start_bit=start_bit,
            amount=amount,
        )
        return cls(address=variable_address, transport_size=transport_size)


class VariableRequestParameter:
    LENGTH = 2

    def __init__(self, function_code: FunctionCode, items: list[RequestParameterItem]):
        self.function_code = function_code
        self.items: list[RequestParameterItem] = items

    def serialize(self) -> bytes:
        parameter = struct.pack("!BB", self.function_code, len(self.items))
        for item in self.items:
            parameter += item.serialize()
        self.length = len(parameter)
        return parameter


class VariableReadRequest(S7Packet):
    MESSAGE_TYPE = MessageType.JobRequest
    FUNCTION_CODE = FunctionCode.ReadVariable

    def __init__(self, parameter: VariableRequestParameter) -> None:
        self.header = None
        self.parameter = parameter
        self.data = None

    @classmethod
    def create(cls, items: list[VariableAddress]) -> "VariableReadRequest":
        parameter_items = cls._create_parameter_item_list(items=items)
        parameter = VariableRequestParameter(function_code=cls.FUNCTION_CODE, items=parameter_items)
        return cls(parameter=parameter)

    @staticmethod
    def _create_parameter_item_list(items: list[VariableAddress]) -> list[RequestParameterItem]:
        parameter_items = []
        for item in items:
            parameter_item = RequestParameterItem(address=item)
            parameter_item.validate()
            parameter_items.append(parameter_item)
        return parameter_items

    def serialize(self) -> bytes:
        return self.parameter.serialize()

    def request_generator(self, pdu_length: int) -> Generator["VariableReadRequest", None, None]:
        """Generate multiple read requests that fit within the PDU length limit.

        Splits a large read request into smaller chunks based on the maximum
        payload size allowed by the negotiated PDU length.

        Args:
            pdu_length: Maximum PDU size negotiated during connection setup.

        Yields:
            VariableReadRequest: Individual requests sized to fit within pdu_length.
        """
        fixed_packet_part = S7AckDataHeader.LENGTH + VariableRequestParameter.LENGTH + DataItem.HEADER_LENGTH
        max_payload_length = pdu_length - fixed_packet_part
        parameter_item = self.parameter.items[0]
        max_elements = int(max_payload_length / parameter_item.word_size)
        total_elements = parameter_item.address.amount
        offset = parameter_item.address.start

        while total_elements > 0:
            number_of_elements = max_elements if total_elements > max_elements else total_elements
            parameter_item.address.start = offset
            parameter_item.length = number_of_elements

            request = VariableReadRequest(
                parameter=VariableRequestParameter(items=[parameter_item], function_code=self.FUNCTION_CODE)
            )
            yield request
            total_elements -= number_of_elements
            offset += number_of_elements * parameter_item.word_size

    async def async_request_generator(self, pdu_length: int) -> AsyncGenerator["VariableReadRequest", None]:
        for item in self.request_generator(pdu_length):
            yield item

    @classmethod
    def parse(cls, packet: bytes) -> "VariableReadRequest":
        function_code, item_count = struct.unpack_from("!BB", packet)
        if function_code != cls.FUNCTION_CODE.value:
            raise ValueError(f"Invalid function code, got: {function_code}, expected: {cls.FUNCTION_CODE.value}")
        offset = 2
        items = []
        for _ in range(item_count):
            item = RequestParameterItem.parse(packet[offset : offset + RequestParameterItem.LENGTH])
            items.append(item)
            offset += RequestParameterItem.LENGTH
        parameter = VariableRequestParameter(function_code=function_code, items=items)
        return cls(parameter=parameter)

    def serialize_parameter(self) -> bytes:
        return self.parameter.serialize()

    def serialize_data(self) -> bytes:
        return b""


class VariableWriteRequest(S7Packet):
    MESSAGE_TYPE = MessageType.JobRequest
    FUNCTION_CODE = FunctionCode.WriteVariable

    def __init__(self, parameter: VariableRequestParameter, data: list[DataItem]) -> None:
        self.header = None
        self.parameter = parameter
        self.data = data

    @classmethod
    def create(cls, items: list[tuple[VariableAddress, bytes]]) -> "VariableWriteRequest":
        parameter_items = []
        data_items = []
        for item, data in items:
            parameter_item = RequestParameterItem(address=item)
            parameter_item.validate()
            parameter_items.append(parameter_item)
            data_item = cls._create_data_item(parameter_item=parameter_item, data=data)
            data_items.append(data_item)
        parameter = VariableRequestParameter(function_code=cls.FUNCTION_CODE, items=parameter_items)
        return cls(parameter=parameter, data=data_items)

    @classmethod
    def _create_data_item(cls, parameter_item: RequestParameterItem, data: bytes) -> DataItem:
        data_transport_size = DataTypeTransportSize[parameter_item.transport_size.name].value
        amount = parameter_item.address.amount
        data_length = parameter_item.word_size * amount
        if parameter_item.transport_size not in [
            ParameterTransportSize.COUNTER,
            ParameterTransportSize.TIMER,
            ParameterTransportSize.REAL,
            ParameterTransportSize.CHAR,
            ParameterTransportSize.BIT,
        ]:
            data_length *= 8
        return DataItem(transport_size=data_transport_size, data_length=data_length, data=data)

    def request_generator(self, pdu_length: int) -> Generator["VariableWriteRequest", None, None]:
        """Generate multiple write requests that fit within the PDU length limit.

        Splits a large write request into smaller chunks based on the maximum
        payload size allowed by the negotiated PDU length.

        Args:
            pdu_length: Maximum PDU size negotiated during connection setup.

        Yields:
            VariableWriteRequest: Individual requests sized to fit within pdu_length.
        """
        max_payload_length = (
            pdu_length
            - S7Header.LENGTH
            - VariableRequestParameter.LENGTH
            - RequestParameterItem.LENGTH
            - DataItem.HEADER_LENGTH
        )
        parameter_item: RequestParameterItem = self.parameter.items[0]

        max_elements = max_payload_length // parameter_item.word_size
        total_elements = parameter_item.address.amount
        offset = parameter_item.address.start

        data_item = self.data[0]
        size_unit = parameter_item.word_size.value

        if parameter_item.transport_size not in [
            ParameterTransportSize.COUNTER,
            ParameterTransportSize.TIMER,
            ParameterTransportSize.REAL,
            ParameterTransportSize.CHAR,
            ParameterTransportSize.BIT,
        ]:
            size_unit *= 8
        data_offset = 0

        while total_elements > 0:
            number_of_elements = max_elements if total_elements > max_elements else total_elements
            parameter_item.address.start = offset
            parameter_item.length = number_of_elements
            data_length = number_of_elements * size_unit

            data = data_item.data[data_offset : data_offset + data_length]

            request = VariableWriteRequest(
                parameter=VariableRequestParameter(items=[parameter_item], function_code=self.FUNCTION_CODE),
                data=[DataItem(transport_size=data_item.transport_size, data_length=data_length, data=data)],
            )
            yield request
            total_elements -= number_of_elements
            offset += number_of_elements * parameter_item.word_size
            data_offset += number_of_elements * size_unit

    async def async_request_generator(self, pdu_length: int) -> AsyncGenerator["VariableWriteRequest", None]:
        for item in self.request_generator(pdu_length):
            yield item

    @classmethod
    def parse(cls, packet: bytes) -> "VariableWriteRequest":
        function_code, item_count = struct.unpack_from("!BB", packet)
        if function_code != cls.FUNCTION_CODE.value:
            raise ValueError(f"Invalid function code, got: {function_code}, expected: {cls.FUNCTION_CODE.value}")
        offset = 2
        parameter_items = []
        data_items = []
        for _ in range(item_count):
            item = RequestParameterItem.parse(packet[offset : offset + RequestParameterItem.LENGTH])
            parameter_items.append(item)
            offset += RequestParameterItem.LENGTH
        for _ in range(item_count):
            data_item = DataItem.parse(packet[offset:])
            offset += data_item.HEADER_LENGTH + data_item.data_length
            data_items.append(data_item)
        parameter = VariableRequestParameter(function_code=function_code, items=parameter_items)
        return cls(parameter=parameter, data=data_items)

    def serialize_parameter(self) -> bytes:
        return self.parameter.serialize()

    def serialize_data(self) -> bytes:
        return b"".join(data_item.serialize() for data_item in self.data)


@dataclass
class VariableParameter:
    function_code: FunctionCode
    item_count: int

    def serialize(self) -> bytes:
        return struct.pack("!BB", self.function_code, self.item_count)


class ReadVariableResponse(S7Packet):
    MESSAGE_TYPE = MessageType.Response

    def __init__(self, parameter: VariableParameter, data: list[DataItem]) -> None:
        self.header = None
        self.parameter = parameter
        self.data = data

    @classmethod
    def parse(cls, packet: bytes) -> "ReadVariableResponse":
        function_code, item_count = struct.unpack_from("!BB", packet)
        offset = 2
        items = []
        for _ in range(item_count):
            return_code, transport_size, data_length = struct.unpack_from("!BBH", packet, offset=offset)
            offset = offset + 4
            data_size = data_length
            if transport_size not in [
                DataTransportSize.OCTET,
                DataTransportSize.REAL,
                DataTransportSize.BIT,
            ]:
                data_size = int(data_length / 8)

            data = packet[offset : offset + data_size]
            item = DataItem(
                return_code=ItemReturnCode(return_code),
                transport_size=DataTransportSize(transport_size),
                data_length=data_size,
                data=data,
            )
            items.append(item)
            offset = offset + data_size
        parameter = VariableParameter(function_code=function_code, item_count=item_count)
        return cls(parameter=parameter, data=items)

    def values(self) -> list[bytes]:
        """Extract and return data values from the response.

        Returns:
            list[bytes]: List of raw data bytes from each response item.

        Raises:
            ReadVariableException: If any item in the response has a non-success return code.
        """
        result = []
        for item in self.data:
            if item.return_code != ItemReturnCode.SUCCESS:
                raise ReadVariableException(f"ReadVariableResponseItem return code: {item.return_code}", response=self)
            result.append(item.data)
        return result

    def serialize_parameter(self) -> bytes:
        return self.parameter.serialize()

    def serialize_data(self) -> bytes:
        return b"".join(item.serialize() for item in self.data)


class WriteVariableResponse(S7Packet):
    MESSAGE_TYPE = MessageType.Response

    def __init__(self, parameter: VariableParameter, data: list[ItemReturnCode]) -> None:
        self.header = None
        self.parameter = parameter
        self.data: list[ItemReturnCode] = data

    @classmethod
    def parse(cls, packet: bytes) -> "WriteVariableResponse":
        function_code, item_count = struct.unpack_from("!BB", packet)
        offset = 2
        parameter = VariableParameter(function_code=function_code, item_count=item_count)
        data: list[ItemReturnCode] = [ItemReturnCode(item) for item in packet[offset:item_count]]
        return cls(parameter=parameter, data=data)

    def serialize_parameter(self) -> bytes:
        return self.parameter.serialize()

    def serialize_data(self) -> bytes:
        return struct.pack(f"!{len(self.data)}B", *self.data)

    def check_result(self) -> None:
        """Check the result of each write operation.

        Raises:
            WriteVariableException: If any item in the response has a non-success return code.
        """
        for item in self.data:
            if item != ItemReturnCode.SUCCESS:
                raise WriteVariableException(f"WriteVariableResponseItem return code: {item}", response=self)
        return None