pyMC_Repeater/repeater/handler_helpers/protocol_request.py

"""
Protocol request (REQ) handling helper for pyMC Repeater.

Provides repeater-specific callbacks for status and telemetry requests.
"""

import asyncio
import logging
import struct
import time

from pymc_core.node.handlers.protocol_request import (
    REQ_TYPE_GET_ACCESS_LIST,
    REQ_TYPE_GET_NEIGHBOURS,
    REQ_TYPE_GET_STATUS,
    REQ_TYPE_GET_TELEMETRY_DATA,
    SERVER_RESPONSE_DELAY_MS,
    ProtocolRequestHandler,
)

logger = logging.getLogger("ProtocolRequestHelper")


class ProtocolRequestHelper:
    """Provides repeater-specific protocol request handlers."""

    def __init__(
        self,
        identity_manager,
        packet_injector=None,
        acl_dict=None,
        radio=None,
        engine=None,
        neighbor_tracker=None,
    ):

        self.identity_manager = identity_manager
        self.packet_injector = packet_injector
        self.acl_dict = acl_dict or {}
        self.radio = radio
        self.engine = engine
        self.neighbor_tracker = neighbor_tracker
        
        # Dictionary of core handlers keyed by dest_hash
        self.handlers = {}
        
    def register_identity(self, name: str, identity, identity_type: str = "repeater"):

        hash_byte = identity.get_public_key()[0]
        
        # Get ACL for this identity
        identity_acl = self.acl_dict.get(hash_byte)
        if not identity_acl:
            logger.warning(f"Cannot register identity '{name}': no ACL for hash 0x{hash_byte:02X}")
            return
        
        # Create ACL contacts wrapper
        acl_contacts = self._create_acl_contacts_wrapper(identity_acl)
        
        # Build request handlers dict
        request_handlers = {
            REQ_TYPE_GET_STATUS: self._handle_get_status,
        }
        
        # Create core handler
        handler = ProtocolRequestHandler(
            local_identity=identity,
            contacts=acl_contacts,
            get_client_fn=lambda src_hash: self._get_client_from_acl(identity_acl, src_hash),
            request_handlers=request_handlers,
            log_fn=logger.info,
        )
        
        self.handlers[hash_byte] = {
            "handler": handler,
            "identity": identity,
            "name": name,
            "type": identity_type,
        }
        
        logger.info(f"Registered protocol request handler for '{name}': hash=0x{hash_byte:02X}")

    def _create_acl_contacts_wrapper(self, acl):
        """Create contacts wrapper from ACL."""

        class ACLContactsWrapper:
            def __init__(self, identity_acl):
                self._acl = identity_acl
            
            @property
            def contacts(self):
                return self._acl.get_all_clients()
        
        return ACLContactsWrapper(acl)
    
    def _get_client_from_acl(self, acl, src_hash: int):
        """Get client from ACL by source hash."""
        for client_info in acl.get_all_clients():
            if client_info.id.get_public_key()[0] == src_hash:
                return client_info
        return None
    
    async def process_request_packet(self, packet):

        try:
            if len(packet.payload) < 2:
                return False
            
            dest_hash = packet.payload[0]
            
            handler_info = self.handlers.get(dest_hash)
            if not handler_info:
                return False
            
            # Let core handler build response
            response_packet = await handler_info["handler"](packet)
            
            # Send response after delay
            if response_packet and self.packet_injector:
                await asyncio.sleep(SERVER_RESPONSE_DELAY_MS / 1000.0)
                await self.packet_injector(response_packet, wait_for_ack=False)
            
            packet.mark_do_not_retransmit()
            return True
            
        except Exception as e:
            logger.error(f"Error processing protocol request: {e}", exc_info=True)
            return False
    
    def _handle_get_status(self, client, timestamp: int, req_data: bytes):

        # C++ struct RepeaterStats (44 bytes total):
        # uint16_t batt_milli_volts;
        # uint16_t curr_tx_queue_len;
        # int16_t noise_floor;
        # int16_t last_rssi;
        # uint32_t n_packets_recv;
        # uint32_t n_packets_sent;
        # uint32_t total_air_time_secs;
        # uint32_t total_up_time_secs;
        # uint32_t n_sent_flood;
        # uint32_t n_sent_direct;
        # uint32_t n_recv_flood;
        # uint32_t n_recv_direct;
        # uint32_t err_events;
        # int16_t last_snr;
        # uint32_t n_direct_dups;
        # uint32_t n_flood_dups;
        # uint32_t total_rx_air_time_secs;

        # Get stats from radio/engine
        noise_floor = int(self.radio.get_noise_floor() * 1.0) if self.radio else -120
        last_rssi = (
            int(self.radio.last_rssi) if self.radio and hasattr(self.radio, "last_rssi") else -120
        )
        last_snr = int(
            (self.radio.last_snr * 4.0) if self.radio and hasattr(self.radio, "last_snr") else 0
        )

        # Get packet counts
        n_packets_recv = (
            self.radio.packets_received
            if self.radio and hasattr(self.radio, "packets_received")
            else 0
        )
        n_packets_sent = (
            self.radio.packets_sent if self.radio and hasattr(self.radio, "packets_sent") else 0
        )

        # Get airtime stats
        total_air_time_secs = 0
        total_rx_air_time_secs = 0
        if self.engine and hasattr(self.engine, "airtime_manager"):
            total_air_time_secs = int(self.engine.airtime_manager.total_tx_airtime_ms / 1000)

        # Get routing stats
        n_sent_flood = 0
        n_sent_direct = 0
        n_recv_flood = 0
        n_recv_direct = 0
        n_direct_dups = 0
        n_flood_dups = 0

        if self.engine:
            n_sent_flood = getattr(self.engine, "sent_flood_count", 0)
            n_sent_direct = getattr(self.engine, "sent_direct_count", 0)
            n_recv_flood = getattr(self.engine, "recv_flood_count", 0)
            n_recv_direct = getattr(self.engine, "recv_direct_count", 0)
            n_direct_dups = getattr(self.engine, "direct_dup_count", 0)
            n_flood_dups = getattr(self.engine, "flood_dup_count", 0)

        # Pack struct (little-endian)
        stats = struct.pack(
            "<HHhhIIIIIIIIIhIII",
            0,  # batt_milli_volts (not available on Pi)
            0,  # curr_tx_queue_len (TODO)
            noise_floor,
            last_rssi,
            n_packets_recv,
            n_packets_sent,
            total_air_time_secs,
            int(time.time()),  # total_up_time_secs
            n_sent_flood,
            n_sent_direct,
            n_recv_flood,
            n_recv_direct,
            0,  # err_events
            last_snr,
            n_direct_dups,
            n_flood_dups,
            total_rx_air_time_secs,
        )
        
        logger.debug(f"GET_STATUS: noise={noise_floor}dBm, rssi={last_rssi}dBm, snr={last_snr/4}dB")
        
        return stats