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

class DemandPattern(Enum):
    STEADY_INCREASE = "steady_increase"
    SPIKE = "spike"
    SEASONAL = "seasonal"
    CYCLICAL = "cyclical"
    ANOMALOUS = "anomalous"

class ScalingStrategy(Enum):
    PREDICTIVE = "predictive"
    REACTIVE = "reactive"
    SCHEDULED = "scheduled"
    HYBRID = "hybrid"

class ResourceType(Enum):
    COMPUTE = "compute"
    DATABASE = "database"
    CACHE = "cache"
    STORAGE = "storage"
    NETWORK = "network"

@dataclass
class DemandForecast:
    forecast_id: str
    resource_type: ResourceType
    current_demand: float
    predicted_demand: float
    confidence_score: float
    time_horizon: int  # minutes
    demand_pattern: DemandPattern
    recommended_capacity: int
    scaling_urgency: str

@dataclass
class ScalingDecision:
    decision_id: str
    resource_type: ResourceType
    current_capacity: int
    target_capacity: int
    scaling_strategy: ScalingStrategy
    trigger_metrics: Dict[str, float]
    estimated_completion_time: int
    cost_impact: float

class DemandBasedScalingSystem:
    """Intelligent demand-based resource provisioning system"""
    
    def __init__(self, config: Dict[str, Any]):
        self.config = config
        
        # AWS clients
        self.autoscaling = boto3.client('autoscaling')
        self.cloudwatch = boto3.client('cloudwatch')
        self.rds = boto3.client('rds')
        self.elasticache = boto3.client('elasticache')
        self.lambda_client = boto3.client('lambda')
        self.application_autoscaling = boto3.client('application-autoscaling')
        self.dynamodb = boto3.resource('dynamodb')
        
        # Storage
        self.forecasts_table = self.dynamodb.Table(config.get('forecasts_table', 'demand-forecasts'))
        self.decisions_table = self.dynamodb.Table(config.get('decisions_table', 'scaling-decisions'))
        
        # Scaling configuration
        self.scaling_policies = config.get('scaling_policies', {})
        self.demand_thresholds = config.get('demand_thresholds', {})
        self.historical_data = []
        
    async def analyze_demand_patterns(self) -> List[DemandForecast]:
        """Analyze current demand and predict future resource needs"""
        forecasts = []
        
        try:
            # Collect current metrics
            current_metrics = await self._collect_current_metrics()
            
            # Get historical data for pattern analysis
            historical_data = await self._get_historical_data()
            
            # Analyze each resource type
            for resource_type in ResourceType:
                forecast = await self._generate_demand_forecast(
                    resource_type, current_metrics, historical_data
                )
                if forecast:
                    forecasts.append(forecast)
            
            # Store forecasts
            for forecast in forecasts:
                await self._store_forecast(forecast)
            
            return forecasts
            
        except Exception as e:
            logging.error(f"Failed to analyze demand patterns: {str(e)}")
            return []
    
    async def _collect_current_metrics(self) -> Dict[str, Any]:
        """Collect current performance and utilization metrics"""
        try:
            end_time = datetime.utcnow()
            start_time = end_time - timedelta(minutes=30)
            
            metrics = {}
            
            # CPU utilization
            cpu_response = self.cloudwatch.get_metric_statistics(
                Namespace='AWS/EC2',
                MetricName='CPUUtilization',
                StartTime=start_time,
                EndTime=end_time,
                Period=300,
                Statistics=['Average', 'Maximum']
            )
            
            if cpu_response['Datapoints']:
                cpu_data = cpu_response['Datapoints']
                metrics['cpu_average'] = statistics.mean([dp['Average'] for dp in cpu_data])
                metrics['cpu_maximum'] = max([dp['Maximum'] for dp in cpu_data])
            
            # Request count
            request_response = self.cloudwatch.get_metric_statistics(
                Namespace='AWS/ApplicationELB',
                MetricName='RequestCount',
                StartTime=start_time,
                EndTime=end_time,
                Period=300,
                Statistics=['Sum']
            )
            
            if request_response['Datapoints']:
                request_data = request_response['Datapoints']
                metrics['request_rate'] = sum([dp['Sum'] for dp in request_data]) / len(request_data)
            
            # Response time
            latency_response = self.cloudwatch.get_metric_statistics(
                Namespace='AWS/ApplicationELB',
                MetricName='TargetResponseTime',
                StartTime=start_time,
                EndTime=end_time,
                Period=300,
                Statistics=['Average']
            )
            
            if latency_response['Datapoints']:
                latency_data = latency_response['Datapoints']
                metrics['response_time'] = statistics.mean([dp['Average'] for dp in latency_data])
            
            # Database connections
            db_response = self.cloudwatch.get_metric_statistics(
                Namespace='AWS/RDS',
                MetricName='DatabaseConnections',
                StartTime=start_time,
                EndTime=end_time,
                Period=300,
                Statistics=['Average']
            )
            
            if db_response['Datapoints']:
                db_data = db_response['Datapoints']
                metrics['db_connections'] = statistics.mean([dp['Average'] for dp in db_data])
            
            return metrics
            
        except Exception as e:
            logging.error(f"Failed to collect current metrics: {str(e)}")
            return {}
    
    async def _get_historical_data(self) -> List[Dict[str, Any]]:
        """Get historical data for pattern analysis"""
        try:
            end_time = datetime.utcnow()
            start_time = end_time - timedelta(days=7)  # Last 7 days
            
            historical_data = []
            
            # Get historical CPU data
            cpu_response = self.cloudwatch.get_metric_statistics(
                Namespace='AWS/EC2',
                MetricName='CPUUtilization',
                StartTime=start_time,
                EndTime=end_time,
                Period=3600,  # 1 hour periods
                Statistics=['Average']
            )
            
            for datapoint in cpu_response['Datapoints']:
                historical_data.append({
                    'timestamp': datapoint['Timestamp'],
                    'cpu_utilization': datapoint['Average'],
                    'hour_of_day': datapoint['Timestamp'].hour,
                    'day_of_week': datapoint['Timestamp'].weekday()
                })
            
            return sorted(historical_data, key=lambda x: x['timestamp'])
            
        except Exception as e:
            logging.error(f"Failed to get historical data: {str(e)}")
            return []
    
    async def _generate_demand_forecast(self, resource_type: ResourceType, 
                                      current_metrics: Dict[str, Any], 
                                      historical_data: List[Dict[str, Any]]) -> Optional[DemandForecast]:
        """Generate demand forecast for specific resource type"""
        try:
            if resource_type == ResourceType.COMPUTE:
                return await self._forecast_compute_demand(current_metrics, historical_data)
            elif resource_type == ResourceType.DATABASE:
                return await self._forecast_database_demand(current_metrics, historical_data)
            elif resource_type == ResourceType.CACHE:
                return await self._forecast_cache_demand(current_metrics, historical_data)
            else:
                return None
                
        except Exception as e:
            logging.error(f"Failed to generate forecast for {resource_type}: {str(e)}")
            return None
    
    async def _forecast_compute_demand(self, current_metrics: Dict[str, Any], 
                                     historical_data: List[Dict[str, Any]]) -> Optional[DemandForecast]:
        """Forecast compute resource demand"""
        try:
            current_cpu = current_metrics.get('cpu_average', 0)
            current_requests = current_metrics.get('request_rate', 0)
            
            if not historical_data:
                return None
            
            # Analyze historical patterns
            cpu_values = [data['cpu_utilization'] for data in historical_data]
            cpu_trend = self._calculate_trend(cpu_values)
            
            # Predict future demand
            predicted_cpu = current_cpu + (cpu_trend * 60)  # 60 minutes ahead
            predicted_cpu = max(0, min(100, predicted_cpu))  # Clamp between 0-100%
            
            # Determine demand pattern
            pattern = self._identify_demand_pattern(cpu_values, current_cpu)
            
            # Calculate confidence score
            confidence = self._calculate_confidence_score(cpu_values, cpu_trend)
            
            # Determine recommended capacity
            current_capacity = await self._get_current_compute_capacity()
            if predicted_cpu > 70:  # Scale out threshold
                recommended_capacity = int(current_capacity * 1.5)
            elif predicted_cpu < 30:  # Scale in threshold
                recommended_capacity = max(1, int(current_capacity * 0.8))
            else:
                recommended_capacity = current_capacity
            
            # Determine scaling urgency
            if predicted_cpu > 80:
                urgency = "high"
            elif predicted_cpu > 60:
                urgency = "medium"
            else:
                urgency = "low"
            
            return DemandForecast(
                forecast_id=f"compute_{int(time.time())}",
                resource_type=ResourceType.COMPUTE,
                current_demand=current_cpu,
                predicted_demand=predicted_cpu,
                confidence_score=confidence,
                time_horizon=60,
                demand_pattern=pattern,
                recommended_capacity=recommended_capacity,
                scaling_urgency=urgency
            )
            
        except Exception as e:
            logging.error(f"Failed to forecast compute demand: {str(e)}")
            return None
    
    def _calculate_trend(self, values: List[float]) -> float:
        """Calculate trend using linear regression"""
        try:
            if len(values) < 2:
                return 0
            
            x = np.arange(len(values))
            y = np.array(values)
            
            # Simple linear regression
            slope, _ = np.polyfit(x, y, 1)
            return slope
            
        except Exception as e:
            logging.error(f"Failed to calculate trend: {str(e)}")
            return 0
    
    def _identify_demand_pattern(self, historical_values: List[float], current_value: float) -> DemandPattern:
        """Identify the type of demand pattern"""
        try:
            if len(historical_values) < 5:
                return DemandPattern.STEADY_INCREASE
            
            # Calculate recent trend
            recent_values = historical_values[-5:]
            trend = self._calculate_trend(recent_values)
            
            # Calculate volatility
            volatility = np.std(recent_values) if len(recent_values) > 1 else 0
            
            # Identify pattern
            if abs(trend) > 2 and volatility < 5:
                return DemandPattern.STEADY_INCREASE if trend > 0 else DemandPattern.STEADY_INCREASE
            elif volatility > 15:
                return DemandPattern.SPIKE
            elif self._is_seasonal_pattern(historical_values):
                return DemandPattern.SEASONAL
            elif self._is_cyclical_pattern(historical_values):
                return DemandPattern.CYCLICAL
            else:
                return DemandPattern.STEADY_INCREASE
                
        except Exception as e:
            logging.error(f"Failed to identify demand pattern: {str(e)}")
            return DemandPattern.STEADY_INCREASE
    
    def _is_seasonal_pattern(self, values: List[float]) -> bool:
        """Check if data shows seasonal patterns"""
        # Simplified seasonal detection
        if len(values) < 24:  # Need at least 24 hours of data
            return False
        
        # Check for daily patterns (simplified)
        daily_averages = []
        for i in range(0, len(values), 24):
            daily_chunk = values[i:i+24]
            if daily_chunk:
                daily_averages.append(statistics.mean(daily_chunk))
        
        if len(daily_averages) < 2:
            return False
        
        # Check if daily patterns are similar
        daily_std = np.std(daily_averages)
        return daily_std < 10  # Low variation between days indicates seasonality
    
    def _is_cyclical_pattern(self, values: List[float]) -> bool:
        """Check if data shows cyclical patterns"""
        # Simplified cyclical detection
        if len(values) < 12:
            return False
        
        # Check for repeating patterns
        half_length = len(values) // 2
        first_half = values[:half_length]
        second_half = values[half_length:half_length*2]
        
        if len(first_half) != len(second_half):
            return False
        
        # Calculate correlation between halves
        correlation = np.corrcoef(first_half, second_half)[0, 1]
        return correlation > 0.7  # High correlation indicates cyclical pattern
    
    def _calculate_confidence_score(self, historical_values: List[float], trend: float) -> float:
        """Calculate confidence score for the forecast"""
        try:
            if len(historical_values) < 3:
                return 0.3  # Low confidence with limited data
            
            # Calculate prediction accuracy based on historical trend consistency
            recent_values = historical_values[-10:] if len(historical_values) >= 10 else historical_values
            
            # Check trend consistency
            trend_consistency = 1.0 - (np.std(recent_values) / (np.mean(recent_values) + 1))
            trend_consistency = max(0, min(1, trend_consistency))
            
            # Data quality factor
            data_quality = min(1.0, len(historical_values) / 24)  # More data = higher quality
            
            # Combine factors
            confidence = (trend_consistency * 0.7) + (data_quality * 0.3)
            return round(confidence, 2)
            
        except Exception as e:
            logging.error(f"Failed to calculate confidence score: {str(e)}")
            return 0.5
    
    async def _get_current_compute_capacity(self) -> int:
        """Get current compute capacity from Auto Scaling Groups"""
        try:
            response = self.autoscaling.describe_auto_scaling_groups()
            
            total_capacity = 0
            for asg in response['AutoScalingGroups']:
                total_capacity += asg['DesiredCapacity']
            
            return total_capacity if total_capacity > 0 else 1
            
        except Exception as e:
            logging.error(f"Failed to get current compute capacity: {str(e)}")
            return 1
    
    async def make_scaling_decisions(self, forecasts: List[DemandForecast]) -> List[ScalingDecision]:
        """Make scaling decisions based on demand forecasts"""
        decisions = []
        
        try:
            for forecast in forecasts:
                decision = await self._create_scaling_decision(forecast)
                if decision:
                    decisions.append(decision)
            
            # Store decisions
            for decision in decisions:
                await self._store_decision(decision)
            
            return decisions
            
        except Exception as e:
            logging.error(f"Failed to make scaling decisions: {str(e)}")
            return []
    
    async def _create_scaling_decision(self, forecast: DemandForecast) -> Optional[ScalingDecision]:
        """Create scaling decision based on forecast"""
        try:
            current_capacity = await self._get_current_capacity(forecast.resource_type)
            
            # Determine if scaling is needed
            if forecast.recommended_capacity == current_capacity:
                return None  # No scaling needed
            
            # Determine scaling strategy
            if forecast.scaling_urgency == "high":
                strategy = ScalingStrategy.REACTIVE
            elif forecast.demand_pattern == DemandPattern.SEASONAL:
                strategy = ScalingStrategy.SCHEDULED
            elif forecast.confidence_score > 0.8:
                strategy = ScalingStrategy.PREDICTIVE
            else:
                strategy = ScalingStrategy.HYBRID
            
            # Calculate cost impact
            cost_impact = await self._calculate_cost_impact(
                current_capacity, forecast.recommended_capacity, forecast.resource_type
            )
            
            # Estimate completion time
            completion_time = self._estimate_scaling_time(
                current_capacity, forecast.recommended_capacity, forecast.resource_type
            )
            
            return ScalingDecision(
                decision_id=f"decision_{int(time.time())}_{forecast.forecast_id}",
                resource_type=forecast.resource_type,
                current_capacity=current_capacity,
                target_capacity=forecast.recommended_capacity,
                scaling_strategy=strategy,
                trigger_metrics={
                    'current_demand': forecast.current_demand,
                    'predicted_demand': forecast.predicted_demand,
                    'confidence_score': forecast.confidence_score
                },
                estimated_completion_time=completion_time,
                cost_impact=cost_impact
            )
            
        except Exception as e:
            logging.error(f"Failed to create scaling decision: {str(e)}")
            return None
    
    async def _get_current_capacity(self, resource_type: ResourceType) -> int:
        """Get current capacity for resource type"""
        try:
            if resource_type == ResourceType.COMPUTE:
                return await self._get_current_compute_capacity()
            elif resource_type == ResourceType.DATABASE:
                return await self._get_current_database_capacity()
            elif resource_type == ResourceType.CACHE:
                return await self._get_current_cache_capacity()
            else:
                return 1
                
        except Exception as e:
            logging.error(f"Failed to get current capacity for {resource_type}: {str(e)}")
            return 1
    
    async def _calculate_cost_impact(self, current_capacity: int, target_capacity: int, 
                                   resource_type: ResourceType) -> float:
        """Calculate cost impact of scaling decision"""
        try:
            capacity_change = target_capacity - current_capacity
            
            # Simplified cost calculation (per hour)
            if resource_type == ResourceType.COMPUTE:
                hourly_cost_per_unit = 0.10  # $0.10 per instance per hour
            elif resource_type == ResourceType.DATABASE:
                hourly_cost_per_unit = 0.20  # $0.20 per DB instance per hour
            elif resource_type == ResourceType.CACHE:
                hourly_cost_per_unit = 0.05  # $0.05 per cache node per hour
            else:
                hourly_cost_per_unit = 0.10
            
            return capacity_change * hourly_cost_per_unit
            
        except Exception as e:
            logging.error(f"Failed to calculate cost impact: {str(e)}")
            return 0.0
    
    def _estimate_scaling_time(self, current_capacity: int, target_capacity: int, 
                             resource_type: ResourceType) -> int:
        """Estimate time to complete scaling operation (in seconds)"""
        try:
            capacity_change = abs(target_capacity - current_capacity)
            
            # Estimated time per resource type
            if resource_type == ResourceType.COMPUTE:
                time_per_unit = 120  # 2 minutes per EC2 instance
            elif resource_type == ResourceType.DATABASE:
                time_per_unit = 300  # 5 minutes per DB instance
            elif resource_type == ResourceType.CACHE:
                time_per_unit = 180  # 3 minutes per cache node
            else:
                time_per_unit = 120
            
            return capacity_change * time_per_unit
            
        except Exception as e:
            logging.error(f"Failed to estimate scaling time: {str(e)}")
            return 300  # Default 5 minutes
    
    async def execute_scaling_decisions(self, decisions: List[ScalingDecision]) -> List[str]:
        """Execute scaling decisions"""
        execution_ids = []
        
        try:
            for decision in decisions:
                execution_id = await self._execute_scaling_decision(decision)
                if execution_id:
                    execution_ids.append(execution_id)
            
            return execution_ids
            
        except Exception as e:
            logging.error(f"Failed to execute scaling decisions: {str(e)}")
            return []
    
    async def _execute_scaling_decision(self, decision: ScalingDecision) -> Optional[str]:
        """Execute a specific scaling decision"""
        try:
            execution_id = f"exec_{int(time.time())}_{decision.decision_id}"
            
            if decision.resource_type == ResourceType.COMPUTE:
                success = await self._execute_compute_scaling(decision)
            elif decision.resource_type == ResourceType.DATABASE:
                success = await self._execute_database_scaling(decision)
            elif decision.resource_type == ResourceType.CACHE:
                success = await self._execute_cache_scaling(decision)
            else:
                success = False
            
            if success:
                logging.info(f"Successfully executed scaling decision: {decision.decision_id}")
                return execution_id
            else:
                logging.error(f"Failed to execute scaling decision: {decision.decision_id}")
                return None
                
        except Exception as e:
            logging.error(f"Failed to execute scaling decision: {str(e)}")
            return None
    
    async def _execute_compute_scaling(self, decision: ScalingDecision) -> bool:
        """Execute compute scaling decision"""
        try:
            # Get Auto Scaling Groups
            response = self.autoscaling.describe_auto_scaling_groups()
            
            if not response['AutoScalingGroups']:
                return False
            
            # Scale the first ASG (simplified)
            asg = response['AutoScalingGroups'][0]
            asg_name = asg['AutoScalingGroupName']
            
            # Update desired capacity
            self.autoscaling.set_desired_capacity(
                AutoScalingGroupName=asg_name,
                DesiredCapacity=decision.target_capacity,
                HonorCooldown=False
            )
            
            logging.info(f"Scaled {asg_name} to {decision.target_capacity} instances")
            return True
            
        except Exception as e:
            logging.error(f"Failed to execute compute scaling: {str(e)}")
            return False
    
    async def _store_forecast(self, forecast: DemandForecast):
        """Store demand forecast in DynamoDB"""
        try:
            forecast_dict = asdict(forecast)
            self.forecasts_table.put_item(Item=forecast_dict)
            
        except Exception as e:
            logging.error(f"Failed to store forecast: {str(e)}")
    
    async def _store_decision(self, decision: ScalingDecision):
        """Store scaling decision in DynamoDB"""
        try:
            decision_dict = asdict(decision)
            self.decisions_table.put_item(Item=decision_dict)
            
        except Exception as e:
            logging.error(f"Failed to store decision: {str(e)}")

# Usage example
async def main():
    config = {
        'forecasts_table': 'demand-forecasts',
        'decisions_table': 'scaling-decisions',
        'scaling_policies': {
            'compute': {
                'scale_out_threshold': 70,
                'scale_in_threshold': 30,
                'max_capacity': 20
            }
        }
    }
    
    # Initialize scaling system
    scaling_system = DemandBasedScalingSystem(config)
    
    # Analyze demand patterns
    forecasts = await scaling_system.analyze_demand_patterns()
    print(f"Generated {len(forecasts)} demand forecasts")
    
    # Make scaling decisions
    decisions = await scaling_system.make_scaling_decisions(forecasts)
    print(f"Made {len(decisions)} scaling decisions")
    
    # Execute scaling decisions
    executions = await scaling_system.execute_scaling_decisions(decisions)
    print(f"Executed {len(executions)} scaling operations")

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