import ctypes
import struct
import logging
from ..lib1553.headers import UDP1553Header, UDP1553MessageHeader, CommandWordUnion 
from ..lib1553 import messages as msgs

logger = logging.getLogger(__name__)

class MessageDispatcher:
    """
    Parses raw UDP packets into high-level 1553 Message objects.
    Supports multi-message frames as per UDP1553 protocol.
    """

    def __init__(self):
        # Registry: Map (Subaddress, IsTransmit) -> MessageClass
        self._registry = {}
        self._init_registry()

    def _init_registry(self):
        """
        Dynamically registers all message classes defined in lib1553.messages.
        Uses the SUBADDRESS and IS_TRANSMIT constants defined in each class.
        """
        message_classes = [
            msgs.MsgA1, msgs.MsgA2, msgs.MsgA3, msgs.MsgA4, 
            msgs.MsgA5, msgs.MsgA6, msgs.MsgA7, msgs.MsgA8,
            msgs.MsgB1, msgs.MsgB2, msgs.MsgB3, msgs.MsgB4, 
            msgs.MsgB5, msgs.MsgB6, msgs.MsgB7, msgs.MsgB8
        ]

        for cls in message_classes:
            key = (cls.SUBADDRESS, cls.IS_TRANSMIT)
            self._registry[key] = cls

    def parse_packet(self, raw_data):
        """
        Decodes a raw byte buffer containing UDP1553 frame with multiple messages.
        Returns: (udp_header, list_of_messages)
        """
        udp_header_size = ctypes.sizeof(UDP1553Header)
        msg_header_size = ctypes.sizeof(UDP1553MessageHeader)

        if len(raw_data) < udp_header_size:
            return None, []

        # 1. Parse UDP1553 Header
        udp_header = UDP1553Header.from_buffer_copy(raw_data[:udp_header_size])

        if udp_header.marker1553 != UDP1553Header.MARKER_1553:
            return None, []

        # 2. Parse all messages in the frame
        messages = []
        offset = udp_header_size
        
        while offset < len(raw_data) - 1:
            # Check for end marker (0x5315)
            marker = struct.unpack_from('<H', raw_data, offset)[0]
            if marker == UDP1553Header.MARKER_END_1553:
                break
                
            # Check for message begin marker (0x3C3C)
            if marker != UDP1553MessageHeader.MARKER_BEGIN:
                # Skip unknown data
                offset += 2
                continue
            
            # Parse message header
            if offset + msg_header_size > len(raw_data):
                break
                
            msg_header = UDP1553MessageHeader.from_buffer_copy(
                raw_data[offset:offset + msg_header_size]
            )
            offset += msg_header_size
            
            # Extract Command Word info
            cw = msg_header.command_word.struct
            subaddress = cw.subaddress
            is_transmit = (cw.tr_bit == 1)
            word_count = cw.word_count if cw.word_count > 0 else 32
            
            # Mode codes (SA=0 or SA=31) have 1 data word
            if subaddress == 0 or subaddress == 31:
                word_count = 1
            
            # Calculate data size (only if TX message - RT sends data)
            data_size = word_count * 2 if is_transmit else 0
            
            # For TX messages, data comes AFTER header, before ~CW and end marker
            if is_transmit and offset + data_size <= len(raw_data):
                data_bytes = raw_data[offset:offset + data_size]
                offset += data_size
            else:
                data_bytes = b''
            
            # Skip ~CW (inverted command word) and end marker (0x3E3E)
            offset += 4  # 2 bytes ~CW + 2 bytes end marker
            
            # Find message class
            key = (subaddress, is_transmit)
            msg_class = self._registry.get(key)
            
            if msg_class and data_bytes:
                try:
                    message_obj = msg_class(data_bytes)
                    messages.append(message_obj)
                except Exception as e:
                    logger.error(f"Error parsing {msg_class.__name__}: {e}")

        return udp_header, messages
    
    def parse_packet_single(self, raw_data):
        """
        Legacy method - returns first message only for backward compatibility.
        """
        udp_header, messages = self.parse_packet(raw_data)
        if messages:
            return udp_header, messages[0]
        return udp_header, None