from pathlib import Path
import yaml


def check_toc_integrity(doc_folder, output_dir):
    """
    Checks all the MDX files obtained after building the documentation are present in the table of contents.

    Args:
        doc_folder (`str` or `os.PathLike`): The folder where the source files of the documentation lie.
        output_dir (`str` or `os.PathLike`): The folder where the doc is built.
    """
    output_dir = Path(output_dir)
    doc_files = [str(f.relative_to(output_dir).with_suffix("")) for f in output_dir.glob("**/*.mdx")]

    toc_file = Path(doc_folder) / "_toctree.yml"
    with open(toc_file, "r", encoding="utf-8") as f:
        toc = yaml.safe_load(f.read())

    toc_sections = []
    sphinx_refs = []
    # We don't just loop directly in toc as we will add more into it as we un-nest things.
    while len(toc) > 0:
        part = toc.pop(0)
        if "local" in part:
            toc_sections.append(part["local"])
        if "sections" not in part:
            continue
        toc_sections.extend([sec["local"] for sec in part["sections"] if "local" in sec])
        for sec in part["sections"]:
            if "local_fw" in sec:
                toc_sections.extend(sec["local_fw"].values())
        # There should be one sphinx ref per page
        for sec in part["sections"]:
            if "local" in sec:
                sphinx_refs.append(f"{sec['local']} std:doc -1 {sec['local']} {sec['title']}")
        # Toc has some nested sections in the API doc for instance, so we recurse.
        toc.extend([sec for sec in part["sections"] if "sections" in sec])

    files_not_in_toc = [f for f in doc_files if f not in toc_sections]
    if len(files_not_in_toc) > 0:
        message = "\n".join([f"- {f}" for f in files_not_in_toc])
        raise RuntimeError(
            "The following files are not present in the table of contents:\n" + message + f"\nAdd them to {toc_file}."
        )

    files_not_exist = [f for f in toc_sections if f not in doc_files]
    if len(files_not_exist) > 0:
        message = "\n".join([f"- {f}" for f in files_not_exist])
        raise RuntimeError(
            "The following files are present in the table of contents but do not exist:\n"
            + message
            + f"\nRemove them from {toc_file}."
        )

    return sphinx_refs