import requests_mock
import unittest
from unittest.mock import patch

from fargate.run_fargate_task import get_vpc_id, get_security_group, get_subnets, run_task

class TestRunFargateTask(unittest.TestCase):
    @patch('fargate.run_fargate_task.ENVIRONMENT', 'dev')
    @patch('fargate.run_fargate_task.VPC_NAME', 'dev')
    @patch('boto3.client')
    def test_get_vpc_id_success(self, mock_boto_client):
        """Test VPC ID retrieval"""
        # Mock the EC2 client response
        mock_boto_client.return_value.describe_vpcs.return_value = {
            'Vpcs': [
                {
                    'VpcId': 'vpc-12345678',
                    'Tags': [
                        {
                            'Key': 'tag:Name',
                            'Value': 'dev'
                        },
                    ]
                }
            ]
        }

        # Mock the EC2 client
        mock_ec2_client = mock_boto_client.return_value


        vpc_id = get_vpc_id(mock_ec2_client)
        self.assertEqual(vpc_id, 'vpc-12345678')

    @patch('fargate.run_fargate_task.ENVIRONMENT', 'qa')
    @patch('fargate.run_fargate_task.VPC_NAME', 'qa')
    @patch('boto3.client')
    def test_get_vpc_id_failure(self, mock_boto_client):
        """Test VPC ID retrieval failure"""
        # Mock the EC2 client response
        mock_boto_client.return_value.describe_vpcs.return_value = {
            'Vpcs': []
        }

        # Mock the EC2 client
        mock_ec2_client = mock_boto_client.return_value

        with self.assertRaises(SystemExit):
            get_vpc_id(mock_ec2_client)


    @patch('fargate.run_fargate_task.ENVIRONMENT', 'dev')
    @patch('fargate.run_fargate_task.SERVICE_NAME', 'test_service')
    @patch('boto3.client')
    def test_get_security_group_success(self, mock_boto_client):
        """Test security group retrieval"""
        # Mock the EC2 client response
        mock_boto_client.return_value.describe_security_groups.return_value = {
            'SecurityGroups': [
                {
                    'GroupId': 'sg-12345678',
                    'GroupName': 'dev-test_service-task-security-group',
                    'VpcId': 'vpc-12345678'
                }
            ]
        }

        # Mock the EC2 client
        mock_ec2_client = mock_boto_client.return_value

        security_group_id = get_security_group(mock_ec2_client, 'vpc-12345678')
        self.assertEqual(security_group_id, 'sg-12345678')

    @patch('fargate.run_fargate_task.ENVIRONMENT', 'qa')
    @patch('fargate.run_fargate_task.SERVICE_NAME', 'test_service')
    @patch('boto3.client')
    def test_get_security_group_failure(self, mock_boto_client):
        """Test security group retrieval failure"""
        # Mock the EC2 client response
        mock_boto_client.return_value.describe_security_groups.return_value = {
            'SecurityGroups': []
        }

        # Mock the EC2 client
        mock_ec2_client = mock_boto_client.return_value

        with self.assertRaises(SystemExit):
            get_security_group(mock_ec2_client, 'vpc-12345678')

    @patch('boto3.client')
    def test_get_subnets_success(self, mock_boto_client):
        """Test subnet retrieval"""
        # Mock the EC2 client response
        mock_boto_client.return_value.describe_subnets.return_value = {
            'Subnets': [
                {
                    'SubnetId': 'subnet-12345678',
                    'VpcId': 'vpc-12345678',
                    'AvailabilityZoneId': 'use1-az1',
                    'AvailabilityZone': 'us-east-1a',
                    'AvailableIpAddressCount': 254,
                },
                {
                    'SubnetId': 'subnet-23456789',
                    'VpcId': 'vpc-12345678',
                    'AvailabilityZoneId': 'use1-az2',
                    'AvailabilityZone': 'us-east-1b',
                    'AvailableIpAddressCount': 253,
                },
                {
                    'SubnetId': 'subnet-34567890',
                    'VpcId': 'vpc-12345678',
                    'AvailabilityZoneId': 'use1-az3',
                    'AvailabilityZone': 'us-east-1c',
                    'AvailableIpAddressCount': 4091,
                },
                {
                    'SubnetId': 'subnet-45678901',
                    'VpcId': 'vpc-12345678',
                    'AvailabilityZoneId': 'use1-az4',
                    'AvailabilityZone': 'us-east-1d',
                    'AvailableIpAddressCount': 256,
                }
            ]
        }

        # Mock the EC2 client
        mock_ec2_client = mock_boto_client.return_value

        subnets = get_subnets(mock_ec2_client, 'vpc-12345678')
        self.assertEqual(subnets, ['subnet-45678901', 'subnet-12345678'])

    @patch('boto3.client')
    def test_get_subnets_failure(self, mock_boto_client):
        """Test subnet retrieval failure"""
        # Mock the EC2 client response
        mock_boto_client.return_value.describe_subnets.return_value = {
            'Subnets': []
        }

        # Mock the EC2 client
        mock_ec2_client = mock_boto_client.return_value

        with self.assertRaises(SystemExit):
            get_subnets(mock_ec2_client, 'vpc-12345678')

    @patch('fargate.run_fargate_task.ENVIRONMENT', 'dev')
    @patch('fargate.run_fargate_task.VPC_NAME', 'dev')
    @patch('fargate.run_fargate_task.SERVICE_NAME', 'test_service')
    @patch('fargate.run_fargate_task.CLUSTER_NAME', 'test_service')
    @patch('fargate.run_fargate_task.NUM_TASKS_TO_RUN', 1)
    @patch('fargate.run_fargate_task.TASK_STARTED_BY', 'jenkins')
    @patch('boto3.client')
    def test_run_task_success(self, mock_boto_client):
        """Test task execution"""
        # Mock the EC2 client response
        mock_boto_client.return_value.describe_vpcs.return_value = {
            'Vpcs': [
                {
                    'VpcId': 'vpc-12345678',
                    'Tags': [
                        {
                            'Key': 'tag:Name',
                            'Value': 'dev'
                        },
                    ]
                }
            ]
        }

        # Mock the EC2 client response
        mock_boto_client.return_value.describe_security_groups.return_value = {
            'SecurityGroups': [
                {
                    'GroupId': 'sg-12345678',
                    'GroupName': 'dev-test_service-task-security-group',
                    'VpcId': 'vpc-12345678'
                }
            ]
        }

        # Mock the EC2 client response
        mock_boto_client.return_value.describe_subnets.return_value = {
            'Subnets': [
                {
                    'SubnetId': 'subnet-12345678',
                    'VpcId': 'vpc-12345678',
                    'AvailabilityZoneId': 'use1-az1',
                    'AvailabilityZone': 'us-east-1a',
                    'AvailableIpAddressCount': 254,
                },
                {
                    'SubnetId': 'subnet-23456789',
                    'VpcId': 'vpc-12345678',
                    'AvailabilityZoneId': 'use1-az2',
                    'AvailabilityZone': 'us-east-1b',
                    'AvailableIpAddressCount': 253,
                },
                {
                    'SubnetId': 'subnet-34567890',
                    'VpcId': 'vpc-12345678',
                    'AvailabilityZoneId': 'use1-az3',
                    'AvailabilityZone': 'us-east-1c',
                    'AvailableIpAddressCount': 4091,
                },
                {
                    'SubnetId': 'subnet-45678901',
                    'VpcId': 'vpc-12345678',
                    'AvailabilityZoneId': 'use1-az4',
                    'AvailabilityZone': 'us-east-1d',
                    'AvailableIpAddressCount': 256,
                }
            ]
        }

        # Mock the ECS client response
        mock_boto_client.return_value.run_task.return_value = {
            'tasks': [
                {
                    'attachments': [
                        {
                            'id': 'attachment-id-123',
                            'type': 'ElasticNetworkInterface',
                            'status': 'ATTACHED',
                            'details': [
                                {
                                    'name': 'subnetId',
                                    'value': 'subnet-45678901'
                                },
                                {
                                    'name': 'networkInterfaceId',
                                    'value': 'eni-12345678'
                                }
                            ]
                        }
                    ],
                    'attributes': [
                        {
                            'name': 'ecs.availability-zone',
                            'value': 'us-east-1d',
                            'targetType': 'container-instance',
                            'targetId': 'container-instance-123'
                        }
                    ],
                    'availabilityZone': 'us-east-1d',
                    'capacityProviderName': 'FARGATE',
                    'clusterArn': 'arn:aws:ecs:us-east-1:123456789012:cluster/test_cluster',
                    'connectivity': 'CONNECTED',
                    'containers': [
                        {
                            'containerArn': 'arn:aws:ecs:us-east-1:123456789012:container/container-123',
                            'taskArn': 'arn:aws:ecs:us-east-1:123456789012:task/task-123',
                            'name': 'test-container',
                            'image': 'nginx:latest',
                            'imageDigest': 'sha256:abcdef1234567890abcdef1234567890abcdef1234567890abcdef1234567890',
                            'runtimeId': 'runtime-123',
                            'lastStatus': 'RUNNING',
                            'exitCode': 0,
                            'reason': '',
                            'networkBindings': [
                                {
                                    'bindIP': '0.0.0.0',
                                    'containerPort': 80,
                                    'hostPort': 8080,
                                    'protocol': 'tcp',
                                    'containerPortRange': '80',
                                    'hostPortRange': '8080'
                                }
                            ],
                            'networkInterfaces': [
                                {
                                    'attachmentId': 'attachment-id-123',
                                    'privateIpv4Address': '10.0.0.123',
                                    'ipv6Address': '::1'
                                }
                            ],
                            'healthStatus': 'HEALTHY',
                            'managedAgents': [
                                {
                                    'name': 'ExecuteCommandAgent',
                                    'reason': 'Running',
                                    'lastStatus': 'ACTIVE'
                                }
                            ],
                            'cpu': '256',
                            'memory': '512',
                            'memoryReservation': '256',
                            'gpuIds': []
                        }
                    ],
                    'cpu': '256',
                    'desiredStatus': 'RUNNING',
                    'enableExecuteCommand': True,
                    'executionStoppedAt': None,
                    'group': 'family:service:test-service',
                    'healthStatus': 'HEALTHY',
                    'inferenceAccelerators': [],
                    'lastStatus': 'RUNNING',
                    'launchType': 'FARGATE',
                    'memory': '512',
                    'platformVersion': 'LATEST',
                    'platformFamily': 'Linux',
                    'startedBy': 'jenkins',
                    'stopCode': '',
                    'stoppedAt': None,
                    'stoppedReason': '',
                    'stoppingAt': None,
                    'tags': [
                        {'key': 'environment', 'value': 'dev'},
                        {'key': 'service_name', 'value': 'test_service'}
                    ],
                    'taskArn': 'arn:aws:ecs:us-east-1:123456789012:task/test_service/task-123',
                    'taskDefinitionArn': 'arn:aws:ecs:us-east-1:123456789012:task-definition/test-task:1',
                    'version': 1
                }
            ],
            'failures': []
        }

        # Mock the EC2 client
        mock_ec2_client = mock_boto_client.return_value

        # Get VPC ID
        vpc_id = get_vpc_id(mock_ec2_client)

        # Get security group
        security_group_id = get_security_group(mock_ec2_client, vpc_id)

        # Get subnets
        subnets = get_subnets(mock_ec2_client, vpc_id)

        # Mock the ECS client
        mock_ecs_client = mock_boto_client.return_value

        task_arn = run_task(mock_ecs_client, security_group_id, subnets, [], [])

        self.assertEqual(task_arn, 'arn:aws:ecs:us-east-1:123456789012:task/test_service/task-123')