"""
This script reads documentation from /pettingzoo and puts it into md files inside the docs/ directory
"""

import os
import re

from pettingzoo.utils.all_modules import all_environments


def add_frontmatter(text, frontmatter_options):
    frontmatter_text = "---\nautogenerated:"
    for key, value in frontmatter_options.items():
        frontmatter_text += f"\n{key}: {value}"
    frontmatter_text += "\n---\n\n"
    return frontmatter_text + text


def create_docs_md(file_path, text, frontmatter_options):
    text = add_frontmatter(text, frontmatter_options)
    with open(file_path, "w", encoding="utf-8") as file:
        file.write(text)


def get_docs_from_py(file_path):
    print(file_path)
    with open(file_path, encoding="utf-8") as fp:
        text = fp.read()
        regex = re.compile(r'^r?"""\s*((\n|.)*?)("""\s*\n)', re.MULTILINE)
        match = regex.search(text)
        if match:
            return match.group(1)
        else:
            return ""


if __name__ == "__main__":
    ignore_dirs = ["test", "utils"]
    docs_dir = os.path.join(os.path.dirname(__file__), "..")
    envs_dir = os.path.join(os.path.dirname(__file__), "..", "..", "pettingzoo")
    for env_type in os.listdir(envs_dir):
        env_type_path = os.path.join(envs_dir, env_type)
        if not os.path.isdir(env_type_path) or env_type in ignore_dirs:
            continue
        envs_list = os.listdir(env_type_path)
        if "__pycache__" in envs_list:
            envs_list.remove("__pycache__")

        # rlcard_envs don't follow the same folder structure
        if "rlcard_envs" in envs_list:
            envs_list.pop(envs_list.index("rlcard_envs"))
            for i in os.listdir(os.path.join(env_type_path, "rlcard_envs")):
                if (
                    not os.path.isdir(os.path.join(env_type_path, "rlcard_envs", i))
                    and i != "__init__.py"
                    and i != "rlcard_base.py"
                ):
                    envs_list.append(os.path.join("rlcard_envs", i[:-3]))
            envs_list = sorted(envs_list)

        envs_list = list(
            filter(
                lambda x: (
                    os.path.isdir(os.path.join(env_type_path, x))
                    and "utils" not in os.path.join(env_type_path, x)
                )
                or "rlcard_envs" in x,
                envs_list,
            )
        )

        for i, env_name in enumerate(envs_list):
            env_dir_path = os.path.join(env_type_path, env_name)

            if "rlcard_envs" in env_dir_path:
                env_name = env_name.replace("\\", "/").split("/")[1]

            frontmatter_options = {
                "env_icon": f'"../../../_static/img/icons/{env_type}/{env_name}.png"'
            }

            if i == 0:
                frontmatter_options["firstpage"] = ""
            elif i == len(envs_list) - 1:
                frontmatter_options["lastpage"] = ""

            docs_text = get_docs_from_py(
                os.path.join(env_dir_path, env_name + ".py")
                if "rlcard_envs" not in env_dir_path
                else env_dir_path + ".py"
            )
            # Full usage scripts for each environment page
            full_env_names = [
                full_name
                for full_name in all_environments.keys()
                if env_name in full_name
            ]
            env_name_version = full_env_names[0].split("/")[1]
            if env_type == "classic" and env_name != "rps":
                with open(f"{docs_dir}/code_examples/usage_aec_action_mask.py") as f:
                    usage = f.read()
                usage = usage.replace("connect_four_v3", env_name_version)
                docs_text += f"""## Usage
### AEC
```python
{usage}
```
"""
            else:
                with open(f"{docs_dir}/code_examples/usage_aec.py") as f:
                    usage = f.read()
                usage = usage.replace("space_invaders_v2", env_name_version)
                usage = usage.replace("atari", env_type)
                with open(f"{docs_dir}/code_examples/usage_parallel.py") as f:
                    usage_parallel = f.read()
                usage_parallel = usage_parallel.replace(
                    "pistonball_v6", env_name_version
                )
                usage_parallel = usage_parallel.replace("butterfly", env_type)
                docs_text += f"""## Usage
### AEC
```python
{usage}
```
### Parallel
```python
{usage_parallel}
```
"""
            # API information
            if "rlcard_envs" in env_dir_path:
                docs_text += f"""## API
```{{eval-rst}}
.. currentmodule:: pettingzoo.{env_type}.rlcard_envs.{env_name}

.. autoclass:: env
.. autoclass:: raw_env
   :members:
```
"""
            elif env_type in ["mpe", "atari"]:
                docs_text += f"""## API
```{{eval-rst}}
.. currentmodule:: pettingzoo.{env_type}.{env_name}.{env_name}

.. autoclass:: raw_env
   :members:
   :undoc-members:
```
"""
            else:
                docs_text += f"""## API
```{{eval-rst}}
.. currentmodule:: pettingzoo.{env_type}.{env_name}.{env_name}

.. autoclass:: env
.. autoclass:: raw_env
   :members:
```
"""
            docs_env_path = os.path.join(
                docs_dir,
                "environments",
                env_type,
                env_name + ".md",
            )
            create_docs_md(
                docs_env_path,
                docs_text,
                frontmatter_options,
            )
