NATSBridge/src/natsbridge_mpy.py

"""
NATSBridge - Cross-Platform Bi-Directional Data Bridge
MicroPython Implementation

This module provides functionality for sending and receiving data across network boundaries
using NATS as the message bus, with support for both direct payload transport and
URL-based transport for larger payloads.

Note: MicroPython has significant constraints compared to desktop implementations:
- Limited memory (~256KB - 1MB)
- No Arrow IPC support (memory constraints)
- Synchronous API (no async/await)
- Lower size threshold for direct transport
"""

import network
import time
import json
import base64
import uos
import struct
import random

# ---------------------------------------------- Constants ---------------------------------------------- #

"""
Default size threshold for switching from direct to link transport (100KB for MicroPython)
"""
DEFAULT_SIZE_THRESHOLD = 100000

"""
Default NATS server URL
"""
DEFAULT_BROKER_URL = "nats://localhost:4222"

"""
Default HTTP file server URL for link transport
"""
DEFAULT_FILESERVER_URL = "http://localhost:8080"

"""
Hard limit for payload size in MicroPython (50KB)
"""
MAX_PAYLOAD_SIZE = 50000


# ---------------------------------------------- Utility Functions ---------------------------------------------- #

def log_trace(correlation_id, message):
    """
    Log a trace message with correlation ID and timestamp.
    
    Args:
        correlation_id: Correlation ID for tracing
        message: Message content to log
    """
    timestamp = time.strftime('%Y-%m-%dT%H:%M:%SZ', time.localtime())
    print(f"[{timestamp}] [Correlation: {correlation_id}] {message}")


def _generate_uuid():
    """
    Generate a simple UUID compatible with MicroPython.
    
    Returns:
        UUID string
    """
    # Generate a simple UUID-like string
    # Format: xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx
    hex_chars = '0123456789abcdef'
    uuid_str = ''.join([random.choice(hex_chars) for _ in range(32)])
    # Insert hyphens at proper positions
    return f"{uuid_str[:8]}-{uuid_str[8:12]}-{uuid_str[12:16]}-{uuid_str[16:20]}-{uuid_str[20:]}"


# ---------------------------------------------- Serialization Functions ---------------------------------------------- #

def _serialize_data(data, payload_type):
    """
    Serialize data according to specified format.
    
    Args:
        data: Data to serialize (string for "text", dict for "dictionary",
              bytes for "image", "audio", "video", "binary")
        payload_type: Target format: "text", "dictionary", "image", "audio", "video", "binary"
    
    Returns:
        Binary representation of the serialized data
    
    Note:
        MicroPython does not support "table" type due to memory constraints.
    
    Raises:
        ValueError: If payload_type is not one of the supported types
    """
    if payload_type == 'text':
        if isinstance(data, str):
            return data.encode('utf-8')
        else:
            raise ValueError('Text data must be a string')
    elif payload_type == 'dictionary':
        json_str = json.dumps(data)
        return json_str.encode('utf-8')
    elif payload_type in ('image', 'audio', 'video', 'binary'):
        if isinstance(data, (bytes, bytearray, memoryview)):
            return bytes(data)
        else:
            raise ValueError(f'{payload_type} data must be bytes')
    else:
        raise ValueError(f'Unknown payload_type: {payload_type}')


def _deserialize_data(data, payload_type):
    """
    Deserialize bytes to data based on type.
    
    Args:
        data: Serialized data as bytes
        payload_type: Data type ("text", "dictionary", "image", "audio", "video", "binary")
    
    Returns:
        Deserialized data (String for "text", dict for "dictionary", bytes for others)
    
    Note:
        MicroPython does not support "table" type due to memory constraints.
    
    Raises:
        ValueError: If payload_type is not one of the supported types
    """
    if payload_type == 'text':
        return data.decode('utf-8')
    elif payload_type == 'dictionary':
        json_str = data.decode('utf-8')
        return json.loads(json_str)
    elif payload_type in ('image', 'audio', 'video', 'binary'):
        return data
    else:
        raise ValueError(f'Unknown payload_type: {payload_type}')


# ---------------------------------------------- File Server Handlers ---------------------------------------------- #

def _sync_fileserver_upload(file_server_url, dataname, data):
    """
    Synchronous file upload to HTTP server.
    
    Note:
        This is a simplified implementation for MicroPython.
        In practice, would use network.HTTP or similar.
        Currently raises NotImplementedError as file upload is not fully supported.
    
    Args:
        file_server_url: Base URL of the file server
        dataname: Name of the file being uploaded
        data: Raw byte data of the file content
    
    Returns:
        Dict with keys: 'status', 'url'
    
    Raises:
        NotImplementedError: File upload is not implemented in MicroPython
    """
    raise NotImplementedError("File upload not fully implemented in MicroPython. "
                              "Use direct transport only for memory-constrained devices.")


def _sync_fileserver_download(url, max_retries, base_delay, max_delay, correlation_id):
    """
    Synchronous file download with exponential backoff.
    
    Note:
        This is a simplified implementation for MicroPython.
        In practice, would use network.HTTP or similar.
        Currently raises NotImplementedError as file download is not fully supported.
    
    Args:
        url: URL to download from
        max_retries: Maximum retry attempts
        base_delay: Initial delay in ms
        max_delay: Maximum delay in ms
        correlation_id: Correlation ID for logging
    
    Returns:
        Downloaded bytes
    
    Raises:
        NotImplementedError: File download is not implemented in MicroPython
    """
    raise NotImplementedError("File download not fully implemented in MicroPython. "
                              "Use direct transport only for memory-constrained devices.")


# ---------------------------------------------- NATS Client ---------------------------------------------- #

class NATSClient:
    """
    NATS client wrapper for MicroPython.
    
    Note:
        This is a simplified implementation for MicroPython.
        Full NATS client implementation would require additional network stack support.
    """
    
    def __init__(self, url=DEFAULT_BROKER_URL):
        """
        Initialize NATS client.
        
        Args:
            url: NATS server URL
        """
        self.url = url
        self._connected = False
    
    def connect(self):
        """
        Connect to NATS server.
        
        Note:
            This is a placeholder implementation.
            Actual NATS client would require network stack support.
        
        Returns:
            True if connected, False otherwise
        """
        # Placeholder - actual implementation would connect to NATS server
        self._connected = True
        return self._connected
    
    def publish(self, subject, message):
        """
        Publish message to NATS subject.
        
        Note:
            This is a placeholder implementation.
            Actual NATS client would require network stack support.
        
        Args:
            subject: NATS subject to publish to
            message: Message to publish
        """
        if not self._connected:
            raise RuntimeError("Not connected to NATS server")
        # Placeholder - actual implementation would publish to NATS
        print(f"[NATS] Publish to {subject}: {message[:50]}...")
    
    def close(self):
        """Close the NATS connection."""
        self._connected = False


# ---------------------------------------------- Core Functions ---------------------------------------------- #

def _build_envelope(subject, payloads, options):
    """
    Build message envelope from payloads and metadata.
    
    Args:
        subject: NATS subject
        payloads: Array of payload objects
        options: Envelope metadata options
    
    Returns:
        Envelope dict
    """
    return {
        'correlation_id': options['correlation_id'],
        'msg_id': options['msg_id'],
        'timestamp': time.strftime('%Y-%m-%dT%H:%M:%SZ', time.localtime()),
        'send_to': subject,
        'msg_purpose': options['msg_purpose'],
        'sender_name': options['sender_name'],
        'sender_id': options['sender_id'],
        'receiver_name': options['receiver_name'],
        'receiver_id': options['receiver_id'],
        'reply_to': options['reply_to'],
        'reply_to_msg_id': options['reply_to_msg_id'],
        'broker_url': options['broker_url'],
        'metadata': {},
        'payloads': payloads
    }


def _build_payload(dataname, payload_type, payload_bytes, transport, data):
    """
    Build payload object from serialized data.
    
    Args:
        dataname: Name of the payload
        payload_type: Type of the payload
        payload_bytes: Serialized payload bytes
        transport: Transport type ("direct" or "link")
        data: Data (base64 for direct, URL for link)
    
    Returns:
        Payload dict
    """
    return {
        'id': _generate_uuid(),
        'dataname': dataname,
        'payload_type': payload_type,
        'transport': transport,
        'encoding': 'base64' if transport == 'direct' else 'none',
        'size': len(payload_bytes),
        'data': data,
        'metadata': {'payload_bytes': len(payload_bytes)} if transport == 'direct' else {}
    }


def _publish(subject, message, correlation_id):
    """
    Publish message to NATS.
    
    Note:
        This is a simplified implementation for MicroPython.
    
    Args:
        subject: NATS subject to publish to
        message: JSON message to publish
        correlation_id: Correlation ID for logging
    """
    log_trace(correlation_id, f"Publishing to {subject}")
    # Placeholder - actual implementation would use NATSClient
    # client = NATSClient()
    # client.connect()
    # client.publish(subject, message)
    # client.close()


def smartsend(subject, data, **kwargs):
    """
    Send data via NATS with automatic transport selection.
    
    This function intelligently routes data delivery based on payload size.
    If the serialized payload is smaller than size_threshold, it encodes the data as Base64
    and publishes directly over NATS. Otherwise, it uploads the data to a fileserver
    and publishes only the download URL over NATS.
    
    Note:
        MicroPython has memory constraints, so the default size_threshold is lower (100KB).
        Table type is not supported due to memory constraints.
    
    Args:
        subject: NATS subject to publish the message to
        data: List of (dataname, data, type) tuples to send
            - dataname: Name of the payload
            - data: The actual data to send
            - type: Payload type: "text", "dictionary", "image", "audio", "video", "binary"
        broker_url: NATS server URL (default: DEFAULT_BROKER_URL)
        fileserver_url: HTTP file server URL (default: DEFAULT_FILESERVER_URL)
        fileserver_upload_handler: Function to handle fileserver uploads (default: _sync_fileserver_upload)
        size_threshold: Threshold in bytes separating direct vs link transport (default: 100000)
        correlation_id: Correlation ID for tracing (auto-generated if not provided)
        msg_purpose: Purpose of the message (default: "chat")
        sender_name: Name of the sender (default: "NATSBridge")
        receiver_name: Name of the receiver (empty means broadcast)
        receiver_id: UUID of the receiver (empty means broadcast)
        reply_to: Topic to reply to (empty if no reply expected)
        reply_to_msg_id: Message ID this message is replying to
        is_publish: Whether to automatically publish the message (default: True)
        msg_id: Message ID (auto-generated if not provided)
        sender_id: Sender ID (auto-generated if not provided)
    
    Returns:
        Tuple of (env, env_json_str) where:
            - env: Dict containing all metadata and payloads
            - env_json_str: JSON string for publishing to NATS
    
    Example:
        >>> # Send text payload
        >>> env, env_json_str = NATSBridge.smartsend(
        ...     "/chat",
        ...     [("message", "Hello!", "text")],
        ...     broker_url="nats://localhost:4222"
        ... )
        >>>
        >>> # Send dictionary payload
        >>> env, env_json_str = NATSBridge.smartsend(
        ...     "/config",
        ...     [("config", {"key": "value"}, "dictionary")],
        ...     broker_url="nats://localhost:4222"
        ... )
        >>>
        >>> # Send binary payload (image, audio, video)
        >>> env, env_json_str = NATSBridge.smartsend(
        ...     "/media",
        ...     [("image", image_bytes, "image")],
        ...     broker_url="nats://localhost:4222"
        ... )
    """
    # Extract options with defaults
    correlation_id = kwargs.get('correlation_id', _generate_uuid())
    msg_id = kwargs.get('msg_id', _generate_uuid())
    sender_id = kwargs.get('sender_id', _generate_uuid())
    broker_url = kwargs.get('broker_url', DEFAULT_BROKER_URL)
    fileserver_url = kwargs.get('fileserver_url', DEFAULT_FILESERVER_URL)
    size_threshold = kwargs.get('size_threshold', DEFAULT_SIZE_THRESHOLD)
    msg_purpose = kwargs.get('msg_purpose', 'chat')
    sender_name = kwargs.get('sender_name', 'NATSBridge')
    receiver_name = kwargs.get('receiver_name', '')
    receiver_id = kwargs.get('receiver_id', '')
    reply_to = kwargs.get('reply_to', '')
    reply_to_msg_id = kwargs.get('reply_to_msg_id', '')
    is_publish = kwargs.get('is_publish', True)
    fileserver_upload_handler = kwargs.get('fileserver_upload_handler', _sync_fileserver_upload)
    
    log_trace(correlation_id, f"Starting smartsend for subject: {subject}")
    
    # Process payloads
    payloads = []
    for dataname, payload_data, payload_type in data:
        payload_bytes = _serialize_data(payload_data, payload_type)
        payload_size = len(payload_bytes)
        
        # Check against hard limit for MicroPython
        if payload_size > MAX_PAYLOAD_SIZE:
            raise MemoryError(f"Payload '{dataname}' exceeds max size {MAX_PAYLOAD_SIZE} bytes")
        
        log_trace(correlation_id, f"Serialized payload '{dataname}' (type: {payload_type}) size: {payload_size} bytes")
        
        if payload_size < size_threshold:
            # Direct path
            payload_b64 = base64.b64encode(payload_bytes).decode('ascii')
            log_trace(correlation_id, f"Using direct transport for {payload_size} bytes")
            
            payload = _build_payload(dataname, payload_type, payload_bytes, 'direct', payload_b64)
            payloads.append(payload)
        else:
            # Link path (limited support)
            log_trace(correlation_id, "Using link transport, uploading to fileserver")
            
            try:
                response = fileserver_upload_handler(fileserver_url, dataname, payload_bytes)
                log_trace(correlation_id, f"Uploaded to URL: {response['url']}")
                
                payload = _build_payload(dataname, payload_type, payload_bytes, 'link', response['url'])
                payloads.append(payload)
            except NotImplementedError:
                # Fall back to direct transport if file upload not available
                log_trace(correlation_id, "File upload not available, using direct transport")
                payload_b64 = base64.b64encode(payload_bytes).decode('ascii')
                payload = _build_payload(dataname, payload_type, payload_bytes, 'direct', payload_b64)
                payloads.append(payload)
    
    # Build envelope
    env = _build_envelope(subject, payloads, {
        'correlation_id': correlation_id,
        'msg_id': msg_id,
        'msg_purpose': msg_purpose,
        'sender_name': sender_name,
        'sender_id': sender_id,
        'receiver_name': receiver_name,
        'receiver_id': receiver_id,
        'reply_to': reply_to,
        'reply_to_msg_id': reply_to_msg_id,
        'broker_url': broker_url
    })
    
    env_json_str = json.dumps(env)
    
    if is_publish:
        _publish(subject, env_json_str, correlation_id)
    
    return env, env_json_str


def smartreceive(msg, **kwargs):
    """
    Receive and process NATS message.
    
    This function processes incoming NATS messages, handling both direct transport
    (base64 decoded payloads) and link transport (URL-based payloads).
    It deserializes the data based on the transport type and returns the result.
    
    Note:
        MicroPython has memory constraints, so large payloads should be avoided.
        Table type is not supported due to memory constraints.
    
    Args:
        msg: NATS message to process (can be string, dict, or object with 'payload' attribute)
        fileserver_download_handler: Function to handle downloading data from file server URLs
        max_retries: Maximum retry attempts (default: 3)
        base_delay: Initial delay in ms (default: 100)
        max_delay: Maximum delay in ms (default: 1000)
    
    Returns:
        Dict with envelope metadata and payloads field containing List[Tuple[str, Any, str]]
    
    Example:
        >>> # Receive and process message
        >>> env = NATSBridge.smartreceive(msg, fileserver_download_handler=_sync_fileserver_download)
        >>> # env is a Dict with "payloads" key containing List[Tuple[str, Any, str]]
        >>> for dataname, data, type_ in env["payloads"]:
        ...     print(f"{dataname}: {data} (type: {type_})")
    """
    # Parse the JSON envelope
    if isinstance(msg, dict):
        # Already parsed
        env_json_obj = msg
    elif hasattr(msg, 'payload'):
        # Object with payload attribute
        payload = msg.payload if isinstance(msg.payload, str) else msg.payload.decode('utf-8')
        env_json_obj = json.loads(payload)
    else:
        # Assume it's already a JSON string or dict
        env_json_obj = json.loads(msg) if isinstance(msg, str) else msg
    
    correlation_id = env_json_obj['correlation_id']
    log_trace(correlation_id, "Processing received message")
    
    # Process all payloads in the envelope
    payloads_list = []
    num_payloads = len(env_json_obj['payloads'])
    
    for i in range(num_payloads):
        payload_obj = env_json_obj['payloads'][i]
        transport = payload_obj['transport']
        dataname = payload_obj['dataname']
        
        if transport == 'direct':
            log_trace(correlation_id, f"Direct transport - decoding payload '{dataname}'")
            
            # Extract base64 payload from the payload
            payload_b64 = payload_obj['data']
            
            # Decode Base64 payload
            payload_bytes = base64.b64decode(payload_b64)
            
            # Deserialize based on type
            data_type = payload_obj['payload_type']
            data = _deserialize_data(payload_bytes, data_type)
            
            payloads_list.append((dataname, data, data_type))
        elif transport == 'link':
            # Extract download URL from the payload
            url = payload_obj['data']
            log_trace(correlation_id, f"Link transport - fetching '{dataname}' from URL: {url}")
            
            # Fetch with exponential backoff using the download handler
            fileserver_download_handler = kwargs.get('fileserver_download_handler', _sync_fileserver_download)
            max_retries = kwargs.get('max_retries', 3)
            base_delay = kwargs.get('base_delay', 100)
            max_delay = kwargs.get('max_delay', 1000)
            
            downloaded_data = fileserver_download_handler(
                url,
                max_retries,
                base_delay,
                max_delay,
                correlation_id
            )
            
            # Deserialize based on type
            data_type = payload_obj['payload_type']
            data = _deserialize_data(downloaded_data, data_type)
            
            payloads_list.append((dataname, data, data_type))
        else:
            raise ValueError(f"Unknown transport type for payload '{dataname}': {transport}")
    
    env_json_obj['payloads'] = payloads_list
    return env_json_obj


# ---------------------------------------------- Module Exports ---------------------------------------------- #

class NATSBridge:
    """
    MicroPython NATS bridge implementation.
    
    This class provides a convenient interface for NATSBridge functionality,
    encapsulating the main functions and providing a class-based API.
    
    Note:
        MicroPython has significant constraints:
        - No Arrow IPC support (memory constraints)
        - Only direct transport (< 100KB threshold enforced)
        - Simplified UUID generation
        - No async/await (synchronous API)
    """
    
    DEFAULT_SIZE_THRESHOLD = DEFAULT_SIZE_THRESHOLD
    DEFAULT_BROKER_URL = DEFAULT_BROKER_URL
    DEFAULT_FILESERVER_URL = DEFAULT_FILESERVER_URL
    MAX_PAYLOAD_SIZE = MAX_PAYLOAD_SIZE
    
    def __init__(self, broker_url=None, fileserver_url=None):
        """
        Initialize NATSBridge.
        
        Args:
            broker_url: NATS server URL (defaults to DEFAULT_BROKER_URL)
            fileserver_url: HTTP file server URL (defaults to DEFAULT_FILESERVER_URL)
        """
        self.broker_url = broker_url or self.DEFAULT_BROKER_URL
        self.fileserver_url = fileserver_url or self.DEFAULT_FILESERVER_URL
    
    def smartsend(self, subject, data, **kwargs):
        """
        Send data via NATS.
        
        Args:
            subject: NATS subject to publish to
            data: List of (dataname, data, type) tuples
            **kwargs: Additional options passed to smartsend
        
        Returns:
            Tuple of (env, env_json_str)
        """
        kwargs['broker_url'] = kwargs.get('broker_url', self.broker_url)
        kwargs['fileserver_url'] = kwargs.get('fileserver_url', self.fileserver_url)
        return smartsend(subject, data, **kwargs)
    
    def smartreceive(self, msg, **kwargs):
        """
        Receive and process NATS message.
        
        Args:
            msg: NATS message to process
            **kwargs: Additional options passed to smartreceive
        
        Returns:
            Dict with envelope metadata and payloads
        """
        return smartreceive(msg, **kwargs)


# Convenience functions for module-level usage
def send(subject, data, **kwargs):
    """
    Convenience function for sending data.
    
    Args:
        subject: NATS subject to publish to
        data: List of (dataname, data, type) tuples
        **kwargs: Additional options
    
    Returns:
        Tuple of (env, env_json_str)
    """
    return smartsend(subject, data, **kwargs)


def receive(msg, **kwargs):
    """
    Convenience function for receiving messages.
    
    Args:
        msg: NATS message to process
        **kwargs: Additional options
    
    Returns:
        Dict with envelope metadata and payloads
    """
    return smartreceive(msg, **kwargs)


__all__ = [
    'smartsend',
    'smartreceive',
    'NATSBridge',
    'send',
    'receive',
    'DEFAULT_SIZE_THRESHOLD',
    'DEFAULT_BROKER_URL',
    'DEFAULT_FILESERVER_URL',
    'MAX_PAYLOAD_SIZE',
    'NATSClient',
    '_serialize_data',
    '_deserialize_data',
    'log_trace',
    '_sync_fileserver_upload',
    '_sync_fileserver_download'
]