import importlib
import warnings
warnings.filterwarnings('ignore')

import gymnasium as gym

from envs.wrappers.vectorized_multitask import make_vectorized_multitask_env
from envs.wrappers.render import Render


# Every domain used to be imported eagerly here, which means one missing
# dependency took down *all* tasks rather than just its own -- painful on any
# partial install, and fatal when deploying (a hosted demo cannot install a
# Vulkan ICD for ManiSkill, but should still serve the other nine domains).
# Each domain is now loaded behind a guard, generalising the pattern that was
# already used for Atari: on failure the name binds to a stub that raises the
# original error, but only if that domain's task is actually requested.
# `AVAILABLE_DOMAINS` lets a caller filter its task list up front instead.
#
# The guard catches Exception, not just ImportError, because these domains fail
# in more ways than a missing module: a native renderer with no GL/Vulkan
# context available typically raises RuntimeError or OSError at import time.
_IMPORT_ERRORS = {}
AVAILABLE_DOMAINS = set()


def _load_domain(module, domains, hint):
	"""Return envs.<module>.make_env, or a stub explaining why it is unavailable."""
	try:
		fn = importlib.import_module(f'envs.{module}').make_env
	except Exception as err:
		_IMPORT_ERRORS[module] = err

		def _unavailable(cfg, _m=module, _h=hint, _e=err):
			raise ImportError(
				f'The {_m} domain is not available in this installation ({_h}). '
				f'Original import error: {type(_e).__name__}: {_e}'
			)
		return _unavailable
	AVAILABLE_DOMAINS.update(domains)
	return fn


make_dm_control_env = _load_domain(
	'dmcontrol', ('dmcontrol', 'dmcontrol-ext'), 'needs dm-control and a headless MuJoCo GL backend, e.g. MUJOCO_GL=egl')
make_maniskill_env = _load_domain(
	'maniskill', ('maniskill',), 'needs mani_skill, and SAPIEN requires a Vulkan ICD')
make_metaworld_env = _load_domain(
	'metaworld', ('metaworld',), 'needs the metaworld fork and a headless MuJoCo GL backend')
make_mujoco_env = _load_domain(
	'mujoco', ('mujoco',), 'needs mujoco and a headless GL backend')
make_box2d_env = _load_domain(
	'box2d', ('box2d',), 'needs box2d-py (built with swig) and pygame')
make_robodesk_env = _load_domain(
	'robodesk', ('robodesk',), 'needs robodesk and a headless MuJoCo GL backend')
make_ogbench_env = _load_domain(
	'ogbench', ('ogbench',), 'needs ogbench and a headless MuJoCo GL backend')
make_pygame_env = _load_domain(
	'pygame', ('pygame',), 'needs pygame; set SDL_VIDEODRIVER=dummy when headless')
make_atari_env = _load_domain(
	'atari', ('atari',), 'needs ale_py>=0.10 for continuous-action Atari')


def unavailable_domains():
	"""Map of domain module -> the exception that stopped it importing."""
	return dict(_IMPORT_ERRORS)


def make_env(cfg):
	"""
	Make an environment for world-model experiments.
	"""
	gym.logger.set_level(40)
	if not cfg.child_env:
		env = make_vectorized_multitask_env(cfg, make_env)
	else:
		env = None
		for fn in [
			make_dm_control_env, make_maniskill_env, make_metaworld_env,
			make_mujoco_env, make_box2d_env, make_robodesk_env,
			make_ogbench_env, make_pygame_env, make_atari_env,
		]:
			try:
				env = fn(cfg)
				break
			except ValueError as e:
				if 'Unknown task' in str(e):
					continue
				else:
					raise e
		if env is None:
			raise ValueError(f'Failed to make environment "{cfg.task}": please verify that dependencies are installed and that the task exists.')
		assert cfg.num_envs == 1 or cfg.get('obs', 'state') == 'state', \
			'Vectorized environments only support state observations.'
		if cfg.save_video and cfg.get('num_demos', 0) > 0:
			env = Render(env, cfg)
		print(f'[Rank {cfg.rank}] Created env for task {cfg.task}')
	try: # Dict
		cfg.obs_shape = {k: v.shape for k, v in env.observation_space.spaces.items()}
	except: # Box
		cfg.obs_shape = {cfg.get('obs', 'state'): env.observation_space.shape}
	cfg.action_dim = env.action_space.shape[0]
	cfg.episode_length = env.max_episode_steps
	return env