import yaml
import psutil
import datetime

class PolicyManager:
    def __init__(self, config_path='aiva/config/throttling_policies.yaml'):
        self.config_path = config_path
        self.policies = self._load_policies()

    def _load_policies(self):
        try:
            with open(self.config_path, 'r') as f:
                config = yaml.safe_load(f)
            return config['policies']
        except FileNotFoundError:
            print(f"Error: Configuration file not found at {self.config_path}")
            return []
        except yaml.YAMLError as e:
            print(f"Error: Failed to parse YAML file: {e}")
            return []
        except KeyError:
            print("Error: 'policies' key not found in configuration file.")
            return []

    def reload_policies(self):
        self.policies = self._load_policies()

    def get_applicable_policy(self):
        for policy in self.policies:
            if policy.get('enabled', True) and self._is_policy_applicable(policy):
                return policy
        return None

    def _is_policy_applicable(self, policy):
        threshold_type = policy['threshold_type']
        threshold_value = policy['threshold_value']
        threshold_operator = policy['threshold_operator']

        if threshold_type == 'cpu':
            cpu_usage = psutil.cpu_percent(interval=0.1)
            return self._evaluate_condition(cpu_usage, threshold_value, threshold_operator)
        elif threshold_type == 'memory':
            memory_usage = psutil.virtual_memory().percent
            return self._evaluate_condition(memory_usage, threshold_value, threshold_operator)
        elif threshold_type == 'time':
            now = datetime.datetime.now().strftime('%H:%M')
            # Assuming threshold_value is a string like '22:00-06:00'
            start_time, end_time = threshold_value.split('-')
            return self._evaluate_time_condition(now, start_time, end_time, threshold_operator)
        elif threshold_type == 'always':
            return self._evaluate_condition(True, threshold_value, threshold_operator)
        else:
            print(f"Warning: Unknown threshold type: {threshold_type}")
            return False

    def _evaluate_condition(self, actual_value, threshold_value, operator):
        try:
            threshold_value = float(threshold_value)
            actual_value = float(actual_value)
        except ValueError:
            pass # Keep as strings if conversion fails

        if operator == '>=':
            return actual_value >= threshold_value
        elif operator == '<=':
            return actual_value <= threshold_value
        elif operator == '>':
            return actual_value > threshold_value
        elif operator == '<':
            return actual_value < threshold_value
        elif operator == '==':
            return actual_value == threshold_value
        else:
            print(f"Warning: Unknown operator: {operator}")
            return False

    def _evaluate_time_condition(self, current_time, start_time, end_time, operator):
        if operator == 'in':
            # Convert times to datetime.time objects for comparison
            current_time_obj = datetime.datetime.strptime(current_time, '%H:%M').time()
            start_time_obj = datetime.datetime.strptime(start_time, '%H:%M').time()
            end_time_obj = datetime.datetime.strptime(end_time, '%H:%M').time()

            # Handle cases where the time range crosses midnight
            if start_time_obj <= end_time_obj:
                return start_time_obj <= current_time_obj <= end_time_obj
            else:
                # Range crosses midnight, so it's either after start or before end
                return current_time_obj >= start_time_obj or current_time_obj <= end_time_obj
        else:
            print(f"Warning: Unknown operator: {operator}")
            return False

if __name__ == '__main__':
    policy_manager = PolicyManager()
    applicable_policy = policy_manager.get_applicable_policy()
    if applicable_policy:
        print(f"Applicable Policy: {applicable_policy['name']}")
    else:
        print("No applicable policy found.")