#!/usr/bin/env python3
"""Generate `system.asynchronous_metrics` documentation from C++ source.

Each metric in the `system.asynchronous_metrics` table is registered in C++
with a literal string description living next to the value. For example:

    new_values["QueriesMemoryUsage"] = { queries_memory_usage,
        "Memory used by queries, in bytes." };

This script extracts those (name, description) pairs from the source files
and writes them into a Markdown file between
`<!--AUTOGENERATED_METRICS_START-->` and `<!--AUTOGENERATED_METRICS_END-->`
markers. The descriptions in C++ are the single source of truth — to update
the documentation, edit the description in the corresponding source file and
re-run this script.

Usage:
    python3 utils/generate-async-metrics-docs

It rewrites `docs/en/operations/system-tables/asynchronous_metrics.md`
in place.
"""

import os
import re
import sys

SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
REPO_ROOT = os.path.dirname(SCRIPT_DIR)

SOURCE_FILES = [
    "src/Common/AsynchronousMetrics.cpp",
    "src/Interpreters/ServerAsynchronousMetrics.cpp",
    "src/Coordination/KeeperAsynchronousMetrics.cpp",
]

DOC_FILE = "docs/en/operations/system-tables/asynchronous_metrics.md"

AUTOGEN_START = "<!--AUTOGENERATED_METRICS_START-->"
AUTOGEN_END = "<!--AUTOGENERATED_METRICS_END-->"


def parse_macros(text):
    """Find `#define NAME ... \\` blocks where the body is a multi-line string
    literal and return a `{name: joined_string}` map.

    Only macros whose body is a concatenation of string literals are captured,
    which is exactly the pattern used for description-extending macros like
    `BLOCK_DEVICE_EXPLANATION` and `LOAD_AVERAGE_DOCUMENTATION`.
    """
    macros = {}
    pattern = re.compile(
        r"^#define\s+(\w+)\s*\\\n((?:.*\\\n)*.*?)(?=\n)",
        re.MULTILINE,
    )
    for m in pattern.finditer(text):
        name = m.group(1)
        body = m.group(2)
        body = body.replace("\\\n", "\n")
        joined = _join_string_literals(body)
        if joined is not None:
            macros[name] = joined
    return macros


def _join_string_literals(text):
    """Join adjacent C++ string literals into one Python string.

    Returns None if the text contains anything other than whitespace and
    double-quoted string literals.
    """
    pieces = []
    pos = 0
    while pos < len(text):
        c = text[pos]
        if c.isspace():
            pos += 1
            continue
        if c == '"':
            end = pos + 1
            while end < len(text):
                if text[end] == "\\":
                    end += 2
                    continue
                if text[end] == '"':
                    break
                end += 1
            if end >= len(text):
                return None
            literal = text[pos + 1 : end]
            pieces.append(_unescape_c_string(literal))
            pos = end + 1
            continue
        return None
    return "".join(pieces)


def _unescape_c_string(s):
    """Translate the C escape sequences we expect in metric descriptions."""
    out = []
    i = 0
    while i < len(s):
        if s[i] == "\\" and i + 1 < len(s):
            nxt = s[i + 1]
            if nxt == "n":
                out.append("\n")
            elif nxt == "t":
                out.append("\t")
            elif nxt == "\\":
                out.append("\\")
            elif nxt == '"':
                out.append('"')
            elif nxt == "'":
                out.append("'")
            else:
                out.append(s[i + 1])
            i += 2
        else:
            out.append(s[i])
            i += 1
    return "".join(out)


def expand_macros(text, macros):
    """Replace whole-word occurrences of each macro with its string literal form.

    The replacement is the literal text re-quoted as a C string literal so
    that the surrounding adjacent-literal-concatenation continues to work.
    """
    for name, value in macros.items():
        escaped = value.replace("\\", "\\\\").replace('"', '\\"')
        text = re.sub(r"\b" + re.escape(name) + r"\b", '"' + escaped + '"', text)
    return text


def parse_key_expression(expr, locals_map=None):
    """Convert a C++ expression that builds a metric name into a doc name.

    Returns a string with `*var*` placeholders for variable parts, or None if
    the expression is too dynamic to extract statically.
    """
    expr = expr.strip()

    if expr.startswith("fmt::format("):
        return _parse_fmt_format(expr, locals_map)

    parts = _split_plus(expr)
    if parts is None:
        return None
    out = []
    for part in parts:
        part = part.strip()
        if part.startswith('"') and part.endswith('"'):
            out.append(_unescape_c_string(part[1:-1]))
        elif re.match(r"^[A-Za-z_][A-Za-z0-9_]*$", part):
            if locals_map and part in locals_map:
                resolved = locals_map[part]
                if resolved is None:
                    return None
                out.append(resolved)
            else:
                out.append("*" + part + "*")
        else:
            return None
    return "".join(out)


def _split_plus(expr):
    """Split a `+` expression into top-level parts, ignoring `+` inside strings."""
    parts = []
    depth = 0
    current = []
    i = 0
    in_string = False
    while i < len(expr):
        c = expr[i]
        if in_string:
            if c == "\\" and i + 1 < len(expr):
                current.append(c)
                current.append(expr[i + 1])
                i += 2
                continue
            if c == '"':
                in_string = False
            current.append(c)
            i += 1
            continue
        if c == '"':
            in_string = True
            current.append(c)
            i += 1
            continue
        if c in "([":
            depth += 1
            current.append(c)
            i += 1
            continue
        if c in ")]":
            depth -= 1
            if depth < 0:
                return None
            current.append(c)
            i += 1
            continue
        if c == "+" and depth == 0:
            parts.append("".join(current))
            current = []
            i += 1
            continue
        current.append(c)
        i += 1
    parts.append("".join(current))
    return parts


def _parse_fmt_format(expr, locals_map=None):
    """Parse `fmt::format("template", arg1, arg2, ...)`."""
    inner = expr[len("fmt::format(") :].rstrip()
    if not inner.endswith(")"):
        return None
    inner = inner[:-1]
    args = _split_top_commas(inner)
    if args is None or not args:
        return None
    template = args[0].strip()
    if not (template.startswith('"') and template.endswith('"')):
        return None
    template = _unescape_c_string(template[1:-1])

    placeholders = []
    for raw_arg in args[1:]:
        arg = raw_arg.strip()
        if re.match(r"^[A-Za-z_][A-Za-z0-9_]*$", arg):
            if locals_map and arg in locals_map and locals_map[arg]:
                placeholders.append(locals_map[arg])
            else:
                placeholders.append("*" + arg + "*")
        elif re.match(r"^[A-Za-z_][A-Za-z0-9_.]*$", arg):
            placeholders.append("*" + arg.replace(".", "_") + "*")
        else:
            return None

    placeholder_iter = iter(placeholders)
    out = []
    i = 0
    while i < len(template):
        if template[i : i + 2] == "{}":
            out.append(next(placeholder_iter, "*?*"))
            i += 2
        elif template[i] == "{":
            end = template.find("}", i)
            if end < 0:
                return None
            out.append(next(placeholder_iter, "*?*"))
            i = end + 1
        else:
            out.append(template[i])
            i += 1
    return "".join(out)


def _split_top_commas(expr):
    parts = []
    depth = 0
    current = []
    i = 0
    in_string = False
    while i < len(expr):
        c = expr[i]
        if in_string:
            if c == "\\" and i + 1 < len(expr):
                current.append(c)
                current.append(expr[i + 1])
                i += 2
                continue
            if c == '"':
                in_string = False
            current.append(c)
            i += 1
            continue
        if c == '"':
            in_string = True
            current.append(c)
            i += 1
            continue
        if c in "([{":
            depth += 1
            current.append(c)
            i += 1
            continue
        if c in ")]}":
            depth -= 1
            if depth < 0:
                return None
            current.append(c)
            i += 1
            continue
        if c == "," and depth == 0:
            parts.append("".join(current))
            current = []
            i += 1
            continue
        current.append(c)
        i += 1
    parts.append("".join(current))
    return parts


def parse_description_expression(expr):
    """Parse a description argument as concatenated string literals.

    After macro expansion every description is a sequence of adjacent C string
    literals separated only by whitespace. Returns the joined Python string or
    None if the expression contains anything else.
    """
    return _join_string_literals(expr)


# Metric-registration call shapes we recognise. The key and description
# expressions span multiple lines so we match non-greedy up to the closing
# punctuation and then balance brackets in Python.
ASSIGN_PATTERN = re.compile(
    r"(?P<container>new_values|values)\[(?P<key>.+?)\]\s*=\s*"
    r"(?:\{(?P<brace>(?:[^{}]|\{[^{}]*\})*)\}"
    r"|AsynchronousMetricValue\((?P<paren>(?:[^()]|\([^()]*\))*)\))",
    re.DOTALL,
)


JEMALLOC_DESCRIPTION = (
    "An internal metric of the low-level memory allocator (jemalloc)."
    " See https://jemalloc.net/jemalloc.3.html"
)

JEMALLOC_HELPERS = {
    "saveJemallocMetric": "jemalloc.{}",
    "saveAllArenasMetric": "jemalloc.arenas.all.{}",
    "saveJemallocProf": "jemalloc.prof.{}",
}

JEMALLOC_HELPER_PATTERN = re.compile(
    r"(?P<helper>saveJemallocMetric|saveAllArenasMetric|saveJemallocProf)"
    r"\s*<[^>]+>\s*\([^,]+,\s*\"(?P<name>[^\"]+)\"\s*\)"
)

# The MIB-name argument (second arg) is usually a string literal but is sometimes
# a `fmt::format(...)` call when the arena index is variable. The metric name
# (third arg) is always a string literal, so we match the body as a sequence of
# non-paren chars and optional parenthesized blocks (covering `fmt::format(...)`)
# up to the final `,` and then capture the trailing string literal.
JEMALLOC_IMPL_PATTERN = re.compile(
    r"saveJemallocMetricImpl\s*<[^>]+>\s*\("
    r"[^()]*(?:\([^()]*\)[^()]*)*"
    r",\s*\"(?P<name>[^\"]+)\"\s*\)",
    re.DOTALL,
)

# Map-of-pairs literal used to attach a metric name + description to a
# configuration-key string (e.g. the protocol thread/rejected-connection maps in
# `ServerAsynchronousMetrics::updateImpl`). The outer key is irrelevant for the
# docs; we extract the inner (MetricName, Description) pair.
NAME_DOC_PAIR_PATTERN = re.compile(
    r'\{\s*"[^"]+"\s*,\s*\{\s*"([^"]+)"\s*,\s*((?:"(?:[^"\\]|\\.)*"\s*)+)\}\s*\}',
    re.DOTALL,
)


# Capture `String varname = ...;` so we can resolve metric keys that are
# computed into a local variable before being used as the assignment key.
LOCAL_STRING_PATTERN = re.compile(
    r"\b(?:String|std::string|const\s+std::string)\s+(\w+)\s*=\s*(.+?);",
    re.DOTALL,
)


def collect_local_strings(text):
    """Build a map from local string-variable name to its parsed value.

    Values are stored in the same form as parse_key_expression output (with
    `*var*` placeholders for any sub-variables we cannot resolve). Variables
    we cannot parse map to `None` so callers know to skip them.
    """
    locals_map = {}
    for m in LOCAL_STRING_PATTERN.finditer(text):
        name = m.group(1)
        rhs = m.group(2).strip()
        locals_map[name] = parse_key_expression(rhs, locals_map)
    return locals_map


def extract_metrics_from_source(text, macros, source_path):
    """Return a list of (metric_name, description, source_path) tuples."""
    text = expand_macros(text, macros)
    locals_map = collect_local_strings(text)

    results = []

    for m in JEMALLOC_HELPER_PATTERN.finditer(text):
        helper = m.group("helper")
        suffix = m.group("name")
        full_name = JEMALLOC_HELPERS[helper].format(suffix)
        results.append((full_name, JEMALLOC_DESCRIPTION, source_path))

    for m in JEMALLOC_IMPL_PATTERN.finditer(text):
        results.append((m.group("name"), JEMALLOC_DESCRIPTION, source_path))

    for m in NAME_DOC_PAIR_PATTERN.finditer(text):
        name = m.group(1)
        desc = parse_description_expression(m.group(2))
        if desc is None:
            continue
        desc = desc.strip()
        if not desc:
            continue
        results.append((name, desc, source_path))

    for m in ASSIGN_PATTERN.finditer(text):
        key_expr = m.group("key")
        body = m.group("brace") if m.group("brace") is not None else m.group("paren")
        body_args = _split_top_commas(body)
        if body_args is None or len(body_args) < 2:
            continue
        # description is the last argument (value comes first).
        desc_expr = body_args[-1]
        name = parse_key_expression(key_expr, locals_map)
        if name is None:
            continue
        # Skip purely-placeholder keys: these come from generic helper functions
        # (e.g. saveJemallocMetricImpl) where the actual metric names are
        # produced at call sites and captured elsewhere.
        if re.fullmatch(r"\*\w+\*", name):
            continue
        desc = parse_description_expression(desc_expr)
        if desc is None:
            continue
        desc = desc.strip()
        if not desc:
            continue
        results.append((name, desc, source_path))
    return results


def slugify(name):
    """Anchor id matching the existing docs convention."""
    return re.sub(r"[^a-z0-9_]+", "", name.lower())


def render_markdown(metrics):
    """Render a sorted list of `### Name {#anchor}` / description entries."""
    sort_key = lambda item: (item[0].lower(), item[0])
    metrics = sorted(set((n, d) for n, d, _ in metrics), key=sort_key)
    lines = []
    for name, desc in metrics:
        anchor = slugify(name)
        lines.append(f"### {name} {{#{anchor}}}")
        lines.append("")
        lines.append(desc.strip())
        lines.append("")
    return "\n".join(lines)


def update_doc_file(filepath, generated_block):
    with open(filepath, "r") as f:
        content = f.read()

    if AUTOGEN_START not in content or AUTOGEN_END not in content:
        raise SystemExit(
            f"{filepath}: missing {AUTOGEN_START} / {AUTOGEN_END} markers"
        )

    pattern = re.compile(
        re.escape(AUTOGEN_START) + r".*?" + re.escape(AUTOGEN_END),
        re.DOTALL,
    )
    replacement = AUTOGEN_START + "\n\n" + generated_block + "\n" + AUTOGEN_END
    new_content = pattern.sub(replacement, content, count=1)

    if new_content != content:
        with open(filepath, "w") as f:
            f.write(new_content)
        return True
    return False


def main():
    metrics = []
    duplicates = {}
    for relpath in SOURCE_FILES:
        path = os.path.join(REPO_ROOT, relpath)
        with open(path, "r") as f:
            text = f.read()
        macros = parse_macros(text)
        extracted = extract_metrics_from_source(text, macros, relpath)
        for name, desc, _ in extracted:
            existing = duplicates.get(name)
            if existing is not None and existing != desc:
                print(
                    f"warning: metric '{name}' has different descriptions in code",
                    file=sys.stderr,
                )
            duplicates[name] = desc
        metrics.extend(extracted)

    print(f"Extracted {len(set(n for n, _, _ in metrics))} unique metrics", file=sys.stderr)
    block = render_markdown(metrics)

    doc_path = os.path.join(REPO_ROOT, DOC_FILE)
    if update_doc_file(doc_path, block):
        print(f"Updated {DOC_FILE}", file=sys.stderr)
    else:
        print(f"{DOC_FILE} already up to date", file=sys.stderr)


if __name__ == "__main__":
    main()
