import boto3
import json
import logging
import time
from botocore.exceptions import ClientError

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


def create_emr_cluster(
    cluster_name: str,
    release_label: str = "emr-6.15.0",
    applications: list = None,
    instance_fleets: dict = None,
    configurations: list = None,
    service_role: str = "EMR_DefaultRole",
    log_uri: str = None
) -> str:
    """
    Create an EMR cluster with production-grade settings.

    Args:
        cluster_name: Name for the EMR cluster
        release_label: EMR release version
        applications: List of applications to install
        instance_fleets: Instance fleet configuration
        configurations: EMR configuration classifications
        service_role: IAM service role ARN
        log_uri: S3 path for cluster logs

    Returns:
        Cluster ID string
    """
    emr_client = boto3.client('emr')

    if applications is None:
        applications = [{'Name': 'Spark'}, {'Name': 'Hive'}]

    if instance_fleets is None:
        instance_fleets = {
            'MasterFleet': {
                'InstanceTypeConfigs': [
                    {'InstanceType': 'm5.xlarge', 'WeightedCapacity': 1}
                ],
                'TargetOnDemandCapacity': 1
            },
            'CoreFleet': {
                'InstanceTypeConfigs': [
                    {
                        'InstanceType': 'r5.2xlarge',
                        'WeightedCapacity': 1,
                        'EbsConfiguration': {
                            'EbsBlockDeviceConfigs': [{
                                'VolumeSpecification': {
                                    'SizeInGB': 100,
                                    'VolumeType': 'gp3'
                                },
                                'VolumesPerInstance': 2
                            }]
                        }
                    }
                ],
                'TargetOnDemandCapacity': 2,
                'TargetSpotCapacity': 2,
                'AllocationStrategy': 'CAPACITY_OPTIMIZED'
            },
            'TaskFleet': {
                'InstanceTypeConfigs': [
                    {
                        'InstanceType': 'c5.4xlarge',
                        'WeightedCapacity': 1,
                        'BidPrice': '0.50'
                    },
                    {
                        'InstanceType': 'c5.2xlarge',
                        'WeightedCapacity': 1,
                        'BidPrice': '0.30'
                    }
                ],
                'TargetSpotCapacity': 4,
                'AllocationStrategy': 'CAPACITY_OPTIMIZED'
            }
        }

    if configurations is None:
        configurations = [
            {
                'Classification': 'spark-defaults',
                'Properties': {
                    'spark.dynamicAllocation.enabled': 'true',
                    'spark.shuffle.service.enabled': 'true',
                    'spark.sql.shuffle.partitions': '200',
                    'spark.sql.adaptive.enabled': 'true',
                    'spark.sql.adaptive.coalescePartitions.enabled': 'true',
                    'spark.sql.adaptive.skewJoin.enabled': 'true',
                    'spark.executor.memory': '8g',
                    'spark.executor.cores': '4',
                    'spark.driver.memory': '4g'
                }
            },
            {
                'Classification': 'yarn-site',
                'Properties': {
                    'yarn.nodemanager.resource.memory-mb': '24576',
                    'yarn.scheduler.maximum-allocation-mb': '24576'
                }
            }
        ]

    cluster_config = {
        'Name': cluster_name,
        'ReleaseLabel': release_label,
        'Applications': applications,
        'Instances': {
            'MasterFleet': instance_fleets['MasterFleet'],
            'CoreFleet': instance_fleets['CoreFleet'],
            'TaskFleet': instance_fleets['TaskFleet'],
            'Ec2KeyName': 'emr-key-pair',
            'KeepJobFlowAliveWhenNoSteps': True,
            'TerminationProtected': True
        },
        'Configurations': configurations,
        'ServiceRole': service_role,
        'JobFlowRole': 'EMR_EC2_DefaultRole',
        'LogUri': log_uri or 's3://emr-logs-cluster/elasticmapreduce/',
        'VisibleToAllUsers': True,
        'Tags': [
            {'Key': 'Environment', 'Value': 'production'},
            {'Key': 'ManagedBy', 'Value': 'boto3'}
        ]
    }

    try:
        response = emr_client.run_job_flow(**cluster_config)
        cluster_id = response['JobFlowId']
        logger.info(f"Created EMR cluster: {cluster_id}")

        # Wait for cluster to be ready
        waiter = emr_client.get_waiter('cluster_running')
        logger.info("Waiting for cluster to reach WAITING state...")
        waiter.wait(
            ClusterId=cluster_id,
            WaiterConfig={'Delay': 30, 'MaxAttempts': 60}
        )
        logger.info(f"Cluster {cluster_id} is ready")

        return cluster_id

    except ClientError as e:
        logger.error(f"Failed to create cluster: {e.response['Error']['Message']}")
        raise


def add_emr_step(
    cluster_id: str,
    step_name: str,
    script_path: str,
    args: list = None
) -> str:
    """Add a Spark step to an EMR cluster."""
    emr_client = boto3.client('emr')

    step_config = {
        'Name': step_name,
        'ActionOnFailure': 'CONTINUE',
        'HadoopJarStep': {
            'Jar': 'command-runner.jar',
            'Args': [
                'spark-submit',
                '--deploy-mode', 'cluster',
                '--conf', 'spark.dynamicAllocation.enabled=true',
                '--conf', 'spark.sql.adaptive.enabled=true',
                script_path
            ] + (args or [])
        }
    }

    response = emr_client.add_job_flow_steps(
        JobFlowId=cluster_id,
        Steps=[step_config]
    )

    step_id = response['StepIds'][0]
    logger.info(f"Added step {step_name}: {step_id}")
    return step_id


def scale_task_nodes(cluster_id: str, target_capacity: int):
    """Scale task fleet capacity using modify-instance-fleet."""
    emr_client = boto3.client('emr')

    emr_client.modify_instance_fleet(
        ClusterId=cluster_id,
        InstanceFleet={
            'InstanceFleetType': 'TASK',
            'TargetOnDemandCapacity': 0,
            'TargetSpotCapacity': target_capacity
        }
    )
    logger.info(f"Scaled task fleet to {target_capacity} spot instances")

# Optimal Spark configuration for EMR
spark_config = {
    # Adaptive Query Execution (Spark 3.x)
    'spark.sql.adaptive.enabled': 'true',
    'spark.sql.adaptive.coalescePartitions.enabled': 'true',
    'spark.sql.adaptive.skewJoin.enabled': 'true',
    'spark.sql.adaptive.skewJoin.skewedPartitionFactor': '5',
    'spark.sql.adaptive.skewJoin.skewedPartitionThresholdInBytes': '256MB',

    # Dynamic Allocation
    'spark.dynamicAllocation.enabled': 'true',
    'spark.shuffle.service.enabled': 'true',
    'spark.dynamicAllocation.minExecutors': '2',
    'spark.dynamicAllocation.maxExecutors': '100',
    'spark.dynamicAllocation.executorIdleTimeout': '60s',

    # Shuffle Optimization
    'spark.sql.shuffle.partitions': '200',
    'spark.shuffle.compress': 'true',
    'spark.shuffle.spill.compress': 'true',
    'spark.io.compression.codec': 'snappy',

    # Memory Management
    'spark.executor.memory': '16g',
    'spark.executor.memoryOverhead': '4g',
    'spark.driver.memory': '8g',
    'spark.memory.fraction': '0.8',
    'spark.memory.storageFraction': '0.3'
}