import boto3
import json
import logging

logger = logging.getLogger(__name__)

class DMSEndpointManager:
    """Manage AWS DMS endpoints for database migration."""

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

    def create_source_endpoint(self, config):
        """Create a DMS source endpoint with error handling."""
        try:
            response = self.client.create_endpoint(
                EndpointIdentifier=config['identifier'],
                EndpointType='source',
                EngineName=config['engine'],
                ServerName=config['server'],
                Port=config['port'],
                DatabaseName=config['database'],
                Username=config['username'],
                Password=config['password'],
                SslMode=config.get('ssl_mode', 'require'),
                CertificateArn=config.get('cert_arn'),
                Tags=config.get('tags', [])
            )
            endpoint_arn = response['Endpoint']['EndpointArn']
            logger.info(f"Created source endpoint: {endpoint_arn}")
            return endpoint_arn
        except Exception as e:
            logger.error(f"Failed to create source endpoint: {e}")
            raise

    def create_target_endpoint(self, config):
        """Create a DMS target endpoint with error handling."""
        try:
            response = self.client.create_endpoint(
                EndpointIdentifier=config['identifier'],
                EndpointType='target',
                EngineName=config['engine'],
                ServerName=config['server'],
                Port=config['port'],
                DatabaseName=config['database'],
                Username=config['username'],
                Password=config['password'],
                SslMode=config.get('ssl_mode', 'require'),
                ExtraConnectionAttributes=config.get('extra_attributes', ''),
                Tags=config.get('tags', [])
            )
            endpoint_arn = response['Endpoint']['EndpointArn']
            logger.info(f"Created target endpoint: {endpoint_arn}")
            return endpoint_arn
        except Exception as e:
            logger.error(f"Failed to create target endpoint: {e}")
            raise

    def test_endpoint_connection(self, endpoint_arn):
        """Test connectivity to an endpoint."""
        try:
            response = self.client.test_connection(
                EndpointArn=endpoint_arn,
                EndpointType='source'
            )
            status = response['ConnectionStatus']
            logger.info(f"Endpoint test status: {status}")
            return status
        except Exception as e:
            logger.error(f"Endpoint test failed: {e}")
            raise


class DMSReplicationManager:
    """Manage DMS replication instances and tasks."""

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

    def create_replication_instance(self, config):
        """Create a DMS replication instance."""
        try:
            response = self.client.create_replication_instance(
                ReplicationInstanceIdentifier=config['identifier'],
                ReplicationInstanceClass=config.get('instance_class', 'dms.r5.large'),
                AllocatedStorage=config.get('storage', 100),
                MultiAZ=config.get('multi_az', True),
                EngineVersion=config.get('engine_version', '3.5.1'),
                MinAllocatedStorage=config.get('min_storage', 50),
                MaxAllocatedStorage=config.get('max_storage', 500),
                VpcSecurityGroupIds=config.get('security_groups', []),
                ReplicationSubnetGroupIdentifier=config.get('subnet_group'),
                PubliclyAccessible=False,
                Tags=config.get('tags', [])
            )
            instance_arn = response['ReplicationInstance']['ReplicationInstanceArn']
            logger.info(f"Created replication instance: {instance_arn}")
            return instance_arn
        except Exception as e:
            logger.error(f"Failed to create replication instance: {e}")
            raise

    def create_migration_task(self, config):
        """Create a DMS migration task with full load + CDC."""
        try:
            response = self.client.create_replication_task(
                ReplicationTaskIdentifier=config['identifier'],
                SourceEndpointArn=config['source_arn'],
                TargetEndpointArn=config['target_arn'],
                ReplicationInstanceArn=config['replication_instance_arn'],
                MigrationType=config.get('migration_type', 'full-load-and-cdc'),
                TableMappings=json.dumps(config['table_mappings']),
                ReplicationTaskSettings=json.dumps(config.get('task_settings', {})),
                Tags=config.get('tags', [])
            )
            task_arn = response['ReplicationTask']['ReplicationTaskArn']
            logger.info(f"Created migration task: {task_arn}")
            return task_arn
        except Exception as e:
            logger.error(f"Failed to create migration task: {e}")
            raise

    def get_task_status(self, task_arn):
        """Get current status of a migration task."""
        try:
            response = self.client.describe_replication_tasks(
                Filters=[{'Name': 'replication-task-arn', 'Values': [task_arn]}]
            )
            tasks = response.get('ReplicationTasks', [])
            if tasks:
                status = tasks[0]['Status']
                stats = tasks[0].get('ReplicationTaskStats', {})
                logger.info(f"Task status: {status}, Latency: {stats.get('CDCLatencyTarget', 'N/A')}")
                return {'status': status, 'stats': stats}
        except Exception as e:
            logger.error(f"Failed to get task status: {e}")
            raise

import boto3
import hashlib
import logging
from datetime import datetime

logger = logging.getLogger(__name__)

class MigrationValidator:
    """Validate data integrity after database migration."""

    def __init__(self, redshift_cluster, database, db_user):
        self.redshift = boto3.client('redshift-data')
        self.cluster = redshift_cluster
        self.database = database
        self.db_user = db_user

    def execute_query(self, sql):
        """Execute SQL on Redshift and return results."""
        try:
            response = self.redshift.execute_statement(
                ClusterIdentifier=self.cluster,
                Database=self.database,
                DbUser=self.db_user,
                Sql=sql
            )
            statement_id = response['Id']

            # Wait for completion
            while True:
                status = self.redshift.describe_statement(Id=statement_id)
                if status['Status'] in ['FINISHED', 'FAILED', 'ABORTED']:
                    break

            if status['Status'] == 'FAILED':
                raise Exception(f"Query failed: {status.get('Error')}")

            return self.redshift.get_statement_result(Id=statement_id)
        except Exception as e:
            logger.error(f"Query execution failed: {e}")
            raise

    def validate_row_counts(self, source_counts):
        """Compare row counts between source and target."""
        results = []
        for table, expected_count in source_counts.items():
            result = self.execute_query(f"SELECT COUNT(*) FROM {table}")
            actual_count = int(result['Records'][0][0]['longValue'])
            match = actual_count == expected_count
            results.append({
                'table': table,
                'expected': expected_count,
                'actual': actual_count,
                'match': match
            })
            logger.info(f"{table}: expected={expected_count}, actual={actual_count}, match={match}")
        return results

    def validate_checksums(self, tables):
        """Compare checksums for critical tables."""
        results = []
        for table in tables:
            result = self.execute_query(
                f"SELECT MD5(CAST(COUNT(*) AS VARCHAR) || MAX(updated_at)) FROM {table}"
            )
            checksum = result['Records'][0][0]['stringValue']
            results.append({'table': table, 'checksum': checksum})
            logger.info(f"{table} checksum: {checksum}")
        return results

    def validate_sample_data(self, table, key_column, sample_ids):
        """Compare sample rows between source and target."""
        ids_str = ','.join(str(i) for i in sample_ids)
        result = self.execute_query(
            f"SELECT * FROM {table} WHERE {key_column} IN ({ids_str}) ORDER BY {key_column}"
        )
        return result['Records']