Capacity Planning.

# capacity/baseline.py

"""
Capacity baseline measurement and tracking.

Establishes what "normal" looks like for your system.
"""

from dataclasses import dataclass, field
from typing import Dict, List, Optional
from datetime import datetime, timedelta, timezone
from enum import Enum
import statistics


@dataclass
class ResourceMetrics:
    """Metrics for a single resource."""
    name: str
    current_value: float
    max_capacity: float
    unit: str
    
    @property
    def utilization(self) -> float:
        """Current utilization percentage."""
        return (self.current_value / self.max_capacity) * 100 if self.max_capacity > 0 else 0
    
    @property
    def headroom(self) -> float:
        """Available headroom percentage."""
        return 100 - self.utilization
    
    @property
    def status(self) -> str:
        """Health status based on headroom."""
        if self.headroom > 50:
            return "healthy"
        elif self.headroom > 30:
            return "adequate"
        elif self.headroom > 20:
            return "warning"
        elif self.headroom > 10:
            return "critical"
        else:
            return "emergency"


@dataclass
class CapacityBaseline:
    """Baseline capacity measurements for a service."""
    service_name: str
    measured_at: datetime
    
    # Traffic metrics
    avg_rps: float
    peak_rps: float
    p50_latency_ms: float
    p99_latency_ms: float
    
    # Resource metrics
    resources: Dict[str, ResourceMetrics] = field(default_factory=dict)
    
    # Derived limits
    max_sustainable_rps: float = 0  # RPS before latency degrades
    breaking_point_rps: float = 0   # RPS when errors start


class CapacityMeasurer:
    """
    Measures and tracks system capacity.
    """
    
    def __init__(self, metrics_client, config):
        self.metrics = metrics_client
        self.config = config
    
    async def measure_baseline(
        self,
        service_name: str,
        window: str = "7d"
    ) -> CapacityBaseline:
        """
        Measure baseline capacity from historical data.
        """
        # Traffic metrics
        avg_rps = await self.metrics.query(
            f'avg(rate(http_requests_total{{service="{service_name}"}}[5m]))',
            window=window
        )
        
        peak_rps = await self.metrics.query(
            f'max(rate(http_requests_total{{service="{service_name}"}}[5m]))',
            window=window
        )
        
        p50_latency = await self.metrics.query(
            f'histogram_quantile(0.5, sum by (le) (rate(http_request_duration_seconds_bucket{{service="{service_name}"}}[5m])))',
            window=window
        ) * 1000  # Convert to ms
        
        p99_latency = await self.metrics.query(
            f'histogram_quantile(0.99, sum by (le) (rate(http_request_duration_seconds_bucket{{service="{service_name}"}}[5m])))',
            window=window
        ) * 1000
        
        # Resource metrics
        resources = {}
        
        # CPU
        cpu_usage = await self.metrics.query(
            f'avg(rate(container_cpu_usage_seconds_total{{service="{service_name}"}}[5m]))',
            window=window
        )
        cpu_limit = await self.metrics.query(
            f'avg(container_spec_cpu_quota{{service="{service_name}"}} / container_spec_cpu_period{{service="{service_name}"}})',
            window=window
        )
        resources['cpu'] = ResourceMetrics(
            name='CPU',
            current_value=cpu_usage,
            max_capacity=cpu_limit,
            unit='cores'
        )
        
        # Memory
        memory_usage = await self.metrics.query(
            f'avg(container_memory_working_set_bytes{{service="{service_name}"}})',
            window=window
        ) / (1024 ** 3)  # Convert to GB
        memory_limit = await self.metrics.query(
            f'avg(container_spec_memory_limit_bytes{{service="{service_name}"}})',
            window=window
        ) / (1024 ** 3)
        resources['memory'] = ResourceMetrics(
            name='Memory',
            current_value=memory_usage,
            max_capacity=memory_limit,
            unit='GB'
        )
        
        # Database connections
        db_connections = await self.metrics.query(
            f'avg(db_connections_active{{service="{service_name}"}})',
            window=window
        )
        db_max_connections = await self.metrics.query(
            f'avg(db_connections_max{{service="{service_name}"}})',
            window=window
        )
        resources['db_connections'] = ResourceMetrics(
            name='Database Connections',
            current_value=db_connections,
            max_capacity=db_max_connections,
            unit='connections'
        )
        
        return CapacityBaseline(
            service_name=service_name,
            measured_at=datetime.now(timezone.utc),
            avg_rps=avg_rps,
            peak_rps=peak_rps,
            p50_latency_ms=p50_latency,
            p99_latency_ms=p99_latency,
            resources=resources
        )
    
    async def generate_capacity_report(
        self,
        service_name: str
    ) -> str:
        """
        Generate a human-readable capacity report.
        """
        baseline = await self.measure_baseline(service_name)
        
        report = f"""
CAPACITY REPORT: {service_name}
Generated: {baseline.measured_at.isoformat()}
{'=' * 60}

TRAFFIC METRICS
---------------
Average RPS:      {baseline.avg_rps:.1f}
Peak RPS:         {baseline.peak_rps:.1f}
P50 Latency:      {baseline.p50_latency_ms:.1f}ms
P99 Latency:      {baseline.p99_latency_ms:.1f}ms

RESOURCE UTILIZATION
--------------------
"""
        for name, resource in baseline.resources.items():
            status_emoji = {
                'healthy': '🟢',
                'adequate': '🟡',
                'warning': '🟠',
                'critical': '🔴',
                'emergency': '⚫'
            }.get(resource.status, '⚪')
            
            report += f"""
{resource.name}:
  Current:   {resource.current_value:.2f} {resource.unit}
  Maximum:   {resource.max_capacity:.2f} {resource.unit}
  Usage:     {resource.utilization:.1f}%
  Headroom:  {resource.headroom:.1f}%
  Status:    {status_emoji} {resource.status.upper()}
"""
        
        # Identify bottleneck
        bottleneck = min(
            baseline.resources.values(),
            key=lambda r: r.headroom
        )
        
        report += f"""
{'=' * 60}
BOTTLENECK ANALYSIS
-------------------
Primary bottleneck: {bottleneck.name}
  - Only {bottleneck.headroom:.1f}% headroom remaining
  - Will likely be first to saturate under load

RECOMMENDATIONS
---------------
"""
        for name, resource in baseline.resources.items():
            if resource.headroom < 30:
                report += f"  ⚠️  {resource.name}: Prepare or start scaling (headroom < 30%)\n"
            elif resource.headroom < 50:
                report += f"  📋 {resource.name}: Start planning scaling\n"
        
        return report

# capacity/forecasting.py

"""
Demand forecasting for capacity planning.

Predicts future load based on historical trends and known events.
"""

from dataclasses import dataclass
from typing import List, Optional, Tuple
from datetime import datetime, timedelta, timezone
import math


@dataclass
class TrafficForecast:
    """Forecast for future traffic."""
    date: datetime
    predicted_rps: float
    confidence_low: float
    confidence_high: float
    factors: List[str]  # What's driving this prediction


@dataclass
class GrowthModel:
    """Model for traffic growth."""
    base_rps: float
    monthly_growth_rate: float  # e.g., 0.10 for 10% MoM
    seasonality: dict  # Month -> multiplier
    day_of_week_pattern: dict  # Day -> multiplier
    hour_of_day_pattern: dict  # Hour -> multiplier


class DemandForecaster:
    """
    Forecasts future demand based on historical patterns.
    """
    
    def __init__(self, metrics_client):
        self.metrics = metrics_client
    
    async def build_growth_model(
        self,
        service_name: str,
        lookback_months: int = 6
    ) -> GrowthModel:
        """
        Build growth model from historical data.
        """
        # Get monthly averages
        monthly_data = []
        for i in range(lookback_months, 0, -1):
            start = datetime.now(timezone.utc) - timedelta(days=30 * i)
            end = start + timedelta(days=30)
            
            avg_rps = await self.metrics.query_range(
                f'avg(rate(http_requests_total{{service="{service_name}"}}[1h]))',
                start=start,
                end=end
            )
            monthly_data.append(avg_rps.mean())
        
        # Calculate growth rate
        if len(monthly_data) >= 2:
            # Compound monthly growth rate
            start_value = monthly_data[0]
            end_value = monthly_data[-1]
            months = len(monthly_data) - 1
            
            if start_value > 0:
                growth_rate = (end_value / start_value) ** (1 / months) - 1
            else:
                growth_rate = 0
        else:
            growth_rate = 0
        
        # Get seasonality (month of year patterns)
        seasonality = await self._calculate_seasonality(service_name)
        
        # Get day-of-week patterns
        dow_pattern = await self._calculate_dow_pattern(service_name)
        
        # Get hour-of-day patterns
        hod_pattern = await self._calculate_hod_pattern(service_name)
        
        return GrowthModel(
            base_rps=monthly_data[-1] if monthly_data else 0,
            monthly_growth_rate=growth_rate,
            seasonality=seasonality,
            day_of_week_pattern=dow_pattern,
            hour_of_day_pattern=hod_pattern
        )
    
    def forecast(
        self,
        model: GrowthModel,
        target_date: datetime,
        include_events: List[dict] = None
    ) -> TrafficForecast:
        """
        Forecast traffic for a target date.
        """
        now = datetime.now(timezone.utc)
        months_ahead = (target_date - now).days / 30
        
        # Base growth projection
        base_projection = model.base_rps * ((1 + model.monthly_growth_rate) ** months_ahead)
        
        # Apply seasonality
        month = target_date.month
        seasonality_factor = model.seasonality.get(month, 1.0)
        
        # Apply day-of-week pattern
        dow = target_date.weekday()
        dow_factor = model.day_of_week_pattern.get(dow, 1.0)
        
        # Apply hour pattern for peak hour
        peak_hour_factor = max(model.hour_of_day_pattern.values())
        
        # Calculate predicted RPS
        predicted_rps = base_projection * seasonality_factor * dow_factor * peak_hour_factor
        
        # Apply event multipliers
        factors = ["organic_growth", f"seasonality_{month}", f"dow_{dow}"]
        
        if include_events:
            for event in include_events:
                if event['date'].date() == target_date.date():
                    predicted_rps *= event['traffic_multiplier']
                    factors.append(event['name'])
        
        # Confidence intervals (simplified)
        confidence_low = predicted_rps * 0.8
        confidence_high = predicted_rps * 1.4
        
        return TrafficForecast(
            date=target_date,
            predicted_rps=predicted_rps,
            confidence_low=confidence_low,
            confidence_high=confidence_high,
            factors=factors
        )
    
    def forecast_range(
        self,
        model: GrowthModel,
        start_date: datetime,
        end_date: datetime,
        known_events: List[dict] = None
    ) -> List[TrafficForecast]:
        """
        Forecast traffic for a date range.
        """
        forecasts = []
        current = start_date
        
        while current <= end_date:
            forecast = self.forecast(model, current, known_events)
            forecasts.append(forecast)
            current += timedelta(days=1)
        
        return forecasts


# =============================================================================
# EXAMPLE: PLANNING FOR KNOWN EVENTS
# =============================================================================

"""
Known events that affect traffic:

BLACK_FRIDAY = {
    'name': 'Black Friday',
    'date': datetime(2024, 11, 29),
    'traffic_multiplier': 5.0,  # 5x normal traffic
    'duration_hours': 24
}

PRODUCT_LAUNCH = {
    'name': 'Product Launch',
    'date': datetime(2024, 10, 15),
    'traffic_multiplier': 3.0,
    'duration_hours': 48
}

MARKETING_CAMPAIGN = {
    'name': 'TV Ad Campaign',
    'date': datetime(2024, 9, 1),
    'traffic_multiplier': 2.0,
    'duration_hours': 168  # 1 week
}
"""

# capacity/load_testing.py

"""
Load testing orchestration and analysis.

Uses k6, Locust, or similar tools under the hood.
"""

from dataclasses import dataclass, field
from typing import List, Dict, Optional
from datetime import datetime, timedelta, timezone
from enum import Enum
import asyncio
import subprocess
import json


class LoadPattern(Enum):
    CONSTANT = "constant"
    RAMP_UP = "ramp_up"
    SPIKE = "spike"
    STRESS = "stress"


@dataclass
class LoadTestConfig:
    """Configuration for a load test."""
    name: str
    target_url: str
    
    # Load profile
    pattern: LoadPattern
    virtual_users: int
    ramp_up_time_seconds: int = 60
    duration_seconds: int = 300
    
    # For spike tests
    spike_users: int = 0
    spike_duration_seconds: int = 60
    
    # Thresholds
    max_latency_p99_ms: int = 500
    max_error_rate: float = 0.01
    
    # Scenarios
    scenarios: List[dict] = field(default_factory=list)


@dataclass
class LoadTestResult:
    """Results from a load test."""
    config: LoadTestConfig
    started_at: datetime
    completed_at: datetime
    
    # Summary metrics
    total_requests: int
    successful_requests: int
    failed_requests: int
    
    # Latency
    latency_p50_ms: float
    latency_p95_ms: float
    latency_p99_ms: float
    latency_max_ms: float
    
    # Throughput
    requests_per_second: float
    
    # Thresholds
    passed: bool
    threshold_violations: List[str]
    
    # Raw data for analysis
    timeseries_data: List[dict] = field(default_factory=list)


class LoadTestRunner:
    """
    Runs and analyzes load tests.
    """
    
    def __init__(self, metrics_client, alerter):
        self.metrics = metrics_client
        self.alerter = alerter
    
    async def run_test(
        self,
        config: LoadTestConfig,
        notify_on_failure: bool = True
    ) -> LoadTestResult:
        """
        Run a load test and analyze results.
        """
        started_at = datetime.now(timezone.utc)
        
        # Generate k6 script
        script = self._generate_k6_script(config)
        
        # Run k6
        result_json = await self._execute_k6(script, config)
        
        # Parse results
        result = self._parse_results(config, result_json, started_at)
        
        # Check thresholds
        result.passed, result.threshold_violations = self._check_thresholds(
            config, result
        )
        
        # Notify if failed
        if not result.passed and notify_on_failure:
            await self.alerter.send_warning(
                f"Load test '{config.name}' failed: {result.threshold_violations}"
            )
        
        # Store results
        await self._store_results(result)
        
        return result
    
    def _generate_k6_script(self, config: LoadTestConfig) -> str:
        """Generate k6 test script."""
        
        if config.pattern == LoadPattern.CONSTANT:
            stages = f"""
                stages: [
                    {{ duration: '{config.ramp_up_time_seconds}s', target: {config.virtual_users} }},
                    {{ duration: '{config.duration_seconds}s', target: {config.virtual_users} }},
                    {{ duration: '30s', target: 0 }},
                ]
            """
        elif config.pattern == LoadPattern.RAMP_UP:
            stages = f"""
                stages: [
                    {{ duration: '{config.duration_seconds}s', target: {config.virtual_users} }},
                    {{ duration: '60s', target: 0 }},
                ]
            """
        elif config.pattern == LoadPattern.SPIKE:
            stages = f"""
                stages: [
                    {{ duration: '60s', target: {config.virtual_users} }},
                    {{ duration: '{config.spike_duration_seconds}s', target: {config.spike_users} }},
                    {{ duration: '60s', target: {config.virtual_users} }},
                    {{ duration: '30s', target: 0 }},
                ]
            """
        elif config.pattern == LoadPattern.STRESS:
            # Gradually increase until we find the breaking point
            stages = f"""
                stages: [
                    {{ duration: '60s', target: {config.virtual_users} }},
                    {{ duration: '120s', target: {config.virtual_users * 2} }},
                    {{ duration: '120s', target: {config.virtual_users * 4} }},
                    {{ duration: '120s', target: {config.virtual_users * 8} }},
                    {{ duration: '60s', target: 0 }},
                ]
            """
        
        script = f"""
import http from 'k6/http';
import {{ check, sleep }} from 'k6';
import {{ Rate, Trend }} from 'k6/metrics';

const errorRate = new Rate('errors');
const latency = new Trend('latency');

export let options = {{
    {stages},
    thresholds: {{
        'http_req_duration': ['p(99)<{config.max_latency_p99_ms}'],
        'errors': ['rate<{config.max_error_rate}'],
    }},
}};

export default function() {{
    const res = http.get('{config.target_url}');
    
    const success = check(res, {{
        'status is 200': (r) => r.status === 200,
        'latency < 500ms': (r) => r.timings.duration < 500,
    }});
    
    errorRate.add(!success);
    latency.add(res.timings.duration);
    
    sleep(1);
}}
"""
        return script
    
    async def _execute_k6(self, script: str, config: LoadTestConfig) -> dict:
        """Execute k6 and return JSON results."""
        # Write script to temp file
        script_path = f'/tmp/k6_script_{config.name}.js'
        with open(script_path, 'w') as f:
            f.write(script)
        
        # Run k6 and export the end-of-test summary as one JSON document
        # (--out json=... writes one JSON object per line, not a single document)
        result = subprocess.run(
            ['k6', 'run', '--summary-export', '/tmp/k6_summary.json', script_path],
            capture_output=True,
            text=True
        )
        
        # Parse JSON summary
        with open('/tmp/k6_summary.json', 'r') as f:
            return json.load(f)
    
    def _check_thresholds(
        self,
        config: LoadTestConfig,
        result: LoadTestResult
    ) -> tuple[bool, List[str]]:
        """Check if results meet thresholds."""
        violations = []
        
        if result.latency_p99_ms > config.max_latency_p99_ms:
            violations.append(
                f"P99 latency {result.latency_p99_ms}ms > threshold {config.max_latency_p99_ms}ms"
            )
        
        error_rate = result.failed_requests / result.total_requests
        if error_rate > config.max_error_rate:
            violations.append(
                f"Error rate {error_rate:.2%} > threshold {config.max_error_rate:.2%}"
            )
        
        return len(violations) == 0, violations


# =============================================================================
# EXAMPLE LOAD TEST SCENARIOS
# =============================================================================

API_LOAD_TEST = LoadTestConfig(
    name="api_baseline",
    target_url="https://api.example.com/health",
    pattern=LoadPattern.CONSTANT,
    virtual_users=100,
    duration_seconds=300,
    max_latency_p99_ms=200,
    max_error_rate=0.01
)

BLACK_FRIDAY_TEST = LoadTestConfig(
    name="black_friday_simulation",
    target_url="https://api.example.com/checkout",
    pattern=LoadPattern.SPIKE,
    virtual_users=500,
    spike_users=5000,
    spike_duration_seconds=120,
    duration_seconds=600,
    max_latency_p99_ms=1000,
    max_error_rate=0.05
)

STRESS_TEST = LoadTestConfig(
    name="find_breaking_point",
    target_url="https://api.example.com/api",
    pattern=LoadPattern.STRESS,
    virtual_users=100,
    duration_seconds=600,
    max_latency_p99_ms=5000,
    max_error_rate=0.10
)

# capacity/autoscaling.py

"""
Auto-scaling configuration and policies.

Defines when and how to scale automatically.
"""

from dataclasses import dataclass
from typing import List, Optional
from enum import Enum


class ScalingMetric(Enum):
    CPU = "cpu"
    MEMORY = "memory"
    RPS = "requests_per_second"
    QUEUE_DEPTH = "queue_depth"
    LATENCY = "latency_p99"
    CUSTOM = "custom"


@dataclass
class ScalingPolicy:
    """Defines when to scale."""
    name: str
    metric: ScalingMetric
    
    # Thresholds
    scale_up_threshold: float
    scale_down_threshold: float
    
    # Behavior
    scale_up_cooldown_seconds: int = 300
    scale_down_cooldown_seconds: int = 600
    
    # Limits
    min_replicas: int = 2
    max_replicas: int = 100
    
    # Aggressiveness
    scale_up_by: int = 2  # Add 2 pods at a time
    scale_down_by: int = 1  # Remove 1 pod at a time


@dataclass
class AutoScalingConfig:
    """Complete auto-scaling configuration."""
    deployment_name: str
    namespace: str
    
    # Policies (can have multiple)
    policies: List[ScalingPolicy]
    
    # Predictive scaling (for known events)
    scheduled_scaling: List[dict] = None
    
    # Protection
    disable_scale_down_during_deployment: bool = True
    

# =============================================================================
# EXAMPLE CONFIGURATIONS
# =============================================================================

# Web API: Scale on CPU and latency
WEB_API_SCALING = AutoScalingConfig(
    deployment_name="api-server",
    namespace="production",
    policies=[
        # Primary: CPU-based scaling
        ScalingPolicy(
            name="cpu-scaling",
            metric=ScalingMetric.CPU,
            scale_up_threshold=70,  # Scale up when CPU > 70%
            scale_down_threshold=30,  # Scale down when CPU < 30%
            min_replicas=3,
            max_replicas=50,
            scale_up_by=3,
            scale_down_by=1,
            scale_down_cooldown_seconds=900  # 15 min cooldown
        ),
        # Secondary: Latency-based scaling
        ScalingPolicy(
            name="latency-scaling",
            metric=ScalingMetric.LATENCY,
            scale_up_threshold=200,  # Scale up when p99 > 200ms
            scale_down_threshold=50,  # Scale down when p99 < 50ms
            min_replicas=3,
            max_replicas=50,
            scale_up_by=5,  # More aggressive for latency
            scale_up_cooldown_seconds=60  # React faster to latency
        ),
    ],
    scheduled_scaling=[
        # Pre-scale for known traffic patterns
        {
            "name": "morning_ramp",
            "schedule": "0 8 * * 1-5",  # 8 AM weekdays
            "min_replicas": 10
        },
        {
            "name": "evening_scale_down",
            "schedule": "0 22 * * *",  # 10 PM daily
            "min_replicas": 3
        },
        {
            "name": "black_friday",
            # Cron can't express "Thanksgiving week" ("0 0 * 11 5" would mean
            # every Friday in November), so pin the dates and update yearly
            "schedule": "0 0 25-29 11 *",  # Nov 25-29, 2024 (Thanksgiving week)
            "min_replicas": 50,
            "max_replicas": 200
        },
    ]
)

# Worker pool: Scale on queue depth
WORKER_SCALING = AutoScalingConfig(
    deployment_name="worker-pool",
    namespace="production",
    policies=[
        ScalingPolicy(
            name="queue-scaling",
            metric=ScalingMetric.QUEUE_DEPTH,
            scale_up_threshold=1000,  # Scale up when queue > 1000
            scale_down_threshold=100,  # Scale down when queue < 100
            min_replicas=2,
            max_replicas=20,
            scale_up_by=2,
            scale_up_cooldown_seconds=60  # React quickly to queue buildup
        ),
    ]
)


def generate_kubernetes_hpa(config: AutoScalingConfig) -> dict:
    """Generate Kubernetes HPA manifest from config."""
    
    hpa = {
        "apiVersion": "autoscaling/v2",
        "kind": "HorizontalPodAutoscaler",
        "metadata": {
            "name": f"{config.deployment_name}-hpa",
            "namespace": config.namespace
        },
        "spec": {
            "scaleTargetRef": {
                "apiVersion": "apps/v1",
                "kind": "Deployment",
                "name": config.deployment_name
            },
            "minReplicas": min(p.min_replicas for p in config.policies),
            "maxReplicas": max(p.max_replicas for p in config.policies),
            "metrics": [],
            "behavior": {
                "scaleUp": {
                    "stabilizationWindowSeconds": 60,
                    "policies": [
                        {
                            "type": "Pods",
                            "value": max(p.scale_up_by for p in config.policies),
                            "periodSeconds": 60
                        }
                    ]
                },
                "scaleDown": {
                    "stabilizationWindowSeconds": 300,
                    "policies": [
                        {
                            "type": "Pods",
                            "value": 1,
                            "periodSeconds": 120
                        }
                    ]
                }
            }
        }
    }
    
    # Add metrics
    for policy in config.policies:
        if policy.metric == ScalingMetric.CPU:
            hpa["spec"]["metrics"].append({
                "type": "Resource",
                "resource": {
                    "name": "cpu",
                    "target": {
                        "type": "Utilization",
                        "averageUtilization": int(policy.scale_up_threshold)
                    }
                }
            })
        elif policy.metric == ScalingMetric.MEMORY:
            hpa["spec"]["metrics"].append({
                "type": "Resource",
                "resource": {
                    "name": "memory",
                    "target": {
                        "type": "Utilization",
                        "averageUtilization": int(policy.scale_up_threshold)
                    }
                }
            })
    
    return hpa