infrastructure/cost_estimator.py

"""
Cost Estimation Engine – Deterministic pricing with Bayesian uncertainty.
"""

import os
from functools import lru_cache
from typing import Dict, Optional, Union, Any
import yaml

from .intents import ProvisionResourceIntent, ResourceType

class CostEstimator:
    DEFAULT_PRICING = {
        ResourceType.VM: {
            "Standard_D2s_v3": 70.0,
            "Standard_D4s_v3": 140.0,
            "Standard_D8s_v3": 280.0,
            "Standard_D16s_v3": 560.0,
        },
        ResourceType.STORAGE_ACCOUNT: {
            "50GB": 5.0,
            "100GB": 10.0,
            "1TB": 100.0,
            "10TB": 900.0,
        },
        ResourceType.DATABASE: {
            "Basic": 15.0,
            "Standard": 50.0,
            "Premium": 200.0,
        },
        ResourceType.KUBERNETES_CLUSTER: {
            "Small": 100.0,
            "Medium": 300.0,
            "Large": 600.0,
        },
        ResourceType.FUNCTION_APP: {
            "Consumption": 0.0,
            "Premium": 75.0,
        },
        ResourceType.VIRTUAL_NETWORK: {
            "default": 0.0,
        },
    }

    def __init__(self, pricing_file: Optional[str] = None):
        if pricing_file and os.path.exists(pricing_file):
            with open(pricing_file, 'r') as f:
                raw = yaml.safe_load(f)
            self._pricing = {}
            for res_str, sizes in raw.items():
                try:
                    res_type = ResourceType(res_str)
                except ValueError:
                    continue
                self._pricing[res_type] = sizes
        else:
            self._pricing = self.DEFAULT_PRICING.copy()

        # Cache for estimate_monthly_cost using a hashable key
        self._cost_cache = {}

    def estimate_monthly_cost(self, intent: ProvisionResourceIntent) -> Optional[float]:
        # Create a hashable key from the intent's immutable fields
        # (resource_type, size, region, environment) – configuration is omitted as it doesn't affect cost in this demo.
        key = (intent.resource_type, intent.size, intent.region, intent.environment)
        if key in self._cost_cache:
            return self._cost_cache[key]

        resource_pricing = self._pricing.get(intent.resource_type)
        if not resource_pricing:
            self._cost_cache[key] = None
            return None
        cost = resource_pricing.get(intent.size)
        self._cost_cache[key] = cost
        return cost

    def cost_delta_vs_baseline(
        self,
        intent: ProvisionResourceIntent,
        baseline_intent: Optional[ProvisionResourceIntent] = None,
    ) -> Optional[float]:
        proposed = self.estimate_monthly_cost(intent)
        if proposed is None:
            return None

        if baseline_intent:
            baseline = self.estimate_monthly_cost(baseline_intent)
            if baseline is None:
                return None
            return proposed - baseline
        else:
            resource_pricing = self._pricing.get(intent.resource_type)
            if not resource_pricing:
                return None
            min_cost = min(resource_pricing.values())
            return proposed - min_cost

    def estimate_cost_distribution(self, intent: ProvisionResourceIntent) -> Dict[str, float]:
        cost = self.estimate_monthly_cost(intent)
        if cost is None:
            return {}
        return {str(cost): 1.0}