"""
Production Kinesis Data Streams producer with batching, retries, and error handling.
"""
import json
import logging
import time
from typing import List, Dict, Any
import boto3
from botocore.config import Config

logger = logging.getLogger(__name__)

class KinesisProducer:
    def __init__(self, stream_name: str, region: str = 'us-east-1'):
        self.stream_name = stream_name
        self.client = boto3.client(
            'kinesis',
            config=Config(
                retries={'max_attempts': 3, 'mode': 'adaptive'},
                connect_timeout=5,
                read_timeout=10
            ),
            region_name=region
        )
    
    def put_records(self, records: List[Dict[str, Any]], 
                    batch_size: int = 500) -> Dict[str, Any]:
        """Put records to Kinesis with retry logic."""
        total_sent = 0
        total_failed = 0
        
        for i in range(0, len(records), batch_size):
            batch = records[i:i + batch_size]
            
            kinesis_records = [
                {
                    'Data': json.dumps(record).encode('utf-8'),
                    'PartitionKey': record.get('partition_key', str(i))
                }
                for i, record in enumerate(batch)
            ]
            
            try:
                response = self.client.put_records(
                    StreamName=self.stream_name,
                    Records=kinesis_records
                )
                
                failed = response.get('FailedRecordCount', 0)
                total_sent += len(batch) - failed
                total_failed += failed
                
                if failed > 0:
                    logger.warning(f"Batch {i // batch_size}: {failed} records failed")
                    self._retry_failed(kinesis_records, response['Records'])
                    
            except Exception as e:
                logger.error(f"Batch put failed: {str(e)}")
                total_failed += len(batch)
        
        return {
            'total_sent': total_sent,
            'total_failed': total_failed,
            'success_rate': total_sent / max(len(records), 1) * 100
        }
    
    def _retry_failed(self, original_records: List, 
                      responses: List[Dict], max_retries: int = 3):
        """Retry failed records with exponential backoff."""
        for attempt in range(max_retries):
            failed_records = [
                original_records[i]
                for i, resp in enumerate(responses)
                if 'ErrorCode' in resp
            ]
            
            if not failed_records:
                break
            
            time.sleep(2 ** attempt)
            
            response = self.client.put_records(
                StreamName=self.stream_name,
                Records=failed_records
            )
            
            responses = response['Records']
            logger.info(f"Retry {attempt + 1}: {response.get('FailedRecordCount', 0)} still failed")

"""
Production Kinesis Consumer using KCL with checkpointing and error handling.
"""
import logging
from amazon_kclpy import processor
from amazon_kclpy.v2 import kcl_process

logger = logging.getLogger(__name__)

class StreamRecordProcessor(processor.RecordProcessorBase):
    def __init__(self):
        self.checkpoint_counter = 0
        self.records_processed = 0
        self.batch_size = 100
    
    def initialize(self, context):
        logger.info(f"Initializing processor for shard: {context.shard_id}")
        self.checkpoint_counter = 0
        self.records_processed = 0
    
    def process_records(self, records, checkpointer):
        for record in records:
            try:
                self._process_single_record(record)
                self.records_processed += 1
            except Exception as e:
                logger.error(f"Failed to process record: {str(e)}")
                continue
            
            self.checkpoint_counter += 1
            
            if self.checkpoint_counter >= self.batch_size:
                checkpointer.checkpoint(record.sequence_number)
                self.checkpoint_counter = 0
                logger.info(f"Checkpointed at {self.records_processed} records")
    
    def _process_single_record(self, record):
        """Process a single record with business logic."""
        data = json.loads(record.data.decode('utf-8'))
        
        # Validate required fields
        required_fields = ['event_id', 'event_type', 'timestamp']
        for field in required_fields:
            if field not in data:
                raise ValueError(f"Missing required field: {field}")
        
        # Process based on event type
        if data['event_type'] == 'purchase':
            self._handle_purchase(data)
        elif data['event_type'] == 'page_view':
            self._handle_page_view(data)
        else:
            logger.warning(f"Unknown event type: {data['event_type']}")
    
    def _handle_purchase(self, data):
        """Handle purchase events."""
        logger.info(f"Processing purchase: {data['event_id']}")
    
    def _handle_page_view(self, data):
        """Handle page view events."""
        logger.debug(f"Processing page view: {data['event_id']}")
    
    def shutdown(self, checkpointer, reason):
        logger.info(f"Shutting down. Reason: {reason}")
        checkpointer.checkpoint()

if __name__ == "__main__":
    kcl_process.main(
        processor=StreamRecordProcessor(),
        log_level='INFO'
    )

"""
Lambda function for Kinesis Data Firehose transformation.
Transforms, validates, and enriches streaming records.
"""
import json
import base64
import logging
from typing import Dict, List, Any

logger = logging.getLogger(__name__)

def transform_record(record: Dict[str, Any]) -> Dict[str, Any]:
    """Transform a single record from Firehose."""
    try:
        # Decode the data
        raw_data = base64.b64decode(record['data']).decode('utf-8')
        data = json.loads(raw_data)
        
        # Apply transformations
        transformed = {
            'event_id': data.get('event_id'),
            'user_id': data.get('user_id'),
            'event_type': data.get('event_type', 'unknown'),
            'amount': float(data.get('amount', 0)),
            'timestamp': data.get('timestamp'),
            'processed_at': datetime.utcnow().isoformat()
        }
        
        # Validate required fields
        if not transformed['event_id'] or not transformed['user_id']:
            raise ValueError("Missing required fields")
        
        # Enrichment example: add computed fields
        transformed['is_high_value'] = transformed['amount'] > 1000
        
        return {
            'recordId': record['recordId'],
            'result': 'Ok',
            'data': base64.b64encode(
                json.dumps(transformed).encode('utf-8')
            ).decode('utf-8')
        }
        
    except Exception as e:
        logger.error(f"Transform failed: {str(e)}")
        return {
            'recordId': record['recordId'],
            'result': 'ProcessingFailed',
            'data': record['data']
        }

def lambda_handler(event, context):
    """Main Lambda handler for Firehose transformation."""
    output_records = []
    
    for record in event['records']:
        output_records.append(transform_record(record))
    
    logger.info(f"Processed {len(output_records)} records")
    
    return {'records': output_records}