import os
import tensorflow as tf

class Config:
    # Core Environment Framework
    MODE = 'competitive'  # Options: 'competitive' or 'cooperative'
    USE_GPU = True
    DEVICE = '/GPU:0' if USE_GPU else '/CPU:0'
    NUM_ENVS = 10 if MODE == 'competitive' else 20 #10 for Air and Fire, 20 for Air and Ice
    MAX_STEPS_PER_EPISODE = 2048 
    
    # File Management Shards 
    LOG_DIR = "./DefenseAI_Competitive/logs"
    CHECKPOINT_DIR = "./DefenseAI_Competitive/checkpoints"
    CHECKPOINT_INTERVAL = 10
    
    # Neural Network Topologies 
    ACTOR_LAYERS = (128, 64)
    CRITIC_LAYERS = (128, 64)
    
    # Reinforcement Learning Core Hyperparameters
    GAMMA = 0.99
    BUFFER_MAX_LENGTH = 100000
    BATCH_SIZE = 64
    TOTAL_EPISODES = 5000

    # Programmable Cyber Range Network Topology 
    # 9 Total Node Count but can scale dynamically 
    NETWORK_TOPOLOGY = {
        "subnets": {
            "public_dmz":      {"num_hosts": 2, "base_vulnerability": 0.7},  # Entry point web targets
            "dns_services":    {"num_hosts": 1, "base_vulnerability": 0.5},  # Core DNS infrastructure (Spoofing target)
            "corporate":       {"num_hosts": 3, "base_vulnerability": 0.4},  # Standard employee workstations
            "active_directory": {"num_hosts": 2, "base_vulnerability": 0.3},  # Crown jewels: Domain Controllers (AD DC)
            "secure_core":     {"num_hosts": 1, "base_vulnerability": 0.1}   # Isolated accounting/backend databases
        }
    }

    @classmethod
    def TOTAL_HOSTS(cls):
        """Dynamically computes total node/host counts across your custom footprint."""
        return sum(subnet["num_hosts"] for subnet in cls.NETWORK_TOPOLOGY["subnets"].values())

    # Symmetrical Learning Rate Schedulers
    BASE_LR_RED = 1e-4 
    BASE_LR_BLUE = 1e-4 
    
    @classmethod
    def get_red_lr_schedule(cls):
        """Returns an exponential decay scheduler for the Red Agent."""
        return tf.keras.optimizers.schedules.ExponentialDecay(
            initial_learning_rate=cls.BASE_LR_RED,
            decay_steps=1000,
            decay_rate=0.96,
            staircase=True
        )

    @classmethod
    def get_blue_lr_schedule(cls):
        """Returns an exponential decay scheduler for the Blue Agent."""
        return tf.keras.optimizers.schedules.ExponentialDecay(
            initial_learning_rate=cls.BASE_LR_BLUE,
            decay_steps=1000,
            decay_rate=0.96,
            staircase=True
        )