import boto3
import json
import logging
import time
from datetime import datetime

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


def create_kinesis_stream(
    stream_name: str,
    shard_count: int = 4,
    retention_hours: int = 24,
    enable_enhanced_fan_out: bool = False
) -> dict:
    """
    Create a Kinesis Data Stream with production-grade settings.

    Args:
        stream_name: Unique stream identifier
        shard_count: Number of shards for throughput
        retention_hours: Data retention period (24-8760 hours)
        enable_enhanced_fan_out: Enable enhanced fan-out for consumers

    Returns:
        dict with stream details
    """
    kinesis_client = boto3.client('kinesis')

    stream_config = {
        'StreamName': stream_name,
        'ShardCount': shard_count,
        'StreamModeDetails': {
            'StreamMode': 'PROVISIONED'
        },
        'Tags': {
            'Environment': 'production',
            'ManagedBy': 'boto3'
        }
    }

    try:
        response = kinesis_client.create_stream(**stream_config)
        logger.info(f"Created stream: {stream_name} with {shard_count} shards")

        # Wait for stream to become active
        waiter = kinesis_client.get_waiter('stream_exists')
        logger.info("Waiting for stream to become active...")
        waiter.wait(StreamName=stream_name)

        # Update retention period if needed
        if retention_hours > 24:
            kinesis_client.increase_stream_retention_period(
                StreamName=stream_name,
                RetentionPeriodHours=retention_hours
            )
            logger.info(f"Increased retention to {retention_hours} hours")

        # Enable enhanced fan-out if requested
        if enable_enhanced_fan_out:
            register_enhanced_consumer(stream_name, f"{stream_name}-consumer")

        return {
            'stream_name': stream_name,
            'shard_count': shard_count,
            'retention_hours': retention_hours,
            'status': 'ACTIVE'
        }

    except kinesis_client.exceptions.ResourceInUseException:
        logger.warning(f"Stream {stream_name} already exists")
        return describe_stream(stream_name)
    except Exception as e:
        logger.error(f"Failed to create stream: {str(e)}")
        raise


def put_records_to_stream(
    stream_name: str,
    records: list,
    partition_key_extractor: callable = None
) -> dict:
    """
    Put records to a Kinesis stream with error handling.

    Args:
        stream_name: Target stream name
        records: List of record dictionaries with 'data' and optional 'partition_key'
        partition_key_extractor: Function to extract partition key from record

    Returns:
        dict with success count and failed records
    """
    kinesis_client = boto3.client('kinesis')

    formatted_records = []
    for i, record in enumerate(records):
        partition_key = record.get('partition_key')
        if not partition_key and partition_key_extractor:
            partition_key = partition_key_extractor(record)
        elif not partition_key:
            partition_key = f"pk-{i}"

        formatted_records.append({
            'Data': json.dumps(record['data']).encode('utf-8'),
            'PartitionKey': partition_key
        })

    failed_records = []
    success_count = 0

    # Process in batches of 500 (Kinesis limit)
    for i in range(0, len(formatted_records), 500):
        batch = formatted_records[i:i + 500]

        try:
            response = kinesis_client.put_records(
                StreamName=stream_name,
                Records=batch
            )

            success_count += response['Records'].__len__() - response['FailedRecordCount']
            failed_records.extend([
                {'index': i + idx, 'error_code': r.get('ErrorCode'), 'error_msg': r.get('ErrorMessage')}
                for idx, r in enumerate(response['Records'])
                if 'ErrorCode' in r
            ])

            if response['FailedRecordCount'] > 0:
                logger.warning(f"Batch {i // 500}: {response['FailedRecordCount']} records failed")

        except Exception as e:
            logger.error(f"Batch {i // 500} failed completely: {str(e)}")
            failed_records.extend([{'index': i + j, 'error': str(e)} for j in range(len(batch))])

    result = {
        'total_records': len(records),
        'success_count': success_count,
        'failed_count': len(failed_records),
        'failed_records': failed_records
    }

    logger.info(f"Put {success_count}/{len(records)} records to {stream_name}")
    return result


def register_enhanced_consumer(stream_name: str, consumer_name: str) -> str:
    """Register an enhanced fan-out consumer."""
    kinesis_client = boto3.client('kinesis')

    response = kinesis_client.register_stream_consumer(
        StreamName=stream_name,
        ConsumerName=consumer_name
    )

    consumer_arn = response['Consumer']['ConsumerARN']
    logger.info(f"Registered enhanced consumer: {consumer_arn}")
    return consumer_arn


def describe_stream(stream_name: str) -> dict:
    """Get stream details including shard information."""
    kinesis_client = boto3.client('kinesis')

    response = kinesis_client.describe_stream(StreamName=stream_name)
    stream_desc = response['StreamDescription']

    return {
        'stream_name': stream_desc['StreamName'],
        'stream_arn': stream_desc['StreamARN'],
        'shard_count': len(stream_desc['Shards']),
        'retention_hours': stream_desc['RetentionPeriodHours'],
        'status': stream_desc['StreamStatus'],
        'creation_time': str(stream_desc['StreamCreationTimestamp'])
    }


def get_shard_iterator(stream_name: str, shard_id: str, iterator_type: str = 'LATEST') -> str:
    """Get a shard iterator for reading records."""
    kinesis_client = boto3.client('kinesis')

    response = kinesis_client.get_shard_iterator(
        StreamName=stream_name,
        ShardId=shard_id,
        ShardIteratorType=iterator_type
    )

    return response['ShardIterator']

import boto3
import json
import logging

logger = logging.getLogger(__name__)


def create_firehose_delivery_stream(
    stream_name: str,
    s3_destination: str,
    iam_role_arn: str,
    buffer_size_mb: int = 5,
    buffer_interval_sec: int = 300,
    compression: str = "GZIP",
    format_conversion: str = "PARQUET"
) -> dict:
    """
    Create a Kinesis Data Firehose delivery stream.

    Args:
        stream_name: Delivery stream name
        s3_destination: S3 bucket ARN for delivery
        iam_role_arn: IAM role ARN for Firehose
        buffer_size_mb: Buffer size in MB (1-128)
        buffer_interval_sec: Buffer interval in seconds (60-900)
        compression: Compression format (GZIP, SNAPPY, ZIP)
        format_conversion: Target format (PARQUET, ORC, AVRO)

    Returns:
        dict with delivery stream details
    """
    firehose_client = boto3.client('firehose')

    s3_config = {
        'RoleARN': iam_role_arn,
        'BucketARN': s3_destination,
        'Prefix': 'raw/year=!{timestamp:yyyy}/month=!{timestamp:MM}/day=!{timestamp:dd}/',
        'ErrorOutputPrefix': 'errors/',
        'BufferingHints': {
            'SizeInMBs': buffer_size_mb,
            'IntervalInSeconds': buffer_interval_sec
        },
        'CompressionFormat': compression,
        'CloudWatchLoggingOptions': {
            'Enabled': True,
            'LogGroupName': '/aws/kinesis-firehose',
            'LogStreamName': stream_name
        }
    }

    if format_conversion:
        s3_config['ProcessingConfiguration'] = {
            'Enabled': True,
            'Processors': [{
                'Type': 'ConvertRecord',
                'Parameters': [{
                    'ParameterName': 'OutputFormat',
                    'ParameterValue': format_conversion
                }]
            }]
        }

    try:
        response = firehose_client.create_delivery_stream(
            DeliveryStreamName=stream_name,
            DeliveryStreamType='DirectPut',
            S3DestinationConfiguration=s3_config,
            Tags={
                'Environment': 'production',
                'ManagedBy': 'boto3'
            }
        )

        stream_arn = response['DeliveryStreamARN']
        logger.info(f"Created Firehose stream: {stream_arn}")
        return {'stream_arn': stream_arn, 'status': 'CREATING'}

    except firehose_client.exceptions.ResourceInUseException:
        logger.warning(f"Firehose stream {stream_name} already exists")
        return describe_firehose_stream(stream_name)
    except Exception as e:
        logger.error(f"Failed to create Firehose stream: {str(e)}")
        raise


def describe_firehose_stream(stream_name: str) -> dict:
    """Get Firehose delivery stream details."""
    firehose_client = boto3.client('firehose')

    response = firehose_client.describe_delivery_stream(
        DeliveryStreamName=stream_name
    )

    stream_desc = response['DeliveryStreamDescription']
    return {
        'stream_name': stream_desc['DeliveryStreamName'],
        'stream_arn': stream_desc['DeliveryStreamARN'],
        'status': stream_desc['DeliveryStreamStatus'],
        'creation_time': str(stream_desc['CreateTimestamp']),
        'destinations': len(stream_desc.get('Destinations', []))
    }