import time

import boto3
import logging
from botocore.exceptions import ClientError
from typing import List, Dict

# Set up logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class ECSRedeployer:
    def __init__(self, region: str = None):
        self.ecs_client = boto3.client('ecs', region_name=region)
        self.region = region or boto3.session.Session().region_name

    def get_all_clusters(self) -> List[str]:
        """Retrieve all ECS cluster ARNs in the account."""
        try:
            clusters = []
            paginator = self.ecs_client.get_paginator('list_clusters')

            for page in paginator.paginate():
                clusters.extend(page['clusterArns'])

            return clusters
        except ClientError as e:
            logger.error(f"Failed to get clusters: {str(e)}")
            raise

    def get_services_for_cluster(self, cluster_arn: str) -> List[str]:
        """Get all services in a specific cluster."""
        try:
            services = []
            paginator = self.ecs_client.get_paginator('list_services')

            for page in paginator.paginate(cluster=cluster_arn):
                services.extend(page['serviceArns'])

            return services
        except ClientError as e:
            logger.error(f"Failed to get services for cluster {cluster_arn}: {str(e)}")
            raise

    def force_redeploy_service(self, cluster_arn: str, service_arn: str) -> bool:
        """Force redeploy a specific service by updating its task definition."""
        try:
            # Get current service details
            service = self.ecs_client.describe_services(
                cluster=cluster_arn,
                services=[service_arn]
            )['services'][0]

            task_definition = service['taskDefinition']

            # Update the service to force new deployment
            response = self.ecs_client.update_service(
                cluster=cluster_arn,
                service=service_arn,
                taskDefinition=task_definition,
                forceNewDeployment=True
            )

            logger.info(f"Initiated force redeploy for service {service_arn} in cluster {cluster_arn}")
            return True

        except ClientError as e:
            logger.error(f"Failed to redeploy service {service_arn} in cluster {cluster_arn}: {str(e)}")
            return False

    def redeploy_all_services(self) -> Dict[str, Dict[str, bool]]:
        """Force redeploy all services in all clusters."""
        results = {}

        try:
            clusters = self.get_all_clusters()
            logger.info(f"Found {len(clusters)} clusters")

            for cluster_arn in clusters:
                cluster_name = cluster_arn.split('/')[-1]
                results[cluster_name] = {}

                services = self.get_services_for_cluster(cluster_arn)
                logger.info(f"Found {len(services)} services in cluster {cluster_name}")

                for service_arn in services:
                    service_name = service_arn.split('/')[-1]
                    success = self.force_redeploy_service(cluster_arn, service_arn)
                    results[cluster_name][service_name] = success

                    # Add a small delay to avoid API throttling
                    time.sleep(1.5)

            return results

        except Exception as e:
            logger.error(f"An error occurred during redeployment: {str(e)}")
            raise


def main():
    try:
        # Initialize the redeployer
        redeployer = ECSRedeployer()

        # Start the redeployment process
        logger.info("Starting force redeployment of all ECS services...")
        results = redeployer.redeploy_all_services()

        # Print summary
        logger.info("\nRedeployment Summary:")
        for cluster, services in results.items():
            successful = sum(1 for s in services.values() if s)
            total = len(services)
            logger.info(f"Cluster {cluster}: {successful}/{total} services redeployed successfully")

            # Print failed services if any
            failed_services = [s for s, success in services.items() if not success]
            if failed_services:
                logger.warning(f"Failed services in {cluster}: {', '.join(failed_services)}")

    except Exception as e:
        logger.error(f"Script execution failed: {str(e)}")
        raise


if __name__ == "__main__":
    main()