from functools import partial
from smac.env import MultiAgentEnv, StarCraft2Env
from smacv2.env.starcraft2.wrapper import StarCraftCapabilityEnvWrapper
from .lavapath_wrapper import LavaPathWrapper
import sys
import os
import pdb

try:
    gfootball = True
    
    from .gfootball import GoogleFootballEnv
    
except Exception as e:
    gfootball = False
    print(e)

def env_fn(env, **kwargs) -> MultiAgentEnv:
    return env(**kwargs)

REGISTRY = {}
REGISTRY["sc2"] = partial(env_fn, env=StarCraft2Env)
REGISTRY['sc2wrapped'] = partial(env_fn, env=StarCraftCapabilityEnvWrapper)
REGISTRY['lavapath'] = partial(env_fn, env=LavaPathWrapper)


if gfootball:
    REGISTRY["gfootball"] = partial(env_fn, env=GoogleFootballEnv)

if sys.platform == "linux":
    os.environ.setdefault("SC2PATH",
                          os.path.join("~", "StarCraftII"))
