snakepit

#!/usr/bin/env python3
"""
Example of creating a custom bridge by extending the snakepit_bridge package.

This demonstrates how to create your own adapter without modifying
the core bridge logic, using proper package imports instead of
fragile sys.path manipulation.

This version showcases:
- Proper relative imports from the snakepit_bridge package
- Extending the generic command set with custom commands
- Production-ready error handling and graceful shutdown
"""

import sys
import os
import hashlib
import base64
from datetime import datetime

# Add the bridge package to Python path if not already installed
if __name__ == "__main__":
    bridge_dir = os.path.dirname(os.path.abspath(__file__))
    if bridge_dir not in sys.path:
        sys.path.insert(0, bridge_dir)

from snakepit_bridge import BaseCommandHandler, ProtocolHandler
from snakepit_bridge.adapters.generic import GenericCommandHandler
from snakepit_bridge.core import setup_graceful_shutdown, setup_broken_pipe_suppression


class CustomCommandHandler(BaseCommandHandler):
    """
    Example custom command handler that adds new functionality
    while keeping all the existing generic commands.
    """
    
    def __init__(self):
        super().__init__()
        # Custom initialization - could load ML models, connect to DBs, etc.
        self.custom_state = {"initialized_at": datetime.now().isoformat()}
    
    def _register_commands(self):
        """Register both inherited and new commands."""
        # First, register the generic commands by creating a temporary instance
        generic = GenericCommandHandler()
        
        # Copy all generic command registrations
        for cmd, handler in generic._command_registry.items():
            # Bind the handler methods to self to maintain state
            self.register_command(cmd, getattr(self, handler.__name__))
        
        # Add our custom commands
        self.register_command("hash", self.handle_hash)
        self.register_command("encode", self.handle_encode)
        self.register_command("custom_info", self.handle_custom_info)
    
    # Include all the generic command handlers
    def handle_ping(self, args):
        """Enhanced ping with custom data."""
        result = {
            "status": "ok",
            "bridge_type": "custom",  # Changed from "generic"
            "uptime": self.get_uptime(),
            "requests_handled": self.request_count,
            "timestamp": datetime.now().timestamp(),
            "python_version": sys.version,
            "worker_id": args.get("worker_id", "unknown"),
            "echo": args,
            "custom_data": {
                "handler_type": "CustomCommandHandler",
                "initialized_at": self.custom_state["initialized_at"]
            }
        }
        return result
    
    def handle_echo(self, args):
        """Echo command - unchanged from generic."""
        return {
            "status": "ok",
            "echoed": args,
            "timestamp": datetime.now().timestamp()
        }
    
    def handle_compute(self, args):
        """Compute command - unchanged from generic."""
        try:
            operation = args.get("operation", "add")
            a = args.get("a", 0)
            b = args.get("b", 0)
            
            if operation == "add":
                result = a + b
            elif operation == "subtract":
                result = a - b
            elif operation == "multiply":
                result = a * b
            elif operation == "divide":
                if b == 0:
                    raise ValueError("Division by zero")
                result = a / b
            else:
                raise ValueError(f"Unsupported operation: {operation}")
            
            return {
                "status": "ok",
                "operation": operation,
                "inputs": {"a": a, "b": b},
                "result": result,
                "timestamp": datetime.now().timestamp()
            }
        except Exception as e:
            return {
                "status": "error", 
                "error": str(e),
                "timestamp": datetime.now().timestamp()
            }
    
    def handle_info(self, args):
        """Enhanced info command."""
        return {
            "status": "ok",
            "bridge_info": {
                "name": "Custom Snakepit Bridge Example",
                "version": "2.0.0",
                "base_version": "2.0.0",
                "supported_commands": self.get_supported_commands(),
                "uptime": self.get_uptime(),
                "total_requests": self.request_count
            },
            "system_info": {
                "python_version": sys.version,
                "platform": sys.platform
            },
            "custom_info": self.custom_state,
            "timestamp": datetime.now().timestamp()
        }
    
    # New custom commands
    def handle_hash(self, args):
        """Generate hash of input text using specified algorithm."""
        text = args.get("text", "")
        algorithm = args.get("algorithm", "sha256")
        
        try:
            if algorithm == "sha256":
                hash_obj = hashlib.sha256(text.encode())
            elif algorithm == "md5":
                hash_obj = hashlib.md5(text.encode())
            elif algorithm == "sha1":
                hash_obj = hashlib.sha1(text.encode())
            else:
                raise ValueError(f"Unsupported algorithm: {algorithm}")
            
            return {
                "status": "ok",
                "algorithm": algorithm,
                "input": text,
                "hash": hash_obj.hexdigest(),
                "timestamp": datetime.now().timestamp()
            }
        except Exception as e:
            return {
                "status": "error",
                "error": str(e),
                "timestamp": datetime.now().timestamp()
            }
    
    def handle_encode(self, args):
        """Encode/decode text using base64."""
        text = args.get("text", "")
        operation = args.get("operation", "encode")
        
        try:
            if operation == "encode":
                result = base64.b64encode(text.encode()).decode()
            elif operation == "decode":
                result = base64.b64decode(text.encode()).decode()
            else:
                raise ValueError(f"Unsupported operation: {operation}")
            
            return {
                "status": "ok",
                "operation": operation,
                "input": text,
                "result": result,
                "timestamp": datetime.now().timestamp()
            }
        except Exception as e:
            return {
                "status": "error",
                "error": str(e),
                "timestamp": datetime.now().timestamp()
            }
    
    def handle_custom_info(self, args):
        """Demonstrate accessing custom state."""
        return {
            "status": "ok",
            "custom_state": self.custom_state,
            "handler_class": self.__class__.__name__,
            "additional_commands": [
                cmd for cmd in self.get_supported_commands() 
                if cmd not in ["ping", "echo", "compute", "info"]
            ],
            "timestamp": datetime.now().timestamp()
        }
    
    def get_uptime(self):
        """Helper method to calculate uptime."""
        return datetime.now().timestamp() - self.start_time


def main():
    """Main entry point for the custom bridge."""
    
    # Suppress broken pipe errors globally
    setup_broken_pipe_suppression()
    
    if len(sys.argv) > 1 and sys.argv[1] == "--help":
        print("Custom Snakepit Bridge Example V2")
        print("Usage: python example_custom_bridge_v2.py [--mode pool-worker]")
        print("")
        print("This example shows how to extend the generic bridge with custom commands.")
        print("Uses proper package imports instead of fragile sys.path manipulation.")
        print("")
        print("Supported commands:")
        handler = CustomCommandHandler()
        for cmd in sorted(handler.get_supported_commands()):
            print(f"  {cmd}")
        return
    
    # Create and run the custom bridge
    command_handler = CustomCommandHandler()
    protocol_handler = ProtocolHandler(command_handler)
    
    # Set up graceful shutdown handling
    setup_graceful_shutdown(protocol_handler)
    
    try:
        protocol_handler.run()
    except (KeyboardInterrupt, BrokenPipeError, IOError):
        os._exit(0)
    except Exception as e:
        from snakepit_bridge.core import safe_print
        safe_print(f"Bridge error: {e}")
        sys.exit(1)


if __name__ == "__main__":
    main()