import json
import boto3
import requests
import time
from datetime import datetime
from typing import Dict, List, Tuple, Optional
import logging
import os

# Configure logging
logger = logging.getLogger()
logger.setLevel(os.environ.get('LOG_LEVEL', 'INFO'))

# Configuration from environment variables
MAX_RETRIES = int(os.environ.get('MAX_RETRIES', '5'))
RETRY_DELAY = int(os.environ.get('RETRY_DELAY', '10'))
HEALTH_CHECK_TIMEOUT = int(os.environ.get('HEALTH_CHECK_TIMEOUT', '30'))
TEST_TIMEOUT = int(os.environ.get('TEST_TIMEOUT', '300'))
CONTAINER_PORT = int(os.environ.get('CONTAINER_PORT', '8005'))
CONTAINER_NAME = os.environ.get('CONTAINER_NAME', 'StudioApiContainer1')
ENVIRONMENT = os.environ.get('ENVIRONMENT', 'staging')

# Test endpoints
HEALTH_ENDPOINT = os.environ.get('HEALTH_ENDPOINT', '/health/')
DB_CHECK_ENDPOINT = os.environ.get('DB_CHECK_ENDPOINT', '/api/v1/db-check/')
API_STATUS_ENDPOINT = os.environ.get('API_STATUS_ENDPOINT', '/api/v1/status/')
MIGRATION_CHECK_ENDPOINT = os.environ.get('MIGRATION_CHECK_ENDPOINT', '/api/v1/migration-check/')
CACHE_CHECK_ENDPOINT = os.environ.get('CACHE_CHECK_ENDPOINT', '/api/v1/cache-check/')
STATIC_FILES_PATH = os.environ.get('STATIC_FILES_PATH', '/static/admin/css/base.css')

# Performance threshold
PERFORMANCE_THRESHOLD = float(os.environ.get('PERFORMANCE_THRESHOLD', '2.0'))


class DjangoServiceTester:
    """Handles testing of Django service before allowing traffic"""
    
    def __init__(self, deployment_id: str, lifecycle_hook_id: str):
        self.deployment_id = deployment_id
        self.lifecycle_hook_id = lifecycle_hook_id
        self.codedeploy = boto3.client('codedeploy')
        self.ecs = boto3.client('ecs')
        self.elbv2 = boto3.client('elbv2')
        self.ec2 = boto3.client('ec2')
        
    def get_deployment_info(self) -> Dict:
        """Get detailed deployment information"""
        try:
            response = self.codedeploy.get_deployment(
                deploymentId=self.deployment_id
            )
            deployment_info = response['deploymentInfo']
            logger.info(f"Deployment info retrieved: {json.dumps(deployment_info, default=str)}")
            return deployment_info
        except Exception as e:
            logger.error(f"Failed to get deployment info: {str(e)}")
            raise
    
    def get_ecs_service_from_deployment_group(self, deployment_info: Dict) -> Tuple[str, str]:
        """Extract ECS cluster and service from deployment group"""
        try:
            # Get deployment group information
            app_name = deployment_info['applicationName']
            deployment_group_name = deployment_info['deploymentGroupName']
            
            response = self.codedeploy.get_deployment_group(
                applicationName=app_name,
                deploymentGroupName=deployment_group_name
            )
            
            deployment_group = response['deploymentGroupInfo']
            
            # For ECS deployments, get the ECS service info
            if 'ecsServices' in deployment_group and deployment_group['ecsServices']:
                ecs_service = deployment_group['ecsServices'][0]
                cluster_name = ecs_service['clusterName']
                service_name = ecs_service['serviceName']
                
                logger.info(f"Found ECS service: {service_name} in cluster: {cluster_name}")
                return cluster_name, service_name
            else:
                raise Exception("No ECS service found in deployment group")
                
        except Exception as e:
            logger.error(f"Failed to get ECS service from deployment group: {str(e)}")
            raise
    
    def get_target_tasks(self, cluster_name: str, service_name: str, deployment_info: Dict) -> List[Dict]:
        """Get the new ECS tasks that were deployed"""
        try:
            # Get the green target group from deployment info
            green_target_group_name = None
            if 'loadBalancerInfo' in deployment_info:
                target_groups = deployment_info['loadBalancerInfo']['targetGroupPairInfoList'][0]['targetGroups']
                # The green target group is typically the second one
                if len(target_groups) > 1:
                    green_target_group_name = target_groups[1]['name']
                    logger.info(f"Green target group: {green_target_group_name}")
            
            # Get service details
            service_response = self.ecs.describe_services(
                cluster=cluster_name,
                services=[service_name]
            )
            
            if not service_response['services']:
                raise Exception(f"Service {service_name} not found")
            
            service = service_response['services'][0]
            
            # Get all running tasks for the service
            list_tasks_response = self.ecs.list_tasks(
                cluster=cluster_name,
                serviceName=service_name,
                desiredStatus='RUNNING'
            )
            
            if not list_tasks_response['taskArns']:
                # Try to get pending tasks if no running tasks
                list_tasks_response = self.ecs.list_tasks(
                    cluster=cluster_name,
                    serviceName=service_name,
                    desiredStatus='PENDING'
                )
                
                if not list_tasks_response['taskArns']:
                    raise Exception("No tasks found for the service")
            
            # Describe tasks to get details
            describe_tasks_response = self.ecs.describe_tasks(
                cluster=cluster_name,
                tasks=list_tasks_response['taskArns']
            )
            
            new_tasks = []
            
            # If we have a green target group, get tasks registered to it
            if green_target_group_name:
                # Get target group ARN
                tg_response = self.elbv2.describe_target_groups(
                    Names=[green_target_group_name]
                )
                
                if tg_response['TargetGroups']:
                    target_group_arn = tg_response['TargetGroups'][0]['TargetGroupArn']
                    
                    # Get targets in the green target group
                    targets_response = self.elbv2.describe_target_health(
                        TargetGroupArn=target_group_arn
                    )
                    
                    # Extract task IDs from targets
                    target_task_ids = set()
                    for target in targets_response['TargetHealthDescriptions']:
                        # Target ID format for Fargate: IP address
                        target_id = target['Target']['Id']
                        target_task_ids.add(target_id)
                    
                    # Match tasks with targets
                    for task in describe_tasks_response['tasks']:
                        if task['launchType'] == 'FARGATE':
                            # For Fargate tasks, check ENI IP addresses
                            for attachment in task.get('attachments', []):
                                if attachment['type'] == 'ElasticNetworkInterface':
                                    for detail in attachment['details']:
                                        if detail['name'] == 'privateIPv4Address':
                                            if detail['value'] in target_task_ids:
                                                new_tasks.append(task)
                                                break
                        else:
                            # For EC2 tasks, handle differently if needed
                            new_tasks.append(task)
            else:
                # If no green target group info, assume all running tasks are new
                new_tasks = [
                    task for task in describe_tasks_response['tasks']
                    if task['lastStatus'] == 'RUNNING'
                ]
            
            logger.info(f"Found {len(new_tasks)} new tasks to test")
            return new_tasks
            
        except Exception as e:
            logger.error(f"Failed to get target tasks: {str(e)}")
            raise
    
    def get_task_endpoints(self, tasks: List[Dict]) -> List[str]:
        """Get HTTP endpoints for the tasks"""
        endpoints = []
        
        for task in tasks:
            try:
                # For Fargate tasks, we need to get the ENI details
                if task.get('launchType') == 'FARGATE':
                    for attachment in task.get('attachments', []):
                        if attachment['type'] == 'ElasticNetworkInterface':
                            for detail in attachment['details']:
                                if detail['name'] == 'privateIPv4Address':
                                    private_ip = detail['value']
                                    # For Fargate, use the container port directly
                                    endpoints.append(f"http://{private_ip}:{CONTAINER_PORT}")
                                    logger.info(f"Found Fargate endpoint: http://{private_ip}:{CONTAINER_PORT}")
                                    break
                else:
                    # For EC2 launch type, get container instance details
                    container_instance_arn = task.get('containerInstanceArn')
                    if container_instance_arn:
                        # Get EC2 instance details
                        container_instances = self.ecs.describe_container_instances(
                            cluster=task['clusterArn'],
                            containerInstances=[container_instance_arn]
                        )
                        
                        if container_instances['containerInstances']:
                            ec2_instance_id = container_instances['containerInstances'][0]['ec2InstanceId']
                            
                            # Get EC2 instance details
                            ec2_instances = self.ec2.describe_instances(
                                InstanceIds=[ec2_instance_id]
                            )
                            
                            if ec2_instances['Reservations']:
                                private_ip = ec2_instances['Reservations'][0]['Instances'][0]['PrivateIpAddress']
                                
                                # Find the container port
                                for container in task['containers']:
                                    if container['name'] == CONTAINER_NAME:
                                        for network_binding in container.get('networkBindings', []):
                                            if network_binding.get('containerPort') == CONTAINER_PORT:
                                                host_port = network_binding['hostPort']
                                                endpoints.append(f"http://{private_ip}:{host_port}")
                                                logger.info(f"Found EC2 endpoint: http://{private_ip}:{host_port}")
                                                break
                
            except Exception as e:
                logger.warning(f"Failed to get endpoint for task {task.get('taskArn', 'unknown')}: {str(e)}")
                continue
        
        return endpoints
    
    def wait_for_service_ready(self, endpoints: List[str]) -> bool:
        """Wait for Django service to be ready"""
        logger.info(f"Waiting for {len(endpoints)} endpoints to be ready...")
        
        start_time = time.time()
        ready_endpoints = set()
        
        while time.time() - start_time < HEALTH_CHECK_TIMEOUT:
            for endpoint in endpoints:
                if endpoint in ready_endpoints:
                    continue
                
                try:
                    response = requests.get(
                        f"{endpoint}{HEALTH_ENDPOINT}",
                        timeout=5
                    )
                    if response.status_code == 200:
                        ready_endpoints.add(endpoint)
                        logger.info(f"Endpoint {endpoint} is ready")
                except Exception as e:
                    logger.debug(f"Endpoint {endpoint} not ready yet: {str(e)}")
            
            if len(ready_endpoints) == len(endpoints):
                logger.info("All endpoints are ready!")
                return True
            
            time.sleep(2)
        
        logger.error(f"Timeout waiting for endpoints. Ready: {len(ready_endpoints)}/{len(endpoints)}")
        return len(ready_endpoints) > 0  # Return True if at least one endpoint is ready
    
    def run_health_check(self, endpoint: str) -> Dict:
        """Run health check test"""
        logger.info(f"Running health check for {endpoint}")
        
        try:
            response = requests.get(
                f"{endpoint}{HEALTH_ENDPOINT}",
                timeout=10
            )
            
            success = response.status_code == 200
            
            if success:
                try:
                    data = response.json()
                    return {
                        'success': True,
                        'message': 'Health check passed',
                        'details': data
                    }
                except:
                    return {
                        'success': True,
                        'message': 'Health check passed',
                        'status_code': response.status_code
                    }
            else:
                return {
                    'success': False,
                    'message': f'Health check failed with status {response.status_code}',
                    'status_code': response.status_code
                }
                
        except Exception as e:
            return {
                'success': False,
                'message': f'Health check failed: {str(e)}'
            }
    
    def run_cache_check(self, endpoint: str) -> Dict:
        """Test cache connectivity"""
        logger.info(f"Running cache check for {endpoint}")
        
        try:
            response = requests.get(
                f"{endpoint}{CACHE_CHECK_ENDPOINT}",
                timeout=10
            )
            
            if response.status_code == 200:
                data = response.json()
                return {
                    'success': True,
                    'message': 'Cache connection verified',
                    'details': data
                }
            else:
                return {
                    'success': False,
                    'message': f'Cache check endpoint returned {response.status_code}'
                }
                
        except Exception as e:
            return {
                'success': False,
                'message': f'Cache check failed: {str(e)}'
            }

    def run_database_check(self, endpoint: str) -> Dict:
        """Test database connectivity"""
        logger.info(f"Running database check for {endpoint}")
        
        try:
            response = requests.get(
                f"{endpoint}{DB_CHECK_ENDPOINT}",
                timeout=10
            )
            
            if response.status_code == 200:
                data = response.json()
                return {
                    'success': True,
                    'message': 'Database connection verified',
                    'details': data
                }
            else:
                return {
                    'success': False,
                    'message': f'Database check endpoint returned {response.status_code}'
                }
                
        except Exception as e:
            return {
                'success': False,
                'message': f'Database check failed: {str(e)}'
            }
    
    def run_api_tests(self, endpoint: str) -> Dict:
        """Test critical API endpoints"""
        logger.info(f"Running API tests for {endpoint}")
        
        # Define your critical endpoints
        critical_endpoints = [
            {'path': API_STATUS_ENDPOINT, 'expected_status': [200]},
            # Add more critical endpoints as needed
        ]
        
        results = []
        all_passed = True
        
        for ep in critical_endpoints:
            try:
                response = requests.get(
                    f"{endpoint}{ep['path']}",
                    timeout=10,
                    headers={'Accept': 'application/json'}
                )
                
                passed = response.status_code in ep['expected_status']
                results.append({
                    'endpoint': ep['path'],
                    'status_code': response.status_code,
                    'passed': passed
                })
                
                if not passed:
                    all_passed = False
                    logger.error(f"API test failed for {ep['path']}: got {response.status_code}, expected {ep['expected_status']}")
                    
            except Exception as e:
                results.append({
                    'endpoint': ep['path'],
                    'error': str(e),
                    'passed': False
                })
                all_passed = False
                logger.error(f"API test failed for {ep['path']}: {str(e)}")
        
        return {
            'success': all_passed,
            'message': 'All API tests passed' if all_passed else 'Some API tests failed',
            'results': results
        }
    
    def run_performance_test(self, endpoint: str) -> Dict:
        """Basic performance test"""
        logger.info(f"Running performance test for {endpoint}")
        
        try:
            start_time = time.time()
            response = requests.get(
                f"{endpoint}{API_STATUS_ENDPOINT}",
                timeout=10
            )
            response_time = time.time() - start_time
            
            success = response.status_code == 200 and response_time < PERFORMANCE_THRESHOLD
            
            return {
                'success': success,
                'message': f'Response time: {response_time:.2f}s',
                'response_time': response_time,
                'threshold': PERFORMANCE_THRESHOLD
            }
            
        except Exception as e:
            return {
                'success': False,
                'message': f'Performance test failed: {str(e)}'
            }
    
    def run_all_tests(self, endpoints: List[str]) -> Dict:
        """Run all tests against all endpoints"""
        logger.info(f"Running all tests against {len(endpoints)} endpoints")
        
        all_results = []
        
        for endpoint in endpoints:
            logger.info(f"Testing endpoint: {endpoint}")
            
            endpoint_results = {
                'endpoint': endpoint,
                'tests': {
                    'health_check': self.run_health_check(endpoint),
                    'cache_check': self.run_cache_check(endpoint),
                    'database_check': self.run_database_check(endpoint),
                    # 'api_tests': self.run_api_tests(endpoint),
                    # 'performance': self.run_performance_test(endpoint)
                }
            }
            
            # Determine if all tests passed for this endpoint
            endpoint_results['all_passed'] = all(
                test_result['success'] 
                for test_result in endpoint_results['tests'].values()
            )
            
            all_results.append(endpoint_results)
        
        # Overall success if at least one endpoint passes all tests
        overall_success = any(result['all_passed'] for result in all_results)
        
        return {
            'success': overall_success,
            'total_endpoints': len(endpoints),
            'passed_endpoints': sum(1 for r in all_results if r['all_passed']),
            'results': all_results
        }
    
    def update_deployment_status(self, status: str, message: str = None):
        """Update CodeDeploy with the deployment status"""
        try:
            self.codedeploy.put_lifecycle_event_hook_execution_status(
                deploymentId=self.deployment_id,
                lifecycleEventHookExecutionId=self.lifecycle_hook_id,
                status=status
            )
            logger.info(f"Updated deployment status to: {status}")
            
        except Exception as e:
            logger.error(f"Failed to update deployment status: {str(e)}")
            raise


def lambda_handler(event, context):
    """
    Main Lambda handler for BeforeAllowTraffic hook
    """
    logger.info(f"BeforeAllowTraffic hook triggered: {json.dumps(event)}")
    
    # Extract deployment information
    deployment_id = event['DeploymentId']
    lifecycle_hook_id = event['LifecycleEventHookExecutionId']
    
    # Initialize tester
    tester = DjangoServiceTester(deployment_id, lifecycle_hook_id)
    
    try:
        # Get deployment information
        deployment_info = tester.get_deployment_info()
        
        # Extract ECS service information
        cluster_name, service_name = tester.get_ecs_service_from_deployment_group(deployment_info)
        
        # Get the new tasks
        new_tasks = tester.get_target_tasks(cluster_name, service_name, deployment_info)
        logger.info(f"Found {len(new_tasks)} new tasks to test")
        
        if not new_tasks:
            # If no tasks found, wait a bit and retry once
            logger.warning("No new tasks found, waiting 10 seconds and retrying...")
            time.sleep(10)
            new_tasks = tester.get_target_tasks(cluster_name, service_name, deployment_info)
            
            if not new_tasks:
                raise Exception("No new tasks found to test after retry")
        
        # Get endpoints for the tasks
        endpoints = tester.get_task_endpoints(new_tasks)
        logger.info(f"Found {len(endpoints)} endpoints to test")
        
        if not endpoints:
            raise Exception("No endpoints found to test")
        
        # Wait for services to be ready
        if not tester.wait_for_service_ready(endpoints):
            logger.warning("Not all services became ready, but proceeding with available endpoints")
        
        # Run all tests
        test_results = tester.run_all_tests(endpoints)
        
        # Log detailed results
        logger.info(f"Test results: {json.dumps(test_results, indent=2)}")
        
        if test_results['success']:
            logger.info("All tests passed! Proceeding with deployment.")
            tester.update_deployment_status('Succeeded')
            
            return {
                'statusCode': 200,
                'body': json.dumps({
                    'message': 'All tests passed',
                    'results': test_results
                })
            }
        else:
            logger.error("Tests failed! Rolling back deployment.")
            tester.update_deployment_status('Failed')
            
            return {
                'statusCode': 400,
                'body': json.dumps({
                    'message': 'Tests failed',
                    'results': test_results
                })
            }
            
    except Exception as e:
        logger.error(f"Error during BeforeAllowTraffic hook: {str(e)}", exc_info=True)
        
        try:
            tester.update_deployment_status('Failed')
        except:
            pass
        
        return {
            'statusCode': 500,
            'body': json.dumps({
                'message': f'Error during testing: {str(e)}'
            })
        }