#!/usr/bin/env python3

'''
This script renders Jinja comments found at the beginning
of `.sls` files as state documentation.

The code is horrible, but I'm too lazy to rewrite it properly
because it works for my use cases. There will likely be a lot of bugs.

Full configuration in ``pyproject.toml`` with defaults:

[tool.slsdoc]
  # Choose which renderer to use. `single` or `depth`
renderer = "single"
  # Also list states without documentation
include_empty = false
  # Path prefixes to ignore, relative to formula root
ignore = ["_mapdata"]
  # Given a directory with `package.init`, `package.install`
  # and `package.clean`, only dump init and clean.
avoid_init_dup = True
  # By default, states are ordered alphabetically, but init
  # always comes first and all clean states are separated.
  # This overrides the purely alphabetical order for the roots.
sort_list = ["init", "clean", "package", "config", "service"]

[tool.slsdoc.single_renderer]
  # Include init.sls files in state list
init = true
  # Include .init suffix in when init is true
init_suffix = false
  # Include clean.sls files in state list
clean = true
  # When rendering with prefix `settings.something`,
  # write `else` instead of `settings.something.else`
cut_prefix = false
  # Write `salt.package` instead of `package` when formula root == `salt`
include_formula_root = true
  # Default .rst main header, when not defined in root init.sls
default_header = "Available States"
  # Header, overrides everything
header = ""
  # Description, overrides everything
description = ""
  # Default description, when init.sls is empty
default_description = """
.. contents::
   :local:
"""
  # Always append default description, regardless of
  # presence in init.sls
always_include_default_description = false
  # Character to use to underline main headings
main_heading_char = "="
  # Character to use to underline sub headings
sub_heading_char = "^"
"""
  # File to write to
target = "docs/states.rst"

[tool.slsdoc.depth_renderer]
  # maxdepth for generated index.rst files (toctree)
index_maxdepth = 4
  # Directory to write to
target = "docs/states"

[tool.slsdoc.depth_renderer.single_opts]
  # Include init.sls files in state list
init = false
  # Include .init suffix in when init is true
init_suffix = false
  # Include clean.sls files in state list
clean = false
  # When rendering with prefix `settings.something`,
  # write `else` instead of `settings.something.else`
cut_prefix = true
  # Write `salt.package` instead of `package` when formula root == `salt`
include_formula_root = false
'''

from pathlib import Path

import toml
from jinja2 import Environment
from jinja2.exceptions import TemplateSyntaxError

# import sphinx.util.docstrings as d

base_dir = Path(__file__).parent.parent.absolute()


class State:
    """
    Object representing a documented state.
    """

    def __init__(self, path, comment, formula_root):
        if not isinstance(path, Path):
            path = Path(path)
        self.path = path
        self.comment = str(comment)
        self.formula_root = formula_root

    def to_state_id(self, include_init=True, relative_to="", split=False, formula_root=False):
        ret = self.path.with_suffix("")

        if not include_init and ret.parts[-1] == "init":
            ret = ret.parent

        if relative_to:
            relative_to = Path(relative_to.replace(".", "/"))
            ret = ret.relative_to(relative_to)
        elif formula_root:
            ret = Path(self.formula_root) / ret

        ret = str(ret).replace("/", ".")
        if not split:
            return ret
        return tuple(ret.split("."))

    def is_init(self):
        return self.path.parts[-1] == "init.sls"

    def is_clean(self):
        return self.path.parts[-1] == "clean.sls"

    def __str__(self):
        return self.to_state_id()


class StateRegistry:
    """
    Registers State objects. It generates several indices:

    1. nested dictionary containing the states
    2. flattened list of states
    3. flattened list of prefixes ~ directories containing .sls files

    All of them are guaranteed to be sorted, which is required
    by some of the downstream classes.
    """

    def __init__(self, reg=None, states=None, prefixes=None, default_sort_fn=None, avoid_init_dup=True, sort_list=None):
        # nested dict
        self.reg = reg or {}
        # flat dict with ptr to reg
        self.states = states or {}
        # flat dict with ptr to reg
        self.prefixes = prefixes or {}
        self.default_sort_fn = default_sort_fn or self._default_sort
        # ensure no prefix lists the root
        self.prefixes[()] = self.reg
        self.avoid_init_dup = avoid_init_dup
        self.sort_list = tuple(sort_list or ("init", "clean", "package", "config", "service"))

    def _default_sort(self, a):
        """
        Sort prefixes package > config > service, rest alphabetically.
        Init always comes first. Clean states are separated.
        """
        sort_key = ()
        if a[0] in self.sort_list:
            sort_key += (str(self.sort_list.index(a[0])),)
        else:
            sort_key += (str(len(self.sort_list)),)
        if a[-1] == "init":
            sort_key += a[:-1]
        elif a[-1] == "clean":
            sort_key = ("zzzzzzzzz",) + sort_key + a[:-1]
        else:
            sort_key += a
        return sort_key

    def lookup(self, state, tree=False):
        """
        Looks up State objects by ID ('some.example' or ['some', 'example']).
        If tree, also returns pointers to dicts inside the nested registry.
        """
        ptr = self.reg
        if isinstance(state, str):
            state = state.split(".")
        state = list(state)

        try:
            while state:
                cur = state.pop(0)
                ptr = ptr[cur]
        except KeyError:
            return False

        if isinstance(ptr, State) or tree:
            return ptr
        # map init.sls same way as salt does
        elif "init" in ptr:
            return ptr["init"]

        raise ValueError("State could not be found. Most likely, this is a prefix.")

    def register(self, state):
        try:
            state_id = state.to_state_id(split=True)
        except AttributeError as err:
            raise ValueError("`state` argument must be an instance of `State`") from err

        ptr = self.reg
        prefix = []
        parts = list(state_id)

        while len(parts):
            cur = parts.pop(0)
            prefix = prefix + [cur]
            if not parts:
                ptr[cur] = state
                break
            # make sure subdicts are initialized
            elif cur not in ptr:
                ptr[cur] = {}
                self.prefixes[tuple(prefix)] = ptr[cur]
            ptr = ptr[cur]
        self.states[state_id] = ptr[cur]

    def unregister(self, state):
        try:
            state_id = state.to_state_id(split=True)
        except AttributeError as err:
            raise ValueError("`state` argument must be an instance of `State`") from err
        try:
            self.states.pop(state_id)
        except KeyError as err:
            raise RuntimeError(f"Specified state {state} was not registered") from err

        parts = list(state_id)
        first_run = True

        while parts:
            key = parts.pop()
            parent = self.lookup(parts, tree=True)
            if first_run:
                del parent[key]
            elif not len(parent[key]):
                self.prefixes.pop(tuple(parts + [key]))
                del parent[key]
            first_run = False

    def depth(self, prefix=""):
        """
        Returns depth of nested dict. Optionally from a specified prefix.
        """
        d = self.lookup(prefix, tree=True) if prefix else self.reg
        return self._get_depth(d)

    def _get_depth(self, d):
        """
        Recursively look up depth.
        """
        if isinstance(d, dict) and set(d.keys()) - {"init", "clean"}:
            return 1 + (max(map(self._get_depth, d.values())) if d else 0)
        return 0

    def count(self):
        return len(self.states)

    def dump(self, prefix=None, sort_fn=None):
        """
        Returns all states sorted, optionally beginning from a prefix only.
        """
        prefix = prefix or ()
        if isinstance(prefix, str):
            prefix = tuple(prefix.split("."))
        for state_id in sorted(self.states, key=sort_fn or self.default_sort_fn):
            if state_id[:len(prefix)] == prefix:
                if self.avoid_init_dup and state_id[-1] not in ["init", "clean"]:
                    # given package.init, package.install and package.clean,
                    # only dump init and clean
                    parent = self.lookup(state_id[:-1], tree=True)
                    if len(parent) == 3 and "init" in parent and "clean" in parent:
                        continue
                yield self.states[state_id]


class Renderer:
    """
    Base class for Renderer objects.
    """

    _default_opts = {
        "default_description": """\
The following states are found in {where}:

.. contents::
   :local:
""",
        "always_include_default_description": False,
        "default_header": "Available States",
        "main_heading_char": "=",
        "sub_heading_char": "-",
    }

    def __init__(self, target="docs", prefix=None, basedir=base_dir, opts=None):
        self.opts = self._default_opts | (opts or {})
        self.target = basedir / target
        if isinstance(prefix, str):
            prefix = tuple(prefix.split("."))
        self.prefix = prefix or ()
        self.basedir = basedir
        self._ensurepath(self.target)

    def render(self, states):
        raise NotImplementedError()

    def _write(self, text, ncount=1):
        """
        Convenience method to mimic writeln
        This is bad form though, implicitly requires
        self.f to be set before being called
        """
        if not hasattr(self, "f"):
            raise ValueError(
                "Sorry, you need to set self.f to a file pointer before calling self._write methods."
            )
        self.f.write(text + "\n" * ncount)

    def _write_heading(self, text, uline=None):
        """
        Convenience method to write a subheading
        """
        self._write(text)
        self._write((uline or self.opts["sub_heading_char"]) * len(text))

    def _write_main_heading(self, text, uline=None):
        """
        Convenience method to write the main heading
        """
        self.f.write(text + "\n")
        self.f.write((uline or self.opts["main_heading_char"]) * len(text) + "\n\n")

    def _write_main_description(self, text):
        """
        Convenience method to write the main description
        """
        self.f.write(text + "\n")

    def _get_main_heading(self, states):
        """
        Returns a string to be used as the main heading of an
        rst file. Looks up in the following order:
        1. If there is a documented root init.sls, return its first line.
        2. If a header was specified in opts, return that.
        3. If this renderer was instantiated with a prefix, return the
           capitalized last part of it.
        4. "Available States"
        """
        if header := self.opts.get("header"):
            return header
        if (main := states.lookup(self.prefix + ("init",))) and main.comment:
            lines = main.comment.splitlines()
            header = lines.pop(0)
            # make sure the header does not get printed twice
            main.comment = "\n".join(lines).strip("\n")
            return header
        if self.prefix:
            return self.prefix[-1].title()
        if header := self.opts.get("default_header"):
            return header
        return self.opts["default_header"]

    def _get_main_description(self, states):
        """
        Returns a string to be used as the main description (after the main header).
        Looks up in the following order:
        1. If there is a documented root init.sls, return everything but its first line.
           If it was found, unregister it to not render it twice (bad form again, maybe consider was_rendered on State)
        2. If a description was specified in opts, return that.
        3. general description
        """
        def get_default_desc():
            return self.opts["default_description"].format(
                where=".".join(self.prefix) if self.prefix else "this formula"
            )

        if desc := self.opts.get("description"):
            return desc
        if main := states.lookup(self.prefix + ("init",)):
            desc = main.comment
            if self.opts.get("always_include_default_description"):
                desc += "\n" + get_default_desc()
            return desc
        return get_default_desc()

    def _write_header(self, states):
        """
        Convenience method to write the main heading and description.
        """
        self._write_main_heading(self._get_main_heading(states))
        self._write_main_description(self._get_main_description(states))
        self.f.write("\n")

    def _ensurepath(self, path):
        """
        Convenience method to mimic mkdir -p on the parent dir of a path
        """
        if not path.parent.is_dir():
            path.parent.mkdir(parents=True)


class SingleRenderer(Renderer):
    """
    Renders documented Salt states in a single file. When a prefix
    is specified, it only considers states nested inside that prefix.

    If there is a documented top-layer init.sls, the first line of
    its comment will become the main heading of the file, the stripped
    rest will be rendered directly below before any nested states as
    the main description.

    Should have been a Jinja template @TODO
    """

    default_opts = {
        # include init.sls files
        "init": True,
        # include .init suffix
        "init_suffix": False,
        # include clean.sls files
        "clean": True,
        # when rendering with prefix `settings.something`,
        # write `else` instead of `settings.something.else`
        "cut_prefix": False,
        "include_formula_root": True,
    }

    def __init__(self, target="docs/states.rst", prefix=None, basedir=base_dir, **kwargs):
        if isinstance(target, str):
            target = Path(target)
        target = target.with_suffix(".rst")
        super().__init__(target=target, prefix=prefix, basedir=basedir, opts=self.default_opts | kwargs)

    def render(self, states):
        relative_to = "/".join(self.prefix) if self.opts["cut_prefix"] else ""

        with open(self.target, "w") as f:
            self.f = f

            self._write_header(states)

            for state in states.dump(self.prefix):
                if not self.opts["init"] and state.is_init():
                    continue
                if not self.opts["clean"] and state.is_clean():
                    continue
                state_id = state.to_state_id(self.opts["init_suffix"], relative_to, formula_root=self.opts["include_formula_root"])
                if state_id == ".":
                    continue
                self._write_heading(
                    state.to_state_id(self.opts["init_suffix"], relative_to, formula_root=self.opts["include_formula_root"])
                )
                self._write(state.comment, 3)


class DepthRenderer(Renderer):
    """
    Renders documented Salt states in separate files. They are combined
    according to the maximum nesting depth specified during instantiation.

    Consider an example with max_depth=1.

    This directory structure:
        settings
        ├── animations
        │   ├── cursor_blinking
        │   └── window_resize_time
        └── apps
            └── messages
                └── read_receipts

    will result in these files:
        * docs/settings/index.rst
            - animations
            - apps/index
        * docs/settings/animations.rst
            - cursor_blinking doc
            - window_resize_time doc
        * docs/settings/apps/index.rst
            - messages
        * docs/settings/apps/messages.rst
            - read_receipts doc
    """

    default_opts = {
        # maxdepth for generated index.rst files
        "index_maxdepth": 4,
        "single_opts": {
            "init": True,
            "init_suffix": False,
            "clean": False,
            "cut_prefix": True,
            "include_formula_root": False,
        }
    }

    def __init__(self, max_depth=1, target="docs/states", delegate_cls=None, basedir=base_dir, **kwargs):
        self.max_depth = max_depth
        self.delegate = delegate_cls or SingleRenderer
        super().__init__(target=target, basedir=basedir, opts=self.default_opts | (kwargs or {}))

    def render(self, states):
        found = []
        indices = {}
        cur_top = ()

        # Find at which prefix length the maxdepth is fulfilled
        for prefix in states.prefixes:
            # Filtering out more specific = shallower dicts
            # This requires the list to be sorted!
            # This also goes not too well with rendering clean states
            # by default since they will be separate.
            if cur_top:
                if prefix[:len(cur_top)] == cur_top:
                    continue
                else:
                    cur_top = ()
            if states.depth(prefix) <= self.max_depth:
                found.append(prefix)
                cur_top = prefix

        # fill indices with relative links to
        #   a) other indices if there are nested prefixes
        #   b) the rendered files at that depth
        # example: settings.one settings.two.[subone, subtwo] =>
        # { settings: {'_index': ['one', 'two/index'],
        #              'two': {'_index': [subone, subtwo]}},
        #  '_index': 'settings/index'}

        for prefix in found:
            ptr = indices
            prefix = list(prefix)

            while prefix:
                cur = prefix.pop(0)
                if prefix and cur not in ptr:
                    ptr[cur] = {}
                if "_index" not in ptr:
                    ptr["_index"] = []
                if prefix and cur + "/index" not in ptr["_index"]:
                    ptr["_index"].append(cur + "/index")
                elif not prefix and cur not in ptr["_index"]:
                    ptr["_index"].append(cur)
                if prefix:
                    ptr = ptr[cur]

        if "_index" in indices:
            self._render_index(states, indices)

        for prefix in found:
            self._render_single(states, prefix)

    def _render_single(self, states, prefix):
        target = self._get_target(prefix)
        renderer = self.delegate(target=target, prefix=prefix, **self.opts["single_opts"])
        renderer.render(states)

    def _render_index(self, states, indices, prefix=None):
        prefix = prefix or ()
        target = self._get_target(prefix) / Path("index.rst")
        with open(target, "w") as f:
            self.f = f
            self.prefix = prefix
            self._write_header(states)
            self.prefix = ()
            f.write(f"""
.. toctree::
   :maxdepth: {self.opts['index_maxdepth']}
   :caption: Contents:

""")
            for s in sorted(indices["_index"]):
                f.write("   " + s + "\n")
        for sub in sorted(indices):
            if sub == "_index":
                continue
            if "_index" in indices[sub]:
                new_prefix = prefix + (sub,)
                self._render_index(states, indices[sub], new_prefix)

    def _get_target(self, prefix):
        target_prefix = Path("/".join(prefix))
        target = self.target / target_prefix
        self._ensurepath(target)
        return target


class SaltAutodoc:
    """
    Looks for salt state files (*.sls) relative to base,
    reads them and looks for the first token being a Jinja
    comment. If the comment does not start with the string
    '@internal', it will be saved in the registry.

    The registry is then passed to the renderer.
    """

    SKIP_FILE = "@internal"

    def __init__(self, base=None, renderer=None, registry=None, env=None, include_empty=False, ignore=None):
        """
        base: base path containing Salt states to autodocument
        renderer: renderer to use for rendering
        """
        if base is None:
            base = list(base_dir.glob("*/init.sls"))[0].parent
        self.base = base
        self.renderer = renderer or SingleRenderer()
        self.env = env or Environment()
        self.registry = registry or StateRegistry()
        self.include_empty = include_empty
        self.ignore = ignore or []

    def run(self, src=""):
        """
        1. look for .sls files
        2. lex them with Jinja and see if the first token is a comment
        3. in that case, save it to the registry
        4. pass the registry to the renderer
        """
        for file in self.glob(src):
            if str(file.relative_to(self.base)).startswith(tuple(self.ignore)):
                continue
            with open(file, "r") as f:
                text = "".join(f.readlines())
            try:
                comment = self.find_comment(text)
            except TemplateSyntaxError as e:
                raise RuntimeError(
                    f"There was an error parsing the Jinja syntax in file '{file}':\n\n{e}"
                )

            if (comment or self.include_empty) and not comment.startswith(self.SKIP_FILE):
                path = file.relative_to(self.base)
                self.registry.register(State(path, comment, self.base.name))

        self.renderer.render(self.registry)

    def glob(self, src):
        """
        Returns Path objects pointing to Salt states provided by this formula.
        """
        for path in (self.base / Path(src)).rglob("*.sls"):
            yield path

    def find_comment(self, text):
        """
        Lexes files with Jinja. If the first token is a comment, return it.
        """
        lexed = list(self.env.lex(text))
        try:
            i = 0
            t = ""
            while i < 5 and t != "comment":
                _, t, content = lexed[i]
                i += 1
        except IndexError:
            t = content = ""

        if t == "comment":
            return self.parse(content)

        return ""

    def parse(self, comment):
        """
        Strips indentation from [multiline] strings. Could have used
        docutils.prepare_docstring.
        """
        # return d.prepare_docstring(comment)
        ws = min([len(x) - len(x.lstrip()) for x in comment.splitlines() if len(x) > 0])
        return "\n".join([x[ws:] for x in comment.splitlines()]).strip("\n")


if __name__ == "__main__":
    options = toml.loads((base_dir / "pyproject.toml").read_text()).get("tool", {}).get("slsdoc", {})
    renderer = options.get("renderer", "single")

    if renderer == "single":
        rend = SingleRenderer(**options.get("single_renderer", {}))
    elif renderer == "depth":
        rend = DepthRenderer(**options.get("depth_renderer", {}))
    else:
        raise ValueError(f"Invalid value for renderer: {renderer}")

    registry = StateRegistry(avoid_init_dup=options.get("avoid_init_dup", True), sort_list=options.get("sort_list"))
    doc = SaltAutodoc(renderer=rend, registry=registry, include_empty=options.get("include_empty", False), ignore=options.get("ignore", ["_mapdata"]))
    doc.run()