import boto3
import json
import logging
from datetime import datetime
from typing import Dict, List, Optional
from pymongo import MongoClient
from pymongo.errors import (
    ConnectionFailure,
    OperationFailure,
    BulkWriteError
)

logger = logging.getLogger(__name__)

class DocumentDBManager:
    """Production-grade Amazon DocumentDB manager."""

    def __init__(
        self,
        cluster_endpoint: str,
        port: int = 27017,
        database: str = 'admin',
        username: str = None,
        password: str = None,
        replica_set: str = 'rs0',
        read_preference: str = 'secondaryPreferred'
    ):
        self.connection_string = (
            f"mongodb://{username}:{password}@{cluster_endpoint}:{port}"
            f"/?replicaSet={replica_set}"
            f"&readPreference={read_preference}"
            f"&retryWrites=false"
            f"&ssl=true"
            f"&tlsAllowInvalidCertificates=false"
        )
        self.client = MongoClient(self.connection_string)
        self.db = self.client[database]

    def insert_document(
        self,
        collection: str,
        document: Dict
    ) -> str:
        """Insert a single document with error handling."""
        try:
            result = self.db[collection].insert_one(document)
            logger.info(f"Inserted document: {result.inserted_id}")
            return str(result.inserted_id)
        except OperationFailure as e:
            logger.error(f"Insert failed: {e}")
            raise

    def bulk_insert(
        self,
        collection: str,
        documents: List[Dict],
        ordered: bool = False
    ) -> Dict:
        """Insert multiple documents in bulk with error handling."""
        results = {'success': 0, 'failed': 0, 'errors': []}
        try:
            result = self.db[collection].insert_many(
                documents,
                ordered=ordered
            )
            results['success'] = len(result.inserted_ids)
        except BulkWriteError as e:
            results['success'] = e.details['nInserted']
            results['failed'] = len(e.details['writeErrors'])
            results['errors'] = [
                err['errmsg'] for err in e.details['writeErrors']
            ]
            logger.warning(
                f"Bulk insert: {results['success']} OK, "
                f"{results['failed']} failed"
            )
        except Exception as e:
            logger.error(f"Bulk insert failed: {e}")
            raise

        return results

    def find_with_aggregation(
        self,
        collection: str,
        pipeline: List[Dict],
        max_time_ms: int = 30000
    ) -> List[Dict]:
        """Execute aggregation pipeline with timeout."""
        try:
            cursor = self.db[collection].aggregate(
                pipeline,
                maxTimeMS=max_time_ms,
                allowDiskUse=True
            )
            return list(cursor)
        except OperationFailure as e:
            logger.error(f"Aggregation failed: {e}")
            raise

    def create_indexes(
        self,
        collection: str,
        indexes: List[Dict]
    ) -> List[str]:
        """Create indexes on a collection."""
        index_names = []
        try:
            for index in indexes:
                name = self.db[collection].create_index(
                    index['keys'],
                    unique=index.get('unique', False),
                    name=index.get('name'),
                    background=True
                )
                index_names.append(name)
                logger.info(f"Created index: {name}")
        except Exception as e:
            logger.error(f"Index creation failed: {e}")
            raise
        return index_names

    def watch_change_stream(
        self,
        collection: str,
        resume_token: Optional[Dict] = None,
        pipeline: Optional[List[Dict]] = None
    ):
        """Watch for changes using change streams."""
        try:
            change_stream = self.db[collection].watch(
                pipeline=pipeline or [],
                resume_after=resume_token,
                full_document='updateLookup'
            )
            for change in change_stream:
                yield change
                resume_token = change_cluster._resume_token
        except ConnectionFailure as e:
            logger.error(f"Change stream lost connection: {e}")
            raise


class DocumentDBProvisioner:
    """Provision and manage DocumentDB clusters using boto3."""

    def __init__(self, region: str = 'us-east-1'):
        self.client = boto3.client('docdb', region_name=region)

    def create_cluster(
        self,
        cluster_id: str,
        master_username: str,
        master_password: str,
        instance_class: str = 'db.r5.large',
        instance_count: int = 3,
        backup_retention_days: int = 7,
        encryption_enabled: bool = True,
        enable_cloudwatch_logs: bool = True
    ) -> Dict:
        """Create a production DocumentDB cluster."""
        try:
            # Create subnet group
            self.client.create_db_subnet_group(
                DBSubnetGroupName=f'{cluster_id}-subnet-group',
                DBSubnetGroupDescription=f'Subnet group for {cluster_id}',
                SubnetIds=['subnet-xxx', 'subnet-yyy', 'subnet-zzz']
            )

            # Create cluster
            response = self.client.create_db_cluster(
                DBClusterIdentifier=cluster_id,
                Engine='docdb',
                MasterUsername=master_username,
                MasterUserPassword=master_password,
                DBSubnetGroupName=f'{cluster_id}-subnet-group',
                BackupRetentionPeriod=backup_retention_days,
                StorageEncrypted=encryption_enabled,
                EnableCloudwatchLogsExports=[
                    'audit', 'profiler'
                ] if enable_cloudwatch_logs else [],
                Tags=[
                    {'Key': 'Environment', 'Value': 'production'},
                    {'Key': 'Service', 'Value': 'documentdb'}
                ]
            )

            # Create instances
            for i in range(instance_count):
                self.client.create_db_instance(
                    DBInstanceIdentifier=f'{cluster_id}-instance-{i+1}',
                    DBClusterIdentifier=cluster_id,
                    DBInstanceClass=instance_class,
                    Engine='docdb'
                )

            logger.info(f"Created cluster: {cluster_id}")
            return response['DBCluster']
        except Exception as e:
            logger.error(f"Failed to create cluster: {e}")
            raise

    def create_read_replica(
        self,
        cluster_id: str,
        replica_id: str,
        instance_class: str = 'db.r5.large'
    ) -> Dict:
        """Create a read replica instance."""
        try:
            response = self.client.create_db_instance(
                DBInstanceIdentifier=replica_id,
                DBClusterIdentifier=cluster_id,
                DBInstanceClass=instance_class,
                Engine='docdb',
                Tags=[
                    {'Key': 'Role', 'Value': 'read-replica'}
                ]
            )
            logger.info(f"Created read replica: {replica_id}")
            return response['DBInstance']
        except Exception as e:
            logger.error(f"Failed to create replica: {e}")
            raise

    def enable_cloudwatch_export(
        self,
        cluster_id: str,
        log_types: List[str] = ['audit', 'profiler']
    ) -> Dict:
        """Enable CloudWatch log exports for a cluster."""
        try:
            response = self.client.modify_db_cluster(
                DBClusterIdentifier=cluster_id,
                EnableCloudwatchLogsExports=log_types,
                ApplyImmediately=True
            )
            logger.info(f"Enabled log exports for: {cluster_id}")
            return response['DBCluster']
        except Exception as e:
            logger.error(f"Failed to enable logs: {e}")
            raise


# Production usage
if __name__ == '__main__':
    provisioner = DocumentDBProvisioner()
    cluster = provisioner.create_cluster(
        cluster_id='prod-documentdb',
        master_username='admin',
        master_password='SecureP@ss123!',
        instance_class='db.r5.2xlarge',
        instance_count=4,
        backup_retention_days=14
    )
    print(f"Cluster ARN: {cluster['DBClusterArn']}")

#!/bin/bash
# DocumentDB cluster backup and export script

set -euo pipefail

CLUSTER_ID="${1:-prod-documentdb}"
REGION="${AWS_REGION:-us-east-1}"
BACKUP_BUCKET="${BACKUP_BUCKET:-my-docdb-backups}"
EXPORT_FORMAT="${2:-json}"

echo "=== DocumentDB Export & Backup ==="
echo "Cluster: ${CLUSTER_ID}"
echo "Region: ${REGION}"
echo "Timestamp: $(date -u +%Y-%m-%dT%H:%M:%SZ)"

# Get cluster details
CLUSTER_INFO=$(aws docdb describe-db-clusters \
  --db-cluster-identifier "${CLUSTER_ID}" \
  --region "${REGION}" \
  --query 'DBClusters[0]' \
  --output json)

ENDPOINT=$(echo "$CLUSTER_INFO" | jq -r '.Endpoint')
READ_ENDPOINT=$(echo "$CLUSTER_INFO" | jq -r '.ReaderEndpoint')
STATUS=$(echo "$CLUSTER_INFO" | jq -r '.Status')
ENGINE=$(echo "$CLUSTER_INFO" | jq -r '.EngineVersion')
BACKUP_RETENTION=$(echo "$CLUSTER_INFO" | jq -r '.BackupRetentionPeriod')

echo "Status: ${STATUS}"
echo "Engine: ${ENGINE}"
echo "Endpoint: ${ENDPOINT}"
echo "Reader Endpoint: ${READ_ENDPOINT}"
echo "Backup Retention: ${BACKUP_RETENTION} days"

# Check instance health
INSTANCES=$(aws docdb describe-db-instances \
  --filters "Name=db-cluster-id,Values=${CLUSTER_ID}" \
  --region "${REGION}" \
  --query 'DBInstances[*].[DBInstanceIdentifier,DBInstanceClass,Status]' \
  --output table)

echo ""
echo "=== Cluster Instances ==="
echo "$INSTANCES"

# Monitor IOPS and storage
METRICS=$(aws cloudwatch get-metric-statistics \
  --namespace AWS/DocDB \
  --metric-name FreeStorageSpace \
  --dimensions Name=DBClusterIdentifier,Value="${CLUSTER_ID}" \
  --start-time "$(date -u -d '1 hour ago' +%Y-%m-%dT%H:%M:%SZ)" \
  --end-time "$(date -u +%Y-%m-%dT%H:%M:%SZ)" \
  --period 300 \
  --statistics Average \
  --region "${REGION}")

STORAGE_GB=$(echo "$METRICS" | jq '.Datapoints[-1].Average // 0' | awk '{printf "%.2f", $1/1073741824}')
echo ""
echo "=== Storage Metrics ==="
echo "Free Storage: ${STORAGE_GB} GB"

# Export collection to S3 using Glue
echo ""
echo "=== Starting Glue Export Job ==="
aws glue start-job-run \
  --job-name "docdb-export-${CLUSTER_ID}" \
  --arguments "{
    \"--cluster_endpoint\": \"${ENDPOINT}\",
    \"--database\": \"analytics_db\",
    \"--s3_output\": \"s3://${BACKUP_BUCKET}/exports/\",
    \"--format\": \"${EXPORT_FORMAT}\",
    \"--timestamp\": \"$(date -u +%Y%m%d%H%M%S)\"
  }" \
  --region "${REGION}"

echo "Export job started. Monitor in Glue console."

# Create manual snapshot
SNAPSHOT_ID="${CLUSTER_ID}-manual-$(date -u +%Y%m%d%H%M%S)"
aws docdb create-db-cluster-snapshot \
  --db-cluster-snapshot-identifier "${SNAPSHOT_ID}" \
  --db-cluster-identifier "${CLUSTER_ID}" \
  --region "${REGION}"

echo ""
echo "=== Manual Snapshot Created ==="
echo "Snapshot: ${SNAPSHOT_ID}"

# Cleanup old snapshots (keep last 5 manual)
echo ""
echo "=== Cleaning Up Old Snapshots ==="
SNAPSHOTS=$(aws docdb describe-db-cluster-snapshots \
  --snapshot-type manual \
  --region "${REGION}" \
  --query "DBClusterSnapshots[?starts_with(DBClusterSnapshotIdentifier, '${CLUSTER_ID}-manual')].DBClusterSnapshotIdentifier" \
  --output json)

SNAPSHOT_COUNT=$(echo "$SNAPSHOTS" | jq length)
if [ "$SNAPSHOT_COUNT" -gt 5 ]; then
  echo "Cleaning up $(($SNAPSHOT_COUNT - 5)) old snapshots..."
  echo "$SNAPSHOTS" | jq -r ".[:-5][]" | while read -r snap; do
    aws docdb delete-db-cluster-snapshot \
      --db-cluster-snapshot-identifier "$snap" \
      --region "${REGION}"
    echo "Deleted: $snap"
  done
fi

echo ""
echo "Backup and export complete."