import boto3
import json
import logging
import asyncio
from datetime import datetime, timedelta
from typing import Dict, List, Optional, Any, Tuple
from dataclasses import dataclass, asdict
from enum import Enum
import math
import statistics

class LocationCriteria(Enum):
    LATENCY = "latency"
    COMPLIANCE = "compliance"
    COST = "cost"
    AVAILABILITY = "availability"
    CAPACITY = "capacity"

class ComplianceRequirement(Enum):
    GDPR = "gdpr"
    HIPAA = "hipaa"
    SOX = "sox"
    PCI_DSS = "pci_dss"
    DATA_RESIDENCY = "data_residency"

@dataclass
class LocationProfile:
    region_name: str
    availability_zones: List[str]
    geographic_location: Dict[str, float]  # lat, lng
    compliance_certifications: List[ComplianceRequirement]
    cost_factors: Dict[str, float]
    network_performance: Dict[str, float]
    service_availability: Dict[str, float]
    capacity_limits: Dict[str, int]
    last_updated: datetime

@dataclass
class UserLocation:
    location_id: str
    geographic_location: Dict[str, float]
    user_count: int
    compliance_requirements: List[ComplianceRequirement]
    latency_requirements: Dict[str, float]  # max acceptable latency
    data_classification: str

@dataclass
class LocationRecommendation:
    region_name: str
    score: float
    criteria_scores: Dict[str, float]
    estimated_latency: float
    compliance_match: bool
    estimated_cost: float
    capacity_available: bool
    recommendation_reason: str

class LocationSelectionEngine:
    """Intelligent location selection system for multi-location deployments"""
    
    def __init__(self, config: Dict[str, Any]):
        self.config = config
        
        # AWS clients
        self.ec2 = boto3.client('ec2')
        self.pricing = boto3.client('pricing', region_name='us-east-1')
        self.cloudwatch = boto3.client('cloudwatch')
        self.organizations = boto3.client('organizations')
        self.dynamodb = boto3.resource('dynamodb')
        
        # Storage
        self.location_profiles_table = self.dynamodb.Table(config.get('location_profiles_table', 'location-profiles'))
        self.user_locations_table = self.dynamodb.Table(config.get('user_locations_table', 'user-locations'))
        
        # Location data
        self.location_profiles = {}
        self.user_locations = {}
        
        # Load location data
        asyncio.create_task(self._load_location_data())
        
    async def _load_location_data(self):
        """Load location profiles and user location data"""
        try:
            # Load AWS region information
            await self._load_aws_regions()
            
            # Load user location data
            await self._load_user_locations()
            
            # Update location profiles with real-time data
            await self._update_location_profiles()
            
            logging.info(f"Loaded {len(self.location_profiles)} location profiles")
            
        except Exception as e:
            logging.error(f"Failed to load location data: {str(e)}")
    
    async def _load_aws_regions(self):
        """Load AWS region information and create location profiles"""
        try:
            # Get all available regions
            regions_response = self.ec2.describe_regions()
            
            for region in regions_response['Regions']:
                region_name = region['RegionName']
                
                # Get availability zones for region
                ec2_regional = boto3.client('ec2', region_name=region_name)
                azs_response = ec2_regional.describe_availability_zones()
                
                availability_zones = [az['ZoneName'] for az in azs_response['AvailabilityZones']]
                
                # Create location profile
                profile = LocationProfile(
                    region_name=region_name,
                    availability_zones=availability_zones,
                    geographic_location=self._get_region_coordinates(region_name),
                    compliance_certifications=self._get_region_compliance(region_name),
                    cost_factors=await self._get_region_cost_factors(region_name),
                    network_performance={},
                    service_availability={},
                    capacity_limits={},
                    last_updated=datetime.utcnow()
                )
                
                self.location_profiles[region_name] = profile
                
                # Store in database
                await self._store_location_profile(profile)
            
        except Exception as e:
            logging.error(f"Failed to load AWS regions: {str(e)}")
    
    def _get_region_coordinates(self, region_name: str) -> Dict[str, float]:
        """Get approximate coordinates for AWS region"""
        # Simplified mapping - in practice, this would be more comprehensive
        region_coordinates = {
            'us-east-1': {'lat': 39.0458, 'lng': -77.5081},      # N. Virginia
            'us-west-2': {'lat': 45.5152, 'lng': -122.6784},     # Oregon
            'eu-west-1': {'lat': 53.3498, 'lng': -6.2603},       # Ireland
            'eu-central-1': {'lat': 50.1109, 'lng': 8.6821},     # Frankfurt
            'ap-southeast-1': {'lat': 1.3521, 'lng': 103.8198},  # Singapore
            'ap-northeast-1': {'lat': 35.6762, 'lng': 139.6503}, # Tokyo
            'ap-south-1': {'lat': 19.0760, 'lng': 72.8777},      # Mumbai
            'sa-east-1': {'lat': -23.5505, 'lng': -46.6333},     # São Paulo
        }
        
        return region_coordinates.get(region_name, {'lat': 0.0, 'lng': 0.0})
    
    def _get_region_compliance(self, region_name: str) -> List[ComplianceRequirement]:
        """Get compliance certifications for region"""
        # Simplified mapping - in practice, this would be more comprehensive
        compliance_mapping = {
            'us-east-1': [ComplianceRequirement.HIPAA, ComplianceRequirement.SOX, ComplianceRequirement.PCI_DSS],
            'us-west-2': [ComplianceRequirement.HIPAA, ComplianceRequirement.SOX, ComplianceRequirement.PCI_DSS],
            'eu-west-1': [ComplianceRequirement.GDPR, ComplianceRequirement.PCI_DSS, ComplianceRequirement.DATA_RESIDENCY],
            'eu-central-1': [ComplianceRequirement.GDPR, ComplianceRequirement.PCI_DSS, ComplianceRequirement.DATA_RESIDENCY],
            'ap-southeast-1': [ComplianceRequirement.PCI_DSS],
            'ap-northeast-1': [ComplianceRequirement.PCI_DSS],
        }
        
        return compliance_mapping.get(region_name, [])
    
    async def _get_region_cost_factors(self, region_name: str) -> Dict[str, float]:
        """Get cost factors for region"""
        try:
            # This would typically query AWS Pricing API
            # For now, we'll use simplified cost factors
            cost_factors = {
                'compute_multiplier': 1.0,
                'storage_multiplier': 1.0,
                'network_multiplier': 1.0
            }
            
            # Regional cost adjustments (simplified)
            regional_adjustments = {
                'us-east-1': {'compute_multiplier': 1.0, 'storage_multiplier': 1.0},
                'us-west-2': {'compute_multiplier': 1.1, 'storage_multiplier': 1.05},
                'eu-west-1': {'compute_multiplier': 1.15, 'storage_multiplier': 1.1},
                'eu-central-1': {'compute_multiplier': 1.12, 'storage_multiplier': 1.08},
                'ap-southeast-1': {'compute_multiplier': 1.2, 'storage_multiplier': 1.15},
                'ap-northeast-1': {'compute_multiplier': 1.25, 'storage_multiplier': 1.2},
            }
            
            if region_name in regional_adjustments:
                cost_factors.update(regional_adjustments[region_name])
            
            return cost_factors
            
        except Exception as e:
            logging.error(f"Failed to get cost factors for {region_name}: {str(e)}")
            return {'compute_multiplier': 1.0, 'storage_multiplier': 1.0, 'network_multiplier': 1.0}
    
    async def select_optimal_locations(self, requirements: Dict[str, Any]) -> List[LocationRecommendation]:
        """Select optimal locations based on requirements"""
        try:
            user_locations = requirements.get('user_locations', [])
            compliance_requirements = requirements.get('compliance_requirements', [])
            performance_requirements = requirements.get('performance_requirements', {})
            cost_constraints = requirements.get('cost_constraints', {})
            location_count = requirements.get('location_count', 3)
            
            # Score all available locations
            location_scores = []
            
            for region_name, profile in self.location_profiles.items():
                score = await self._calculate_location_score(
                    profile, user_locations, compliance_requirements, 
                    performance_requirements, cost_constraints
                )
                location_scores.append((region_name, score))
            
            # Sort by score and select top locations
            location_scores.sort(key=lambda x: x[1].score, reverse=True)
            
            # Select diverse locations (avoid same geographic area)
            selected_locations = []
            selected_regions = set()
            
            for region_name, recommendation in location_scores:
                if len(selected_locations) >= location_count:
                    break
                
                # Check geographic diversity
                if self._is_geographically_diverse(region_name, selected_regions):
                    selected_locations.append(recommendation)
                    selected_regions.add(region_name)
            
            # If we don't have enough diverse locations, fill with best remaining
            while len(selected_locations) < location_count and len(selected_locations) < len(location_scores):
                for region_name, recommendation in location_scores:
                    if region_name not in selected_regions:
                        selected_locations.append(recommendation)
                        selected_regions.add(region_name)
                        break
            
            logging.info(f"Selected {len(selected_locations)} optimal locations")
            return selected_locations
            
        except Exception as e:
            logging.error(f"Failed to select optimal locations: {str(e)}")
            return []
    
    async def _calculate_location_score(self, profile: LocationProfile, user_locations: List[Dict[str, Any]],
                                      compliance_requirements: List[str], performance_requirements: Dict[str, Any],
                                      cost_constraints: Dict[str, Any]) -> LocationRecommendation:
        """Calculate score for a specific location"""
        try:
            criteria_scores = {}
            
            # Calculate latency score
            latency_score = await self._calculate_latency_score(profile, user_locations, performance_requirements)
            criteria_scores['latency'] = latency_score
            
            # Calculate compliance score
            compliance_score = self._calculate_compliance_score(profile, compliance_requirements)
            criteria_scores['compliance'] = compliance_score
            
            # Calculate cost score
            cost_score = self._calculate_cost_score(profile, cost_constraints)
            criteria_scores['cost'] = cost_score
            
            # Calculate availability score
            availability_score = await self._calculate_availability_score(profile)
            criteria_scores['availability'] = availability_score
            
            # Calculate capacity score
            capacity_score = self._calculate_capacity_score(profile, performance_requirements)
            criteria_scores['capacity'] = capacity_score
            
            # Calculate weighted overall score
            weights = {
                'latency': 0.3,
                'compliance': 0.25,
                'cost': 0.2,
                'availability': 0.15,
                'capacity': 0.1
            }
            
            overall_score = sum(criteria_scores[criteria] * weights[criteria] for criteria in weights)
            
            # Estimate latency and cost
            estimated_latency = await self._estimate_average_latency(profile, user_locations)
            estimated_cost = self._estimate_monthly_cost(profile, performance_requirements)
            
            return LocationRecommendation(
                region_name=profile.region_name,
                score=overall_score,
                criteria_scores=criteria_scores,
                estimated_latency=estimated_latency,
                compliance_match=compliance_score > 0.8,
                estimated_cost=estimated_cost,
                capacity_available=capacity_score > 0.7,
                recommendation_reason=self._generate_recommendation_reason(criteria_scores, weights)
            )
            
        except Exception as e:
            logging.error(f"Failed to calculate location score: {str(e)}")
            return LocationRecommendation(
                region_name=profile.region_name,
                score=0.0,
                criteria_scores={},
                estimated_latency=1000.0,
                compliance_match=False,
                estimated_cost=0.0,
                capacity_available=False,
                recommendation_reason="Calculation failed"
            )
    
    async def _calculate_latency_score(self, profile: LocationProfile, user_locations: List[Dict[str, Any]],
                                     performance_requirements: Dict[str, Any]) -> float:
        """Calculate latency score for location"""
        try:
            if not user_locations:
                return 0.5  # Neutral score if no user location data
            
            total_weighted_latency = 0
            total_weight = 0
            
            for user_location in user_locations:
                # Calculate distance-based latency estimate
                distance = self._calculate_distance(
                    profile.geographic_location,
                    user_location['geographic_location']
                )
                
                # Estimate latency based on distance (simplified)
                estimated_latency = distance * 0.01  # ~10ms per 1000km
                
                # Weight by user count
                weight = user_location.get('user_count', 1)
                total_weighted_latency += estimated_latency * weight
                total_weight += weight
            
            average_latency = total_weighted_latency / total_weight if total_weight > 0 else 1000
            
            # Score based on latency requirements
            max_acceptable_latency = performance_requirements.get('max_latency_ms', 100)
            
            if average_latency <= max_acceptable_latency:
                return 1.0
            elif average_latency <= max_acceptable_latency * 2:
                return 1.0 - (average_latency - max_acceptable_latency) / max_acceptable_latency
            else:
                return 0.0
                
        except Exception as e:
            logging.error(f"Failed to calculate latency score: {str(e)}")
            return 0.0
    
    def _calculate_distance(self, loc1: Dict[str, float], loc2: Dict[str, float]) -> float:
        """Calculate distance between two geographic locations"""
        try:
            # Haversine formula for great circle distance
            lat1, lng1 = math.radians(loc1['lat']), math.radians(loc1['lng'])
            lat2, lng2 = math.radians(loc2['lat']), math.radians(loc2['lng'])
            
            dlat = lat2 - lat1
            dlng = lng2 - lng1
            
            a = math.sin(dlat/2)**2 + math.cos(lat1) * math.cos(lat2) * math.sin(dlng/2)**2
            c = 2 * math.asin(math.sqrt(a))
            
            # Earth's radius in kilometers
            r = 6371
            
            return c * r
            
        except Exception as e:
            logging.error(f"Failed to calculate distance: {str(e)}")
            return 10000  # Large distance as fallback
    
    def _calculate_compliance_score(self, profile: LocationProfile, requirements: List[str]) -> float:
        """Calculate compliance score for location"""
        try:
            if not requirements:
                return 1.0  # No requirements means full compliance
            
            profile_compliance = set(req.value for req in profile.compliance_certifications)
            required_compliance = set(requirements)
            
            # Calculate overlap
            matching_requirements = profile_compliance.intersection(required_compliance)
            
            return len(matching_requirements) / len(required_compliance)
            
        except Exception as e:
            logging.error(f"Failed to calculate compliance score: {str(e)}")
            return 0.0
    
    def _calculate_cost_score(self, profile: LocationProfile, constraints: Dict[str, Any]) -> float:
        """Calculate cost score for location"""
        try:
            max_budget = constraints.get('max_monthly_budget', float('inf'))
            
            if max_budget == float('inf'):
                return 1.0  # No budget constraint
            
            # Estimate cost based on cost factors
            base_cost = constraints.get('base_monthly_cost', 1000)
            estimated_cost = base_cost * profile.cost_factors.get('compute_multiplier', 1.0)
            
            if estimated_cost <= max_budget:
                # Score based on cost efficiency
                return 1.0 - (estimated_cost / max_budget) * 0.5
            else:
                return 0.0  # Over budget
                
        except Exception as e:
            logging.error(f"Failed to calculate cost score: {str(e)}")
            return 0.5
    
    async def _calculate_availability_score(self, profile: LocationProfile) -> float:
        """Calculate availability score for location"""
        try:
            # This would typically query historical availability data
            # For now, we'll use simplified scoring based on AZ count
            az_count = len(profile.availability_zones)
            
            if az_count >= 3:
                return 1.0
            elif az_count == 2:
                return 0.8
            else:
                return 0.5
                
        except Exception as e:
            logging.error(f"Failed to calculate availability score: {str(e)}")
            return 0.5
    
    def _calculate_capacity_score(self, profile: LocationProfile, requirements: Dict[str, Any]) -> float:
        """Calculate capacity score for location"""
        try:
            # This would typically check actual capacity limits
            # For now, we'll use simplified scoring
            required_capacity = requirements.get('required_capacity', {})
            
            if not required_capacity:
                return 1.0
            
            # Simplified capacity check
            return 0.9  # Assume most regions have adequate capacity
            
        except Exception as e:
            logging.error(f"Failed to calculate capacity score: {str(e)}")
            return 0.5
    
    def _is_geographically_diverse(self, region_name: str, selected_regions: set) -> bool:
        """Check if region provides geographic diversity"""
        try:
            if not selected_regions:
                return True
            
            current_location = self.location_profiles[region_name].geographic_location
            
            for selected_region in selected_regions:
                selected_location = self.location_profiles[selected_region].geographic_location
                distance = self._calculate_distance(current_location, selected_location)
                
                # Require at least 1000km separation
                if distance < 1000:
                    return False
            
            return True
            
        except Exception as e:
            logging.error(f"Failed to check geographic diversity: {str(e)}")
            return True
    
    async def _estimate_average_latency(self, profile: LocationProfile, user_locations: List[Dict[str, Any]]) -> float:
        """Estimate average latency from location to users"""
        try:
            if not user_locations:
                return 50.0  # Default estimate
            
            total_weighted_latency = 0
            total_weight = 0
            
            for user_location in user_locations:
                distance = self._calculate_distance(
                    profile.geographic_location,
                    user_location['geographic_location']
                )
                
                estimated_latency = distance * 0.01  # ~10ms per 1000km
                weight = user_location.get('user_count', 1)
                
                total_weighted_latency += estimated_latency * weight
                total_weight += weight
            
            return total_weighted_latency / total_weight if total_weight > 0 else 50.0
            
        except Exception as e:
            logging.error(f"Failed to estimate average latency: {str(e)}")
            return 100.0
    
    def _estimate_monthly_cost(self, profile: LocationProfile, requirements: Dict[str, Any]) -> float:
        """Estimate monthly cost for location"""
        try:
            base_cost = requirements.get('base_monthly_cost', 1000)
            compute_multiplier = profile.cost_factors.get('compute_multiplier', 1.0)
            storage_multiplier = profile.cost_factors.get('storage_multiplier', 1.0)
            
            # Simplified cost calculation
            estimated_cost = base_cost * ((compute_multiplier + storage_multiplier) / 2)
            
            return estimated_cost
            
        except Exception as e:
            logging.error(f"Failed to estimate monthly cost: {str(e)}")
            return 1000.0
    
    def _generate_recommendation_reason(self, criteria_scores: Dict[str, float], weights: Dict[str, float]) -> str:
        """Generate human-readable recommendation reason"""
        try:
            # Find the highest scoring criteria
            top_criteria = max(criteria_scores.items(), key=lambda x: x[1] * weights[x[0]])
            
            reasons = {
                'latency': 'Excellent latency performance for user base',
                'compliance': 'Strong compliance certification match',
                'cost': 'Cost-effective option within budget',
                'availability': 'High availability with multiple AZs',
                'capacity': 'Adequate capacity for requirements'
            }
            
            return reasons.get(top_criteria[0], 'Good overall score across criteria')
            
        except Exception as e:
            logging.error(f"Failed to generate recommendation reason: {str(e)}")
            return 'Selected based on overall scoring'
    
    async def _store_location_profile(self, profile: LocationProfile):
        """Store location profile in DynamoDB"""
        try:
            profile_dict = asdict(profile)
            profile_dict['last_updated'] = profile.last_updated.isoformat()
            
            self.location_profiles_table.put_item(Item=profile_dict)
            
        except Exception as e:
            logging.error(f"Failed to store location profile: {str(e)}")

# Usage example
async def main():
    config = {
        'location_profiles_table': 'location-profiles',
        'user_locations_table': 'user-locations'
    }
    
    # Initialize location selection engine
    location_engine = LocationSelectionEngine(config)
    
    # Wait for location data to load
    await asyncio.sleep(2)
    
    # Define requirements
    requirements = {
        'user_locations': [
            {
                'location_id': 'us_east_users',
                'geographic_location': {'lat': 40.7128, 'lng': -74.0060},  # New York
                'user_count': 10000,
                'compliance_requirements': ['hipaa'],
                'latency_requirements': {'max_latency_ms': 50}
            },
            {
                'location_id': 'eu_users',
                'geographic_location': {'lat': 51.5074, 'lng': -0.1278},   # London
                'user_count': 5000,
                'compliance_requirements': ['gdpr'],
                'latency_requirements': {'max_latency_ms': 100}
            }
        ],
        'compliance_requirements': ['hipaa', 'gdpr'],
        'performance_requirements': {
            'max_latency_ms': 100,
            'required_capacity': {'compute': 100, 'storage': 1000}
        },
        'cost_constraints': {
            'max_monthly_budget': 10000,
            'base_monthly_cost': 5000
        },
        'location_count': 3
    }
    
    # Select optimal locations
    recommendations = await location_engine.select_optimal_locations(requirements)
    
    print(f"Selected {len(recommendations)} optimal locations:")
    for rec in recommendations:
        print(f"- {rec.region_name}: Score {rec.score:.2f}, Latency {rec.estimated_latency:.1f}ms, Cost ${rec.estimated_cost:.0f}")
        print(f"  Reason: {rec.recommendation_reason}")

if __name__ == "__main__":
    import asyncio
    asyncio.run(main())