import os
os.environ["TF_USE_LEGACY_KERAS"] = "1"
import psutil
import numpy as np
import tensorflow as tf
from tf_agents.environments import py_environment
from tf_agents.specs import array_spec
from tf_agents.trajectories import time_step as ts
from config import Config

try:
    import pynvml
    pynvml.nvmlInit()
    HAS_NVML = True
except ImportError:
    HAS_NVML = False

def configure_gpu():
    """Configures TensorFlow memory growth to prevent VRAM allocation errors."""
    gpus = tf.config.list_physical_devices('GPU')
    if gpus:
        try:
            for gpu in gpus:
                tf.config.experimental.set_memory_growth(gpu, True)
            print(f"✅ Found and configured {len(gpus)} GPU(s).")
        except RuntimeError as e:
            print(e)
    else:
        print("⚠️ No GPU found. Falling back to CPU.")

def get_system_telemetry():
    """Fetches real-time RAM, CPU, and GPU usage metrics for TensorBoard profiling."""
    metrics = {
        'system/ram_percent': psutil.virtual_memory().percent,
        'system/cpu_percent': psutil.cpu_percent(),
    }
    
    if HAS_NVML:
        try:
            handle = pynvml.nvmlDeviceGetHandleByIndex(0)
            util = pynvml.nvmlDeviceGetUtilizationRates(handle)
            temp = pynvml.nvmlDeviceGetTemperature(handle, pynvml.NVML_TEMPERATURE_GPU)
            mem = pynvml.nvmlDeviceGetMemoryInfo(handle)
            
            metrics['hardware/gpu_utilization'] = float(util.gpu)
            metrics['hardware/gpu_temperature_c'] = float(temp)
            metrics['hardware/gpu_memory_percent'] = (mem.used / mem.total) * 100.0
        except Exception:
            pass
            
    return metrics

class PettingZooToTFAgentsWrapper(py_environment.PyEnvironment):
    """Wraps a PettingZoo Parallel environment into a TF-Agents PyEnvironment."""
    
    def __init__(self, pz_env):
        super().__init__()
        self._env = pz_env
        self._current_info = {} 
        
        # Read host array shape sizes straight out of Config
        total_nodes = Config.TOTAL_HOSTS()
        
        # Red Actions: total_nodes * 4 choices | Blue Actions: total_nodes * 2 choices
        self._action_spec = {
            'red': array_spec.BoundedArraySpec(shape=(), dtype=np.int32, minimum=0, maximum=(total_nodes * 4) - 1, name='red_action'),
            'blue': array_spec.BoundedArraySpec(shape=(), dtype=np.int32, minimum=0, maximum=(total_nodes * 2) - 1, name='blue_action')
        }
        
        # Match the complex dictionary nested schema
        self._observation_spec = {
            agent: {
                'node_features': array_spec.BoundedArraySpec(shape=(total_nodes, 2), dtype=np.float32, minimum=0, maximum=1, name='features'),
                'adjacency_matrix': array_spec.BoundedArraySpec(shape=(total_nodes, total_nodes), dtype=np.float32, minimum=0, maximum=1, name='topology')
            } for agent in ['red', 'blue']
        }

    def action_spec(self): 
         return self._action_spec

    def observation_spec(self): 
         return self._observation_spec

    def reward_spec(self):
        """Declares multi-agent reward specs to prevent unpacking errors."""
        return {
            'red': array_spec.ArraySpec(shape=(), dtype=np.float32, name='red_reward'),
            'blue': array_spec.ArraySpec(shape=(), dtype=np.float32, name='blue_reward')
        }

    def get_current_info(self):
        """Safely forwards custom internal metadata down to the trainer."""
        return self._current_info

    def _reset(self):
        obs, infos = self._env.reset()
        self._current_info = infos
        
        # Guarantee observations map explicitly to nested dictionary signatures
        formatted_obs = {
            'red': {
                'node_features': np.array(obs['red']['node_features'], dtype=np.float32),
                'adjacency_matrix': np.array(obs['red']['adjacency_matrix'], dtype=np.float32)
            },
            'blue': {
                'node_features': np.array(obs['blue']['node_features'], dtype=np.float32),
                'adjacency_matrix': np.array(obs['blue']['adjacency_matrix'], dtype=np.float32)
            }
        }
        
        initial_rewards = {
            'red': np.array(0.0, dtype=np.float32),
            'blue': np.array(0.0, dtype=np.float32)
        }
        
        return ts.TimeStep(
            step_type=np.array(ts.StepType.FIRST, dtype=np.int32),
            reward=initial_rewards,
            discount=np.array(1.0, dtype=np.float32),
            observation=formatted_obs
        )

    def _step(self, action):
        # Convert actions into string keys for env.py mapping compliance
        pz_actions = {'red': int(action['red']), 'blue': int(action['blue'])}
        obs, rewards, terminations, truncations, infos = self._env.step(pz_actions)
        
        self._current_info = infos
        
        # Re-enforce nested array typing explicitly on step updates
        formatted_obs = {
            'red': {
                'node_features': np.array(obs['red']['node_features'], dtype=np.float32),
                'adjacency_matrix': np.array(obs['red']['adjacency_matrix'], dtype=np.float32)
            },
            'blue': {
                'node_features': np.array(obs['blue']['node_features'], dtype=np.float32),
                'adjacency_matrix': np.array(obs['blue']['adjacency_matrix'], dtype=np.float32)
            }
        }
        
        formatted_rewards = {
            'red': np.array(rewards['red'], dtype=np.float32),
            'blue': np.array(rewards['blue'], dtype=np.float32)
        }
        
        if terminations['red'] or truncations['red']:
            return ts.termination(formatted_obs, reward=formatted_rewards)
            
        return ts.transition(formatted_obs, reward=formatted_rewards, discount=np.array(0.99, dtype=np.float32))