import os
import re
from collections import defaultdict

import gymnasium as gym
from gymnasium.envs.registration import find_highest_version, get_env_id
from utils import trim


filtered_envs = defaultdict(list)
exclude_env_names = [
    "GymV21Environment",
    "GymV26Environment",
    "FrozenLake8x8",
    "LunarLanderContinuous",
    "BipedalWalkerHardcore",
    "phys2d/CartPole",
    "phys2d/Pendulum",
    "tabular/Blackjack",
    "tabular/CliffWalking",
]
for env_spec in gym.registry.values():
    if get_env_id(env_spec.namespace, env_spec.name, None) not in exclude_env_names:
        highest_version = find_highest_version(env_spec.namespace, env_spec.name)
        env_id = get_env_id(env_spec.namespace, env_spec.name, highest_version)

        env_spec = gym.spec(env_id)
        if (
            isinstance(env_spec.entry_point, str)
            and "gymnasium" in env_spec.entry_point
        ):
            env_module = env_spec.entry_point.split(".")[2]

            if env_spec not in filtered_envs[env_module]:
                filtered_envs[env_module].append(env_spec)

# print(filtered_envs.keys())
for env_module, env_specs in filtered_envs.items():
    env_module_name = env_module.replace("_", " ").title()
    print(env_module_name)
    env_specs = sorted(env_specs, key=lambda spec: spec.name)

    for i, env_spec in enumerate(env_specs):
        print(f"\t{i=}, {env_spec.name}")
        env = gym.make(env_spec)
        env_docstring = trim(env.unwrapped.__doc__)
        assert env_docstring

        snake_env_name = re.sub(r"(?<!^)(?=[A-Z])", "_", env_spec.name).lower()
        title_env_name = re.sub(r"(?<!^)(?=[A-Z])", " ", env_spec.name).title()

        if i == 0:
            related_pages_meta = "firstpage:\n"
        elif i == len(env_specs) - 1:
            related_pages_meta = "lastpage:\n"
        else:
            related_pages_meta = ""

        action_space_table = env.action_space.__repr__().replace("\n", "")
        observation_space_table = env.observation_space.__repr__().replace("\n", "")

        env_page = f"""---
autogenerated:
title: {title_env_name}
{related_pages_meta}---

# {title_env_name}

```{{figure}} ../../_static/videos/{env_module}/{snake_env_name}.gif
:width: 200px
:name: {snake_env_name}
```

This environment is part of the {env_module_name} environments which contains general information about the environment.

|   |   |
|---|---|
| Action Space | `{re.sub(' +', ' ', action_space_table)}` |
| Observation Space | `{re.sub(' +', ' ', observation_space_table)}` |
| import | `gymnasium.make("{env.spec.id}")` |

{env_docstring}
"""

        env_md_path = os.path.join(
            os.path.dirname(__file__),
            "..",
            "environments",
            env_module,
            snake_env_name + ".md",
        )
        file = open(env_md_path, "w", encoding="utf-8")
        file.write(env_page)
        file.close()