from __future__ import annotations

import os
import re
from itertools import chain

from gymnasium.envs.registration import registry
from tqdm import tqdm

from utils import env_name_format, trim

readme_path = os.path.join(
    os.path.dirname(os.path.dirname(os.path.dirname(__file__))),
    "README.md",
)

all_envs = list(registry.values())

filtered_envs_by_type = {}
env_names = []
babyai_envs = {}

# Obtain filtered list
for env_spec in tqdm(all_envs):
    if isinstance(env_spec.entry_point, str):
        # minigrid.envs:Env
        split = env_spec.entry_point.split(".")
        # ignore gymnasium.envs.env_type:Env
        env_module = split[0]

        if len(split) > 2 and "babyai" in split[2]:
            curr_babyai_env = split[2]
            babyai_env_name = curr_babyai_env.split(":")[1]
            babyai_envs[babyai_env_name] = env_spec
        elif env_module == "minigrid":
            env_name = split[1]
            filtered_envs_by_type[env_name] = env_spec
        # if env_module != "minigrid":
    else:
        continue


filtered_envs = {
    env[0]: env[1]
    for env in sorted(
        filtered_envs_by_type.items(),
        key=lambda item: item[1].entry_point.split(".")[1],
    )
}

filtered_babyai_envs = {
    env[0]: env[1]
    for env in sorted(
        babyai_envs.items(),
        key=lambda item: item[1].entry_point.split(".")[1],
    )
}

for env_name, env_spec in chain(filtered_envs.items(), filtered_babyai_envs.items()):
    env = env_spec.make()

    docstring = trim(env.unwrapped.__doc__)

    # minigrid.envs:Env or minigrid.envs.babyai:Env
    split = env_spec.entry_point.split(".")
    # ignore minigrid.envs.env_type:Env
    env_module = split[0]
    env_name = split[-1].split(":")[-1]
    env_type = env_module if len(split) == 2 else split[-1].split(":")[0]

    path_name = ""
    os.makedirs(
        os.path.join(
            os.path.dirname(os.path.dirname(__file__)), "environments", env_type
        ),
        exist_ok=True,
    )

    v_path = os.path.join(
        os.path.dirname(os.path.dirname(__file__)),
        "environments",
        env_type,
        f"{env_name}.md",
    )

    formatted_env_name = env_name_format(env_name)

    # Front matter
    front_matter = f"""---
autogenerated:
title: {formatted_env_name}
---
"""
    # Title and gif
    title = f"# {formatted_env_name}"
    gif = (
        "```{figure} "
        + f"""/_static/videos/{env_type}/{env_name}.gif
:alt: {formatted_env_name}
:width: 200px
```
"""
    )

    # Environment attributes
    action_space_table = env.action_space.__repr__().replace("\n", "")
    observation_space_table = env.observation_space.__repr__().replace("\n", "")
    env_attributes = f"""
|   |   |
|---|---|
| Action Space | `{re.sub(' +', ' ', action_space_table)}` |
| Observation Space | `{re.sub(' +', ' ', observation_space_table)}` |
| Reward Range | `{env.reward_range}` |
| Creation | `gymnasium.make("{env_spec.id}")` |
"""

    # Create Markdown file content
    if docstring is None:
        docstring = "No information provided"
    all_text = f"""{front_matter}

{title}

{gif}

{env_attributes}

{docstring}
"""
    file = open(v_path, "w+", encoding="utf-8")
    file.write(all_text)
    file.close()