# LINE 1-18: MUST BE ABSOLUTE FIRST TO FORCE LEGACY KERAS NAMESPACES
import os
import warnings

os.environ["TF_USE_LEGACY_KERAS"] = "1"      
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"   
os.environ["TF_FUNCTION_NUMERIC_CHECKS"] = "0"

import tensorflow as tf

tf.keras.backend.clear_session()
tf.config.threading.set_intra_op_parallelism_threads(0) 
tf.config.threading.set_inter_op_parallelism_threads(0) 
tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.ERROR)

warnings.filterwarnings("ignore", category=UserWarning, module="gym")
warnings.filterwarnings("ignore", category=RuntimeWarning, module="tf_agents.specs.array_spec")
warnings.filterwarnings("ignore", category=DeprecationWarning)

from tf_agents.environments import tf_py_environment
from config import Config
from env import CyberPettingZooEnv
from utilities import configure_gpu, PettingZooToTFAgentsWrapper
from main import create_ppo_agent

def find_most_secure_checkpoint():
    """Parses era_history.txt and finds the step with the highest Blue score."""
    txt_path = os.path.join(Config.LOG_DIR, "era_history.txt")
    if not os.path.exists(txt_path):
        print(f"🚨 No era records found at {txt_path}. Defaulting to latest available checkpoint.")
        return None
    
    best_blue_score = float('-inf')
    best_step = None
    
    with open(txt_path, "r") as f:
        for line in f:
            if not line.strip(): continue
            parts = {item.split(":")[0]: item.split(":")[1] for item in line.strip().split("|")}
            blue_score = float(parts["BLUE_REWARD"])
            step_val = int(parts["STEP"])
            if blue_score > best_blue_score:
                best_blue_score = blue_score
                best_step = step_val
                
    print(f"🎯 Analysis Complete! Most secure Era found with Blue Reward: {best_blue_score:.2f} at Step: {best_step}")
    return best_step

def evaluate():
    configure_gpu()
    
    # INTERACTIVE USER PROMPT
    print("\n--- GNN Cyber Range Evaluator Setup ---")
    try:
        user_input = input("🎮 Enter the number of trial matches you want to run (default: 5): ").strip()
        num_matches = int(user_input) if user_input else 5
        if num_matches <= 0:
            print("⚠️ Number must be greater than 0. Defaulting to 5 matches.")
            num_matches = 5
    except ValueError:
        print("⚠️ Invalid input detected (must be an integer). Defaulting to 5 matches.")
        num_matches = 5
    
    # 1. Spin up the environment wrapper
    raw_pz_env = CyberPettingZooEnv()
    py_env = PettingZooToTFAgentsWrapper(raw_pz_env)
    tf_env = tf_py_environment.TFPyEnvironment(py_env)
    
    # 2. Reconstruct the structural agent wrappers
    red_agent = create_ppo_agent(tf_env.time_step_spec(), tf_env.action_spec(), 'red', 1e-4)
    blue_agent = create_ppo_agent(tf_env.time_step_spec(), tf_env.action_spec(), 'blue', 1e-4)
    
    # 3. Handle Targeted Weight Restorations
    global_step = tf.Variable(0, dtype=tf.int64)
    checkpoint = tf.train.Checkpoint(
        step=global_step,
        agent_1_actor=red_agent._actor_net,
        agent_1_critic=red_agent._value_net,
        agent_2_actor=blue_agent._actor_net,
        agent_2_critic=blue_agent._value_net
    )
    
    target_step = find_most_secure_checkpoint()
    
    if target_step:
        # Construct path to specific historical shard file matching target_step
        specific_shard = os.path.join(Config.CHECKPOINT_DIR, f"ckpt-{target_step}")
        if os.path.exists(specific_shard + ".index"):
            checkpoint.restore(specific_shard)
            print(f"🔒 Loaded optimal defense policy from: {specific_shard}")
        else:
            print(f"⚠️ Shard file ckpt-{target_step} missing. Restoring latest available checkpoint.")
            tf.train.LatestCheckpointManagement(Config.CHECKPOINT_DIR).restore(checkpoint)
    else:
        # Fallback to standard latest checkpoint manager behavior if no log file exists
        manager = tf.train.CheckpointManager(checkpoint, Config.CHECKPOINT_DIR, max_to_keep=1000)
        if manager.latest_checkpoint:
            checkpoint.restore(manager.latest_checkpoint)
            print(f"🔄 Loaded latest standard fallback checkpoint: {manager.latest_checkpoint}")
            
    # 4. Run Validation Test Matches
    print(f"\n⚔️ Commencing {num_matches} evaluation trials against the targeted network layout...")
    for match in range(1, num_matches + 1):
        time_step = tf_env.reset()
        steps = 0
        while not time_step.is_last() and steps < Config.MAX_STEPS_PER_EPISODE:
            steps += 1
            
            # Isolate both the Observation AND the single-agent Scalar Reward
            red_ts = time_step._replace(
                observation=time_step.observation['red'],
                reward=time_step.reward['red']
            )
            blue_ts = time_step._replace(
                observation=time_step.observation['blue'],
                reward=time_step.reward['blue']
            )
            
            # Query policies with perfectly aligned TimeStep specs
            a1 = red_agent.policy.action(red_ts)
            a2 = blue_agent.policy.action(blue_ts)
            
            time_step = tf_env.step({'red': a1.action, 'blue': a2.action})
            
        env_info = tf_env.pyenv.envs[0].get_current_info()
        print(f"🏁 Trial Match {match}/{num_matches} Complete | Duration: {steps} steps | Final Intrusion Delta: {env_info['score_difference']}")

if __name__ == "__main__":
    evaluate()