"""Generate code reference pages for Zensical (no mkdocs-gen-files).

Writes generated Markdown pages to:
  docs/src/api/...

Also generates a TOML nav snippet (for Zensical) at:
  docs/nav.generated.toml
"""

from __future__ import annotations

import shutil
from pathlib import Path
from typing import Any

PKG_NAME = "afnio"

REPO_ROOT = Path(__file__).resolve().parent.parent.parent
SRC = REPO_ROOT / PKG_NAME

DOCS_DIR = REPO_ROOT / "docs"
DOCS_SRC_DIR = DOCS_DIR / "src"
OUT_ROOT = DOCS_SRC_DIR / "api"  # where generated pages go
NAV_TOML_PATH = DOCS_DIR / "nav.generated.toml"

# Exact module names (dotted import path) to skip
SKIP_MODULES: set[str] = {
    "afnio.autodiff.config",
    "afnio.autodiff.utils",
    "afnio.models.model_registry",
}


Node = dict[str, Any]


class Nav:
    """Tiny replacement for mkdocs_gen_files.Nav, but outputs Zensical TOML nav."""

    def __init__(self) -> None:
        self._items: dict[tuple[str, ...], str] = {}

    def __setitem__(self, parts: tuple[str, ...], doc_path: str) -> None:
        self._items[parts] = doc_path

    def build_toml_nav(self) -> list[dict[str, Any]]:
        """
        Build a nested `nav = [...]` structure compatible with Zensical.

        Leaf nodes:
          { "name" = "path.md" }

        Section nodes (when an index exists for the section):
          { "name" = [ "path/to/index.md", { "child" = ... }, ... ] }

        Section nodes (without index):
          { "name" = [ { "child" = ... }, ... ] }
        """
        root: Node = {"__children__": {}}

        def ensure_child(parent: Node, key: str) -> Node:
            children: dict[str, Node] = parent["__children__"]
            if key not in children:
                children[key] = {"__children__": {}}
            return children[key]

        # Build a tree that can hold both children and an optional index page
        for parts, rel_path in sorted(self._items.items(), key=lambda x: x[0]):
            cursor = root
            for i, part in enumerate(parts):
                is_leaf = i == len(parts) - 1
                if is_leaf:
                    # If this entry is an index.md, store it as the section index
                    if rel_path.endswith("/index.md") or rel_path.endswith(
                        "\\index.md"
                    ):
                        node = ensure_child(cursor, part)
                        node["__index__"] = rel_path
                    else:
                        node = ensure_child(cursor, part)
                        node["__leaf__"] = rel_path
                else:
                    cursor = ensure_child(cursor, part)

        def emit_node(name: str, node: Node) -> dict[str, Any]:
            children: dict[str, Node] = node.get("__children__", {})
            index = node.get("__index__")
            leaf = node.get("__leaf__")

            # If it has no children, it must be a leaf
            if not children:
                if not leaf and index:
                    # e.g. "afnio" pointing to its index only
                    return {name: index}
                if not leaf:
                    raise RuntimeError(f"Missing leaf for {name}")
                return {name: leaf}

            # Section: list of children; prepend index if present
            items: list[Any] = []
            if index:
                items.append(index)

            for child_name in sorted(children.keys()):
                items.append(emit_node(child_name, children[child_name]))

            return {name: items}

        out: list[dict[str, Any]] = []
        for top in sorted(root["__children__"].keys()):
            out.append(emit_node(top, root["__children__"][top]))
        return out

    def dumps_toml(self) -> str:
        nav_list = self.build_toml_nav()
        return "nav = " + _toml_value(nav_list) + "\n"


def _toml_string(s: str) -> str:
    # Minimal TOML string escaping
    return '"' + s.replace("\\", "\\\\").replace('"', '\\"') + '"'


def _toml_value(v: Any) -> str:
    if isinstance(v, str):
        return _toml_string(v)
    if isinstance(v, bool):
        return "true" if v else "false"
    if isinstance(v, (int, float)):
        return str(v)
    if isinstance(v, list):
        inner = ", ".join(_toml_value(x) for x in v)
        return f"[{inner}]"
    if isinstance(v, dict):
        # inline table
        if len(v) != 1:
            raise ValueError("Nav dict entries must have exactly one key")
        ((k, val),) = v.items()
        return "{ " + f"{_toml_string(k)} = {_toml_value(val)}" + " }"
    raise TypeError(f"Unsupported TOML type: {type(v)!r}")


def main() -> None:
    if not SRC.exists():
        raise FileNotFoundError(f"Source package directory not found: {SRC}")

    # Clean output (prevents stale pages)
    if OUT_ROOT.exists():
        shutil.rmtree(OUT_ROOT)
    OUT_ROOT.mkdir(parents=True, exist_ok=True)

    nav = Nav()

    for path in sorted(SRC.rglob("*.py")):
        module_path = path.relative_to(SRC).with_suffix("")
        doc_path = path.relative_to(SRC).with_suffix(".md")
        if not doc_path.parts or doc_path.parts[0] != PKG_NAME:
            doc_path = Path(PKG_NAME) / doc_path
        full_doc_path = OUT_ROOT / doc_path

        parts = tuple(module_path.parts)

        if parts[-1] == "__init__":
            parts = parts[:-1]
            if not parts:
                # top-level package __init__.py -> api/afnio/index.md
                parts = (SRC.name,)
                doc_path = Path(SRC.name) / "index.md"
                full_doc_path = OUT_ROOT / doc_path
            else:
                # package __init__.py -> .../index.md
                doc_path = doc_path.with_name("index.md")
                full_doc_path = full_doc_path.with_name("index.md")
        elif parts[-1] == "__main__":
            continue

        # skip private/internal modules that start with '_'
        if any(p.startswith("_") for p in parts):
            continue

        # nav path should be relative to OUT_ROOT
        href = (Path("api") / doc_path).as_posix()
        nav_parts = parts if parts == (SRC.name,) else (SRC.name,) + parts

        # build full import identifier prefixed with package name
        if parts and parts[0] != SRC.name:
            identifier = ".".join((SRC.name,) + parts)
        else:
            identifier = ".".join(parts)

        # Skip explicit modules listed in SKIP_MODULES (don't add nav entry)
        if identifier in SKIP_MODULES:
            print(f"Skipping explicitly-listed module {identifier}")
            continue

        full_doc_path.parent.mkdir(parents=True, exist_ok=True)
        full_doc_path.write_text(
            f"# `{identifier}`\n\n::: {identifier}\n",
            encoding="utf-8",
        )

        # only add nav entry for non-skipped modules
        nav[nav_parts] = href

        print(f"Generated {full_doc_path.relative_to(DOCS_SRC_DIR)}")

    NAV_TOML_PATH.write_text(nav.dumps_toml(), encoding="utf-8")
    print(f"Generated {NAV_TOML_PATH.relative_to(DOCS_DIR)}")


if __name__ == "__main__":
    main()
