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

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


class MSKManager:
    def __init__(self, region: str = 'us-east-1'):
        self.client = boto3.client('kafka', region_name=region)
        
    def create_cluster(
        self,
        cluster_name: str,
        kafka_version: str = '3.5.1',
        broker_count: int = 3,
        instance_type: str = 'kafka.m5.large',
        storage_gb: int = 500,
        subnet_ids: list = None,
        security_group_ids: list = None
    ) -> Dict[str, Any]:
        try:
            if not subnet_ids:
                raise ValueError("subnet_ids are required")
            if not security_group_ids:
                raise ValueError("security_group_ids are required")
            if broker_count < 3:
                raise ValueError("Minimum 3 brokers required for production")
                
            response = self.client.create_cluster(
                ClusterName=cluster_name,
                KafkaVersion=kafka_version,
                NumberOfBrokerNodes=broker_count,
                BrokerNodeGroupInfo={
                    'InstanceType': instance_type,
                    'ClientSubnets': subnet_ids,
                    'SecurityGroups': security_group_ids,
                    'StorageInfo': {
                        'EBSStorageInfo': {'VolumeSize': storage_gb}
                    }
                },
                EncryptionInfo={
                    'EncryptionAtRest': {
                        'DataVolumeKMSKeyId': 'arn:aws:kms:us-east-1:123456789:key/your-key'
                    },
                    'EncryptionInTransit': {
                        'ClientBroker': 'TLS',
                        'InCluster': True
                    }
                },
                ClientAuthentication={
                    'Sasl': {'Iam': {'Enabled': True}}
                },
                EnhancedMonitoring='PER_BROKER',
                Tags={
                    'Environment': 'production',
                    'Team': 'data-engineering'
                }
            )
            
            cluster_arn = response['ClusterArn']
            logger.info(f"MSK cluster created: {cluster_arn}")
            return response
            
        except Exception as e:
            logger.error(f"Failed to create MSK cluster: {str(e)}")
            raise
            
    def wait_for_cluster(self, cluster_arn: str, timeout: int = 1800):
        start_time = time.time()
        while time.time() - start_time < timeout:
            response = self.client.describe_cluster(ClusterArn=cluster_arn)
            state = response['ClusterInfo']['State']
            if state == 'ACTIVE':
                logger.info(f"Cluster {cluster_arn} is ACTIVE")
                return response
            elif state in ['FAILED', 'DELETING']:
                raise RuntimeError(f"Cluster in unexpected state: {state}")
            time.sleep(30)
        raise TimeoutError(f"Cluster not ACTIVE within {timeout} seconds")
        
    def get_bootstrap_brokers(self, cluster_arn: str) -> Dict[str, str]:
        response = self.client.get_bootstrap_brokers(ClusterArn=cluster_arn)
        return {
            'plaintext': response.get('BootstrapBrokerString', ''),
            'tls': response.get('BootstrapBrokerStringTls', ''),
            'sasl_iam': response.get('BootstrapBrokerStringSaslIam', '')
        }


def main():
    msk = MSKManager(region='us-east-1')
    cluster_arn = msk.create_cluster(
        cluster_name='data-pipeline-production',
        kafka_version='3.5.1',
        broker_count=3,
        instance_type='kafka.m5.large',
        storage_gb=500,
        subnet_ids=['subnet-aaa', 'subnet-bbb', 'subnet-ccc'],
        security_group_ids=['sg-123456']
    )
    msk.wait_for_cluster(cluster_arn)
    brokers = msk.get_bootstrap_brokers(cluster_arn)
    print(f"TLS Endpoints: {brokers['tls']}")


if __name__ == '__main__':
    main()

from confluent_kafka import Producer, KafkaError
from confluent_kafka.admin import AdminClient, NewTopic
import json
import logging
import time
from typing import Dict, List, Optional, Any

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


class KafkaProducer:
    def __init__(self, config: Dict[str, Any]):
        self.producer = Producer({
            'bootstrap.servers': config.get('bootstrap_servers'),
            'security.protocol': config.get('security_protocol', 'SSL'),
            'ssl.ca.location': config.get('ssl_ca_location'),
            'ssl.certificate.location': config.get('ssl_certificate_location'),
            'ssl.key.location': config.get('ssl_key_location'),
            'acks': config.get('acks', 'all'),
            'retries': config.get('retries', 2147483647),
            'max.in.flight.requests.per.connection': 5,
            'enable.idempotence': True,
            'compression.type': 'lz4',
            'batch.size': 16384,
            'linger.ms': 5
        })
        self.admin = AdminClient({
            'bootstrap.servers': config.get('bootstrap_servers'),
            'security.protocol': config.get('security_protocol', 'SSL')
        })
        self.delivery_stats = {'success': 0, 'failure': 0}
        
    def delivery_callback(self, err, msg):
        if err:
            self.delivery_stats['failure'] += 1
            logger.error(f"Delivery failed: {err}")
        else:
            self.delivery_stats['success'] += 1
            
    def create_topic(self, topic_name: str, num_partitions: int = 6, replication_factor: int = 3) -> bool:
        try:
            new_topics = [NewTopic(topic=topic_name, num_partitions=num_partitions, replication_factor=replication_factor)]
            futures = self.admin.create_topics(new_topics)
            for topic, future in futures.items():
                try:
                    future.result()
                    logger.info(f"Topic {topic} created")
                    return True
                except Exception as e:
                    if 'TOPIC_ALREADY_EXISTS' in str(e):
                        return True
                    logger.error(f"Failed to create topic: {e}")
                    return False
        except Exception as e:
            logger.error(f"Topic creation error: {e}")
            return False
            
    def produce(self, topic: str, key: str, value: Dict, headers: Optional[Dict] = None) -> bool:
        try:
            serialized_value = json.dumps(value).encode('utf-8')
            serialized_key = key.encode('utf-8')
            kwargs = {'topic': topic, 'key': serialized_key, 'value': serialized_value, 'callback': self.delivery_callback}
            if headers:
                kwargs['headers'] = [(k, v.encode('utf-8')) for k, v in headers.items()]
            self.producer.produce(**kwargs)
            self.producer.poll(0)
            return True
        except BufferError:
            self.producer.poll(1)
            return False
        except Exception as e:
            logger.error(f"Failed to produce: {e}")
            return False
            
    def flush(self, timeout: float = 30.0) -> int:
        return self.producer.flush(timeout)
        
    def get_stats(self) -> Dict[str, int]:
        return self.delivery_stats.copy()


def main():
    config = {
        'bootstrap_servers': 'b-1.mycluster.abc123.c2.kafka.us-east-1.amazonaws.com:9094',
        'security_protocol': 'SSL',
        'ssl_ca_location': '/opt/kafka/ssl/ca.pem',
        'ssl_certificate_location': '/opt/kafka/ssl/service.cert',
        'ssl_key_location': '/opt/kafka/ssl/service.key'
    }
    producer = KafkaProducer(config)
    producer.create_topic('user-events')
    for i in range(100):
        event = {'event_id': f'evt-{i}', 'user_id': f'user-{i % 10}', 'event_type': 'click', 'timestamp': int(time.time())}
        producer.produce(topic='user-events', key=f'user-{i % 10}', value=event)
    producer.flush()
    stats = producer.get_stats()
    print(f"Produced: {stats['success']}, Failed: {stats['failure']}")


if __name__ == '__main__':
    main()

producer_config = {
    'batch.size': 65536,
    'linger.ms': 10,
    'compression.type': 'lz4',
    'buffer.memory': 67108864,
    'max.in.flight.requests.per.connection': 5,
    'enable.idempotence': True
}