#!/usr/bin/env python3
"""Generate system-table documentation from a ClickHouse binary.

Queries system.documentation via clickhouse-local, then generates or updates
Markdown documentation files. Structured table metadata comments provide the
description, examples and related material; system.documentation combines those
parts with the live column metadata into a consistently structured page.
Asynchronous metric descriptions are extracted directly from their C++
registrations because clickhouse-local has no live metrics instance; the
surrounding `system.asynchronous_metrics` page itself is owned by the source
too, so `system.documentation` exposes it even though the table is attached
only by the server.

Usage:
    python3 utils/generate-system-tables-docs --binary ./build/programs/clickhouse --docs-dir docs/reference/system-tables/
"""

import argparse
import csv
import importlib.util
import io
import os
import re
import subprocess
import sys
from importlib.machinery import SourceFileLoader

csv.field_size_limit(sys.maxsize)

SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
DEFAULT_CONFIG = os.path.join(SCRIPT_DIR, "generate-system-tables-docs-config.xml")

AUTOGENERATED_START = "{/*AUTOGENERATED_START*/}"
AUTOGENERATED_END = "{/*AUTOGENERATED_END*/}"
FRONTMATTER_RE = re.compile(r"\A---\n.*?\n---\n*", re.DOTALL)
GENERATED_SECTION_HEADING_RE = re.compile(
    r"^## (?P<title>Description|Columns|Examples?|See [Aa]lso) \{#[^}\n]+\}[ \t]*$",
    re.MULTILINE,
)

# A few historical pages are published under both a canonical filename and a
# legacy duplicate. Generate every published copy until navigation redirects
# or removes the legacy route, so none of them retain stale content.
DOC_FILENAMES = {
    "delta_lake_metadata_log": (
        "delta_metadata_log.mdx",
        "delta_lake_metadata_log.mdx",
    ),
    "warnings": ("system_warnings.mdx", "warnings.mdx"),
}


def documentation_filenames(table_name):
    return DOC_FILENAMES.get(table_name, (f"{table_name}.mdx",))


def load_async_metrics_generator():
    generator_path = os.path.join(SCRIPT_DIR, "generate-async-metrics-docs")
    loader = SourceFileLoader("generate_async_metrics_docs", generator_path)
    spec = importlib.util.spec_from_file_location(
        "generate_async_metrics_docs", generator_path, loader=loader
    )
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module


def replace_section_catalog(documentation_md, section_heading, generated_catalog):
    """Replace generated H3 entries in an H2 section with a fresh catalog."""
    marker = f"## {section_heading}"
    marker_position = documentation_md.find(marker)
    if marker_position == -1:
        raise ValueError(f"Generated documentation has no {marker} section")

    heading_end = documentation_md.find("\n", marker_position)
    if heading_end == -1:
        raise ValueError(f"Generated documentation has no body after {marker}")
    section_end = documentation_md.find("\n## ", heading_end)
    if section_end == -1:
        section_end = len(documentation_md)

    section = documentation_md[heading_end + 1 : section_end].strip()
    first_entry = re.search(r"(?m)^### ", section)
    introduction = section[: first_entry.start()].rstrip() if first_entry else section
    replacement = "\n\n" + "\n\n".join(
        part for part in (introduction, generated_catalog.strip()) if part
    )
    if section_end < len(documentation_md):
        replacement += "\n"

    return documentation_md[: heading_end + 1] + replacement + documentation_md[section_end:]


def populate_asynchronous_metrics_catalog(full_documentation):
    """Populate the async-metric page from C++ source, independent of runtime state."""
    table_name = "asynchronous_metrics"
    if table_name not in full_documentation:
        raise ValueError(
            "system.documentation has no asynchronous_metrics page. The page is owned by the source"
            " (`ASYNCHRONOUS_METRICS_DOCUMENTATION` in src/Storages/System/StorageSystemAsynchronousMetrics.cpp)"
            " and StorageSystemDocumentation emits it even where the table is not attached, so the binary"
            " is either too old or was built without that row"
        )

    generated_catalog, metric_count = load_async_metrics_generator().generate_markdown()
    if not generated_catalog:
        raise ValueError("No asynchronous metric documentation was extracted from C++ source")
    full_documentation[table_name] = replace_section_catalog(
        full_documentation[table_name],
        "Metric descriptions {#metric-descriptions}",
        generated_catalog,
    )
    return metric_count


def unescape_tsv(value):
    """Unescape TSV escape sequences (\\n, \\t, \\', \\\\).

    Order matters: backslash must be replaced last to avoid
    double-unescaping (e.g., \\\\n -> \\n -> newline).
    """
    result = []
    i = 0
    while i < len(value):
        if value[i] == "\\" and i + 1 < len(value):
            next_char = value[i + 1]
            if next_char == "n":
                result.append("\n")
                i += 2
                continue
            elif next_char == "t":
                result.append("\t")
                i += 2
                continue
            elif next_char == "\\":
                result.append("\\")
                i += 2
                continue
            elif next_char == "'":
                result.append("'")
                i += 2
                continue
        result.append(value[i])
        i += 1
    return "".join(result)


def extract_page_summary(documentation_md):
    """Extract the first paragraph of the generated Description section."""
    marker = "## Description {#description}"
    marker_position = documentation_md.find(marker)
    if marker_position == -1:
        raise ValueError("Generated system-table documentation has no Description section")

    description = documentation_md[marker_position + len(marker):].lstrip("\n")
    section_end = description.find("\n## ")
    if section_end != -1:
        description = description[:section_end]
    first_paragraph = description.split("\n\n", 1)[0]
    summary = " ".join(line.strip() for line in first_paragraph.splitlines()).strip()
    if not summary:
        raise ValueError("Generated system-table documentation has an empty Description section")
    return summary


def run_query(binary, config, query, use_client=False, extra_args=None):
    """Run a query via clickhouse-local (or clickhouse-client) and return TSV rows."""
    if use_client:
        cmd = [binary, "client", "--query", query, "--format", "TSVWithNamesAndTypes"]
    else:
        cmd = [binary, "local", "--query", query, "--format", "TSVWithNamesAndTypes"]
    if config and not use_client:
        cmd.extend(["--config-file", config])
    if extra_args:
        cmd.extend(extra_args)
    stdout = ""
    stderr = ""
    returncode = 0
    timed_out = False

    def _as_text(value):
        if value is None:
            return ""
        return value if isinstance(value, str) else value.decode("utf-8", errors="replace")

    try:
        result = subprocess.run(cmd, capture_output=True, text=True, timeout=300)
        stdout = result.stdout
        stderr = result.stderr
        returncode = result.returncode
    except subprocess.TimeoutExpired as e:
        stdout = _as_text(e.stdout)
        stderr = _as_text(e.stderr)
        returncode = -1
        timed_out = True
        print(f"Error: query timed out after 300 seconds: {query}", file=sys.stderr)
    rows = list(csv.reader(io.StringIO(stdout), delimiter="\t"))
    # `clickhouse-local` can exit with a non-zero status because of races
    # between system log teardown and `Poco::Application` destruction, even
    # after the query has produced complete output to stdout. Tolerate that
    # specific case only: when stdout parses as TSVWithNamesAndTypes (header +
    # types row), retry once and accept only a clean second result. This avoids
    # mistaking output truncated by a real exception for a complete result.
    # Timeouts always propagate.
    if returncode != 0 and not timed_out and len(rows) >= 2:
        print(
            f"Warning: query exited with code {returncode} after producing"
            " parseable output; retrying once.",
            file=sys.stderr,
        )
        if stderr:
            print(f"stderr: {stderr}", file=sys.stderr)
        try:
            result = subprocess.run(cmd, capture_output=True, text=True, timeout=300)
            stdout = result.stdout
            stderr = result.stderr
            returncode = result.returncode
            rows = list(csv.reader(io.StringIO(stdout), delimiter="\t"))
        except subprocess.TimeoutExpired as e:
            stderr = _as_text(e.stderr)
            returncode = -1
            timed_out = True
            print(f"Error: query retry timed out after 300 seconds: {query}", file=sys.stderr)

    if returncode != 0:
        print(f"Error running query: {query}", file=sys.stderr)
        print(f"stderr: {stderr}", file=sys.stderr)
        sys.exit(1)
    if len(rows) < 2:
        return []
    headers = rows[0]
    # rows[1] is the types row, skip it
    return [
        {k: unescape_tsv(v) for k, v in zip(headers, row)}
        for row in rows[2:]
    ]


def preserve_published_section_headings(existing_content, documentation_md):
    """Preserve published headings and anchors for generated top-level sections."""

    def section_key(title):
        key = title.casefold()
        return "examples" if key in ("example", "examples") else key

    existing_headings = {
        section_key(match.group("title")): match.group(0).rstrip()
        for match in GENERATED_SECTION_HEADING_RE.finditer(existing_content)
    }

    def preserve_heading(match):
        return existing_headings.get(
            section_key(match.group("title")), match.group(0).rstrip()
        )

    return GENERATED_SECTION_HEADING_RE.sub(preserve_heading, documentation_md)


def update_full_doc_file(filepath, documentation_md):
    """Replace a page body with its structured embedded documentation.

    Frontmatter and an initial MDX import/component preamble remain maintained
    in the docs tree. Everything after that preamble is a single generated
    region so descriptions, columns and examples cannot drift from the
    documentation embedded in the binary.
    """
    with open(filepath, "r") as f:
        content = f.read()

    documentation_md = preserve_published_section_headings(content, documentation_md)

    frontmatter = FRONTMATTER_RE.match(content)
    prefix = frontmatter.group(0).rstrip("\n") if frontmatter else ""
    body = content[frontmatter.end():] if frontmatter else content

    # Some system-table pages import and render an MDX-only cloud availability
    # notice. It cannot be embedded in C++, but it must remain ahead of the
    # generated Markdown body. Only preserve the well-defined initial preamble;
    # prose and headings are owned by the structured C++ documentation.
    preamble_lines = []
    for line in body.lstrip("\n").splitlines():
        stripped = line.strip()
        if not stripped:
            if preamble_lines:
                preamble_lines.append(line)
            continue
        if line.startswith("import ") or re.fullmatch(r"<[A-Z][^>]*?/>", stripped):
            preamble_lines.append(line)
            continue
        break
    preamble = "\n".join(preamble_lines).strip("\n")

    documentation_md = documentation_md.strip("\n")
    generated = (
        AUTOGENERATED_START
        + "\n"
        + documentation_md
        + "\n"
        + AUTOGENERATED_END
        + "\n"
    )
    parts = [part for part in (prefix, preamble, generated.rstrip("\n")) if part]
    new_content = "\n\n".join(parts) + "\n"

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


def create_doc_file(filepath, table_name, table_comment, documentation_md):
    """Create a new documentation file for a system table."""
    # Escape single quotes in the description for YAML.
    escaped_comment = table_comment.replace("'", "''")

    body = (
        f"{AUTOGENERATED_START}\n"
        f"{documentation_md.strip()}\n"
        f"{AUTOGENERATED_END}\n"
    )

    content = f"""---
description: '{escaped_comment}'
keywords: ['system table', '{table_name}']
sidebarTitle: 'system.{table_name}'
slug: /operations/system-tables/{table_name}
title: 'system.{table_name}'
doc_type: 'reference'
---

{body}"""
    with open(filepath, "w") as f:
        f.write(content)


def generate_pages(docs_dir, full_documentation):
    """Generate every system-table page exposed by system.documentation."""
    updated = 0
    created = 0
    skipped = 0

    for table_name in sorted(full_documentation):
        if re.search(r"_\d+$", table_name):
            continue

        documentation_md = full_documentation[table_name]
        comment = extract_page_summary(documentation_md)

        for doc_filename in documentation_filenames(table_name):
            doc_filepath = os.path.join(docs_dir, doc_filename)

            if os.path.isfile(doc_filepath):
                changed = update_full_doc_file(doc_filepath, documentation_md)
                if changed:
                    updated += 1
                    print(f"  Updated: {doc_filename}")
                else:
                    skipped += 1
            else:
                create_doc_file(
                    doc_filepath,
                    table_name,
                    comment,
                    documentation_md,
                )
                created += 1
                print(f"  Created: {doc_filename}")

    return updated, created, skipped


def main():
    parser = argparse.ArgumentParser(
        description="Generate system tables documentation from a ClickHouse binary."
    )
    parser.add_argument(
        "--binary",
        required=True,
        help="Path to the clickhouse binary.",
    )
    parser.add_argument(
        "--docs-dir",
        required=True,
        help="Path to docs/reference/system-tables/ directory.",
    )
    parser.add_argument(
        "--config",
        default=DEFAULT_CONFIG,
        help="Path to the config file for clickhouse-local (optional).",
    )
    parser.add_argument(
        "--use-client",
        action="store_true",
        help="Use clickhouse-client instead of clickhouse-local (requires a running server).",
    )
    parser.add_argument(
        "--port",
        default=None,
        help="Port for clickhouse-client connection (default: 9000).",
    )
    args = parser.parse_args()

    binary = os.path.abspath(args.binary)
    docs_dir = os.path.abspath(args.docs_dir)
    config = os.path.abspath(args.config) if args.config else None
    use_client = args.use_client
    extra_args = []
    if args.port:
        extra_args.extend(["--port", args.port])

    if not os.path.isfile(binary):
        print(f"Error: binary not found: {binary}", file=sys.stderr)
        sys.exit(1)

    if not os.path.isdir(docs_dir):
        print(f"Error: docs directory not found: {docs_dir}", file=sys.stderr)
        sys.exit(1)

    # Each row is already assembled from the structured table comment and
    # the live schema into the complete generated Markdown body.
    print("Querying structured system table documentation...")
    documentation = run_query(
        binary,
        config,
        "SELECT name, description FROM system.documentation"
        " WHERE type = 'System Table'"
        " ORDER BY name",
        use_client,
        extra_args,
    )
    full_documentation = {row["name"]: row["description"] for row in documentation}
    metric_count = populate_asynchronous_metrics_catalog(full_documentation)
    print(f"Loaded {metric_count} asynchronous metric descriptions from C++ source.")

    updated, created, skipped = generate_pages(docs_dir, full_documentation)

    print(f"\nDone: {updated} updated, {created} created, {skipped} unchanged.")


if __name__ == "__main__":
    main()
