import stomp
import json
import logging
import time
from typing import Dict, Any, Optional
from contextlib import contextmanager

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class ActiveMQProducer:
    """Production-grade ActiveMQ producer with error handling."""
    
    def __init__(
        self,
        host: str,
        port: int = 61613,
        username: str = None,
        password: str = None,
        use_ssl: bool = True
    ):
        self.host = host
        self.port = port
        self.username = username
        self.password = password
        self.use_ssl = use_ssl
        self.conn = None
        self.connected = False
        
    def connect(self) -> bool:
        """Connect to ActiveMQ broker."""
        try:
            self.conn = stomp.Connection(
                host_and_ports=[(self.host, self.port)],
                use_ssl=self.use_ssl
            )
            
            if self.username and self.password:
                self.conn.set_user(self.username, self.password)
            
            self.conn.connect(
                wait=True,
                headers={'accept-version': '1.1', 'heart-beat': '10000,10000'}
            )
            self.connected = True
            logger.info(f"Connected to ActiveMQ at {self.host}:{self.port}")
            return True
            
        except Exception as e:
            logger.error(f"Failed to connect to ActiveMQ: {e}")
            self.connected = False
            return False
            
    def disconnect(self):
        """Disconnect from broker."""
        if self.conn and self.connected:
            try:
                self.conn.disconnect()
                self.connected = False
                logger.info("Disconnected from ActiveMQ")
            except Exception as e:
                logger.error(f"Error disconnecting: {e}")
                
    def send(
        self,
        destination: str,
        message: Dict[str, Any],
        headers: Dict[str, str] = None,
        priority: int = 4,
        time_to_live: int = 0
    ) -> bool:
        """Send a message to ActiveMQ."""
        try:
            if not self.connected:
                if not self.connect():
                    raise ConnectionError("Failed to connect to ActiveMQ")
            
            message_body = json.dumps(message)
            
            send_headers = {
                'content-type': 'application/json',
                'priority': str(priority)
            }
            
            if headers:
                send_headers.update(headers)
                
            if time_to_live > 0:
                send_headers['expires'] = str(
                    int(time.time() * 1000) + (time_to_live * 1000)
                )
            
            self.conn.send(
                destination=destination,
                body=message_body,
                headers=send_headers,
                ack='auto'
            )
            
            logger.info(f"Message sent to {destination}")
            return True
            
        except Exception as e:
            logger.error(f"Failed to send message: {e}")
            return False
            
    def send_batch(
        self,
        destination: str,
        messages: list,
        batch_size: int = 100
    ) -> int:
        """Send a batch of messages."""
        sent_count = 0
        
        for i in range(0, len(messages), batch_size):
            batch = messages[i:i + batch_size]
            
            for message in batch:
                if self.send(destination, message):
                    sent_count += 1
                    
            logger.info(f"Batch {i // batch_size + 1}: sent {len(batch)} messages")
            
        return sent_count


def main():
    """Example usage of ActiveMQ producer."""
    producer = ActiveMQProducer(
        host='b-1.mycluster.abc123.use1-activemq-1.amazonmq.com',
        port=61613,
        username='admin',
        password='password123',
        use_ssl=True
    )
    
    # Connect
    if producer.connect():
        # Send single message
        message = {
            'event_id': 'evt-001',
            'event_type': 'user_signup',
            'timestamp': int(time.time()),
            'user_id': 'user-123',
            'email': 'user@example.com'
        }
        
        producer.send(
            destination='queue://user-events',
            message=message,
            headers={'source': 'web-app', 'version': '1.0'},
            priority=5
        )
        
        # Send batch
        messages = [
            {
                'event_id': f'evt-{i:03d}',
                'event_type': 'page_view',
                'timestamp': int(time.time()),
                'user_id': f'user-{i % 10}',
                'page': f'/products/{i}'
            }
            for i in range(100)
        ]
        
        sent = producer.send_batch(
            destination='queue://user-events',
            messages=messages,
            batch_size=50
        )
        
        print(f"Sent {sent} messages")
        
        # Disconnect
        producer.disconnect()


if __name__ == '__main__':
    main()

import pika
import json
import logging
import signal
import sys
from typing import Dict, Any, Callable, Optional
from contextlib import contextmanager

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class RabbitMQConsumer:
    """Production-grade RabbitMQ consumer with error handling."""
    
    def __init__(
        self,
        host: str,
        port: int = 5671,
        username: str = None,
        password: str = None,
        virtual_host: str = '/',
        use_ssl: bool = True
    ):
        self.host = host
        self.port = port
        self.username = username
        self.password = password
        self.virtual_host = virtual_host
        self.use_ssl = use_ssl
        self.connection = None
        self.channel = None
        self.running = False
        
        signal.signal(signal.SIGINT, self._shutdown)
        signal.signal(signal.SIGTERM, self._shutdown)
        
    def _shutdown(self, signum, frame):
        """Handle shutdown signals gracefully."""
        logger.info("Shutdown signal received, stopping consumer...")
        self.running = False
        if self.channel:
            self.channel.stop_consuming()
            
    def connect(self) -> bool:
        """Connect to RabbitMQ broker."""
        try:
            credentials = None
            if self.username and self.password:
                credentials = pika.PlainCredentials(self.username, self.password)
            
            ssl_options = None
            if self.use_ssl:
                ssl_options = pika.SSLOptions(pika.ssl.PROTOCOL_TLS)
            
            parameters = pika.ConnectionParameters(
                host=self.host,
                port=self.port,
                virtual_host=self.virtual_host,
                credentials=credentials,
                ssl_options=ssl_options,
                heartbeat=600,
                blocked_connection_timeout=300
            )
            
            self.connection = pika.BlockingConnection(parameters)
            self.channel = self.connection.channel()
            
            logger.info(f"Connected to RabbitMQ at {self.host}:{self.port}")
            return True
            
        except Exception as e:
            logger.error(f"Failed to connect to RabbitMQ: {e}")
            return False
            
    def disconnect(self):
        """Disconnect from broker."""
        try:
            if self.connection and self.connection.is_open:
                self.connection.close()
                logger.info("Disconnected from RabbitMQ")
        except Exception as e:
            logger.error(f"Error disconnecting: {e}")
            
    def declare_queue(
        self,
        queue_name: str,
        durable: bool = True,
        exclusive: bool = False,
        auto_delete: bool = False
    ) -> bool:
        """Declare a queue."""
        try:
            if not self.channel:
                return False
                
            self.channel.queue_declare(
                queue=queue_name,
                durable=durable,
                exclusive=exclusive,
                auto_delete=auto_delete
            )
            
            logger.info(f"Queue declared: {queue_name}")
            return True
            
        except Exception as e:
            logger.error(f"Failed to declare queue: {e}")
            return False
            
    def consume(
        self,
        queue_name: str,
        handler: Callable[[Dict], bool],
        prefetch_count: int = 10,
        auto_ack: bool = False
    ):
        """Consume messages from a queue."""
        try:
            if not self.channel:
                if not self.connect():
                    raise ConnectionError("Failed to connect to RabbitMQ")
            
            self.channel.basic_qos(prefetch_count=prefetch_count)
            self.running = True
            
            def callback(ch, method, properties, body):
                try:
                    message = json.loads(body.decode('utf-8'))
                    
                    # Add metadata
                    message['_metadata'] = {
                        'delivery_tag': method.delivery_tag,
                        'exchange': method.exchange,
                        'routing_key': method.routing_key,
                        'timestamp': properties.timestamp,
                        'message_id': properties.message_id
                    }
                    
                    success = handler(message)
                    
                    if success:
                        if not auto_ack:
                            ch.basic_ack(delivery_tag=method.delivery_tag)
                    else:
                        if not auto_ack:
                            ch.basic_nack(
                                delivery_tag=method.delivery_tag,
                                requeue=True
                            )
                            
                except json.JSONDecodeError as e:
                    logger.error(f"Failed to decode message: {e}")
                    if not auto_ack:
                        ch.basic_nack(
                            delivery_tag=method.delivery_tag,
                            requeue=False
                        )
                except Exception as e:
                    logger.error(f"Error processing message: {e}")
                    if not auto_ack:
                        ch.basic_nack(
                            delivery_tag=method.delivery_tag,
                            requeue=True
                        )
            
            self.channel.basic_consume(
                queue=queue_name,
                on_message_callback=callback,
                auto_ack=auto_ack
            )
            
            logger.info(f"Consuming from queue: {queue_name}")
            
            while self.running:
                try:
                    self.connection.process_data_events(time_limit=1)
                except pika.exceptions.AMQPConnectionError:
                    logger.warning("Connection lost, reconnecting...")
                    self.connect()
                    
        except Exception as e:
            logger.error(f"Consumer error: {e}")
        finally:
            self.disconnect()


def message_handler(message: Dict) -> bool:
    """Example message handler."""
    try:
        event = message.get('value', {})
        logger.info(
            f"Processing event: {event.get('event_type')} "
            f"from user {event.get('user_id')}"
        )
        
        # Add processing logic here
        return True
        
    except Exception as e:
        logger.error(f"Handler error: {e}")
        return False


def main():
    """Main consumer example."""
    consumer = RabbitMQConsumer(
        host='b-1.mycluster.abc123.use1-rabbitmq-1.amazonmq.com',
        port=5671,
        username='admin',
        password='password123',
        use_ssl=True
    )
    
    # Declare queue
    consumer.declare_queue(
        queue_name='user-events',
        durable=True
    )
    
    # Consume messages
    consumer.consume(
        queue_name='user-events',
        handler=message_handler,
        prefetch_count=10
    )


if __name__ == '__main__':
    main()

import boto3
import json
import logging
from typing import Dict, Any, List

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class MQMigrationManager:
    """Manages migration to Amazon MQ."""
    
    def __init__(self, region: str = 'us-east-1'):
        self.client = boto3.client('mq', region_name=region)
        self.s3_client = boto3.client('s3', region_name=region)
        
    def export_active_mq_config(
        self,
        broker_id: str,
        s3_bucket: str
    ) -> Dict[str, Any]:
        """Export ActiveMQ configuration from existing broker."""
        try:
            # Get broker configuration
            response = self.client.describe_broker(BrokerId=broker_id)
            config = response.get('Configuration', {})
            
            # Store in S3
            s3_key = f'migration/activemq/config-{broker_id}.json'
            self.s3_client.put_object(
                Bucket=s3_bucket,
                Key=s3_key,
                Body=json.dumps(config, indent=2)
            )
            
            logger.info(f"Configuration exported to s3://{s3_bucket}/{s3_key}")
            return config
            
        except Exception as e:
            logger.error(f"Failed to export configuration: {e}")
            raise
            
    def create_target_broker(
        self,
        broker_name: str,
        engine_type: str,
        engine_version: str,
        instance_type: str,
        subnet_ids: List[str],
        security_group_ids: List[str],
        deployment_mode: str = 'ACTIVE_STANDBY_MULTI_AZ'
    ) -> str:
        """Create target Amazon MQ broker."""
        try:
            response = self.client.create_broker(
                BrokerName=broker_name,
                EngineType=engine_type,
                EngineVersion=engine_version,
                HostInstanceType=instance_type,
                DeploymentMode=deployment_mode,
                SecurityGroups=security_group_ids,
                SubnetIds=subnet_ids,
                Users=[
                    {
                        'Username': 'admin',
                        'Password': self._generate_password()
                    }
                ],
                Configuration={
                    'Id': f'config-{broker_name}',
                    'Revision': 1
                },
                Logs={
                    'General': True
                },
                Tags={
                    'Migration': 'true',
                    'Source': 'on-premises'
                }
            )
            
            broker_arn = response['BrokerArn']
            logger.info(f"Target broker created: {broker_arn}")
            return broker_arn
            
        except Exception as e:
            logger.error(f"Failed to create broker: {e}")
            raise
            
    def import_configuration(
        self,
        config_name: str,
        engine_type: str,
        engine_version: str,
        config_data: Dict[str, Any]
    ) -> str:
        """Import configuration to Amazon MQ."""
        try:
            response = self.client.create_configuration(
                Name=config_name,
                Description=f'Migrated {engine_type} configuration',
                EngineType=engine_type,
                EngineVersion=engine_version,
                ConfigurationOptions=json.dumps(config_data).encode('utf-8')
            )
            
            config_arn = response['Arn']
            logger.info(f"Configuration imported: {config_arn}")
            return config_arn
            
        except Exception as e:
            logger.error(f"Failed to import configuration: {e}")
            raise
            
    def validate_migration(
        self,
        source_broker_id: str,
        target_broker_arn: str
    ) -> Dict[str, Any]:
        """Validate migration completeness."""
        try:
            # Get source broker metrics
            source_metrics = self._get_broker_metrics(source_broker_id)
            
            # Get target broker metrics
            target_metrics = self._get_broker_metrics(target_broker_arn)
            
            validation = {
                'source': source_metrics,
                'target': target_metrics,
                'queues_match': source_metrics.get('queue_count', 0) == target_metrics.get('queue_count', 0),
                'topics_match': source_metrics.get('topic_count', 0) == target_metrics.get('topic_count', 0),
                'messages_migrated': source_metrics.get('message_count', 0) == target_metrics.get('message_count', 0)
            }
            
            logger.info(f"Migration validation: {validation}")
            return validation
            
        except Exception as e:
            logger.error(f"Failed to validate migration: {e}")
            raise
            
    def _get_broker_metrics(self, broker_id: str) -> Dict[str, Any]:
        """Get broker metrics."""
        # Simplified - in production, query CloudWatch
        return {
            'queue_count': 10,
            'topic_count': 5,
            'message_count': 100000,
            'consumer_count': 20
        }
        
    def _generate_password(self) -> str:
        """Generate secure password."""
        import secrets
        import string
        alphabet = string.ascii_letters + string.digits + string.punctuation
        return ''.join(secrets.choice(alphabet) for _ in range(16))


def main():
    """Example migration workflow."""
    migration = MQMigrationManager(region='us-east-1')
    
    # Export configuration
    config = migration.export_active_mq_config(
        broker_id='old-broker-id',
        s3_bucket='migration-bucket'
    )
    
    # Create target broker
    target_broker = migration.create_target_broker(
        broker_name='new-data-pipeline-broker',
        engine_type='ACTIVEMQ',
        engine_version='5.17.6',
        instance_type='mq.m5.large',
        subnet_ids=['subnet-aaa', 'subnet-bbb'],
        security_group_ids=['sg-123456']
    )
    
    # Validate migration
    validation = migration.validate_migration(
        source_broker_id='old-broker-id',
        target_broker_arn=target_broker
    )
    
    print(f"Migration validation: {validation}")


if __name__ == '__main__':
    main()

# Security group configuration
security_group_rules = [
    {
        'Protocol': 'tcp',
        'FromPort': 61613,
        'ToPort': 61613,
        'CidrIp': '10.0.0.0/16',
        'Description': 'ActiveMQ STOMP (internal)'
    },
    {
        'Protocol': 'tcp',
        'FromPort': 61614,
        'ToPort': 61614,
        'CidrIp': '10.0.0.0/16',
        'Description': 'ActiveMQ AMQP (internal)'
    },
    {
        'Protocol': 'tcp',
        'FromPort': 5671,
        'ToPort': 5671,
        'CidrIp': '10.0.0.0/16',
        'Description': 'RabbitMQ AMQP (internal)'
    },
    {
        'Protocol': 'tcp',
        'FromPort': 8162,
        'ToPort': 8162,
        'CidrIp': '10.0.0.0/16',
        'Description': 'Management Console'
    }
]