#!/usr/bin/env python

import argparse
import json

import gym
from gym import envs

import gym_wumpus
from gym_wumpus.envs import WumpusEnv


def run_episode(env: WumpusEnv):
    obs = env.reset()
    total_reward = 0
    print('Observation: {}'.format(env.space_to_percept(obs)))
    for step in range(1000):
        env.render(mode='human')
        action = env.action_space.sample()     # take a random action
        obs, reward, done, info = env.step(action)
        print('{} -> {}, [{}], {}{}'.format(env.space_to_action(action), reward, env.space_to_percept(obs), 'done ' if done else '', info))
        total_reward += reward
        if done:
            break
    print('Reward {} after {} steps'.format(total_reward, step))
    env.close()


def main():
    default_env = 'wumpus-random-v0'
    parser = argparse.ArgumentParser()
    parser.add_argument('id', nargs='?', default=default_env, help='Environment name')
    parser.add_argument('--list', help='Show available environments', action="store_true", default=False)
    parser.add_argument('--file', type=argparse.FileType('r'), help='Read the JSON description of the world from the file')
    args = parser.parse_args()

    if args.list:
        print([e.id for e in envs.registry.all() if str(e.id).lower().find('wumpus') >= 0])
        return

    if args.file is not None:
        env = gym.make('wumpus-custom-v0', desc=json.load(args.file))
    else:
        try:
            env = gym.make(args.id)
        except:
            env = gym.make(default_env)
    run_episode(env)


if __name__ == "__main__":
    main()