import os
from server import AegisGymEnv
from models import AuditAction

def main():
    print("--- AegisGym Final Baseline Inference Simulation (Official API) ---")
    
    env = AegisGymEnv()
    obs = env.reset()
    
    total_score = 0.0
    episodes = 10
    
    for i in range(episodes):
        task_name = f"baseline_episode_{i+1}"
        tier = env.state.current_tier
        print(f"\nEpisode {i+1} | Tier: {tier.upper()}")
        print(f"[START] task={task_name}", flush=True)
        
        # Mocking an agent's decision based on the tier
        if tier == "easy":
            action = AuditAction(
                action_type="FLAG",
                target_id="ACC-BL-001",
                regulation_citation="Sanctions List Match"
            )
        elif tier == "medium":
            action = AuditAction(
                action_type="FLAG",
                target_id="ACC-SMURF-99",
                regulation_citation="Structuring/Smurfing 9000-10000 range"
            )
        else:
            action = AuditAction(
                action_type="FLAG",
                target_id="ACC-REG-VIOLATOR",
                regulation_citation="EU-AI-Act-Art-57"
            )
            
        # step() returns an AuditObservation instance only
        obs = env.step(action)
        print(f"Action Taken: {action.action_type} for {action.target_id}", flush=True)
        print(f"Reward: {obs.reward} | Done: {obs.done}", flush=True)
        
        print(f"[STEP] step=1 reward={obs.reward}", flush=True)
        print(f"[END] task={task_name} score={obs.reward} steps=1", flush=True)
        
        total_score += obs.reward
        if obs.done:
            break
            
    print(f"\n--- Reproducibility Report ---")
    print(f"Total Episodes: {min(episodes, env.state.step_count)}")
    print(f"Custom Agent Mean Score (Reward): {total_score / min(episodes, env.state.step_count)}")

if __name__ == "__main__":
    main()