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

Queries system.tables and system.columns via clickhouse-local to extract
table descriptions and column metadata, then generates/updates markdown
documentation files.

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

import argparse
import csv
import io
import os
import re
import subprocess
import sys

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-->"

# Mapping from base ClickHouse type names to documentation slugs.
TYPE_LINKS = {
    "String": "/sql-reference/data-types/string",
    "FixedString": "/sql-reference/data-types/fixedstring",
    "UInt8": "/sql-reference/data-types/int-uint",
    "UInt16": "/sql-reference/data-types/int-uint",
    "UInt32": "/sql-reference/data-types/int-uint",
    "UInt64": "/sql-reference/data-types/int-uint",
    "UInt128": "/sql-reference/data-types/int-uint",
    "UInt256": "/sql-reference/data-types/int-uint",
    "Int8": "/sql-reference/data-types/int-uint",
    "Int16": "/sql-reference/data-types/int-uint",
    "Int32": "/sql-reference/data-types/int-uint",
    "Int64": "/sql-reference/data-types/int-uint",
    "Int128": "/sql-reference/data-types/int-uint",
    "Int256": "/sql-reference/data-types/int-uint",
    "Float32": "/sql-reference/data-types/float",
    "Float64": "/sql-reference/data-types/float",
    "Decimal": "/sql-reference/data-types/decimal",
    "Decimal32": "/sql-reference/data-types/decimal",
    "Decimal64": "/sql-reference/data-types/decimal",
    "Decimal128": "/sql-reference/data-types/decimal",
    "Decimal256": "/sql-reference/data-types/decimal",
    "Bool": "/sql-reference/data-types/boolean",
    "Date": "/sql-reference/data-types/date",
    "Date32": "/sql-reference/data-types/date32",
    "DateTime": "/sql-reference/data-types/datetime",
    "DateTime64": "/sql-reference/data-types/datetime64",
    "UUID": "/sql-reference/data-types/uuid",
    "Enum8": "/sql-reference/data-types/enum",
    "Enum16": "/sql-reference/data-types/enum",
    "Array": "/sql-reference/data-types/array",
    "Map": "/sql-reference/data-types/map",
    "Tuple": "/sql-reference/data-types/tuple",
    "Nullable": "/sql-reference/data-types/nullable",
    "LowCardinality": "/sql-reference/data-types/lowcardinality",
    "IPv4": "/sql-reference/data-types/ipv4",
    "IPv6": "/sql-reference/data-types/ipv6",
}

# Tables to skip in documentation generation (internal/helper tables).
SKIP_TABLES = {"transposed_metric_log"}


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 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
    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:
        def _as_text(x):
            if x is None:
                return ""
            return x if isinstance(x, str) else x.decode("utf-8", errors="replace")
        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)
    reader = csv.reader(io.StringIO(stdout), delimiter="\t")
    rows = list(reader)
    # `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: the process must have exited cleanly (no timeout)
    # and stdout must parse as TSVWithNamesAndTypes (header + types row).
    # Timeouts always propagate, since partial output may be silently accepted
    # otherwise.
    if returncode != 0:
        if not timed_out and len(rows) >= 2:
            print(
                f"Warning: query exited with code {returncode} but"
                f" produced {len(rows)} rows of output, treating as success.",
                file=sys.stderr,
            )
            if stderr:
                print(f"stderr: {stderr}", file=sys.stderr)
        else:
            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 format_type_link(type_str):
    """Format a ClickHouse type as a markdown link.

    For simple types like 'String', returns '[String](/sql-reference/data-types/string)'.
    For compound types like 'Array(String)', links the outermost type.
    For unknown types, returns the type without a link.
    """
    # Extract the outermost type name (before any parenthesis).
    match = re.match(r"^(\w+)", type_str)
    if not match:
        return f"`{type_str}`"

    base_type = match.group(1)
    link = TYPE_LINKS.get(base_type)
    if link:
        return f"[{type_str}]({link})"
    return f"`{type_str}`"


def indent_continuation(text, indent="  "):
    """Indent continuation lines in a multi-line string for markdown list items."""
    lines = text.split("\n")
    if len(lines) <= 1:
        return text
    # First line stays as-is, subsequent lines get indented.
    result = [lines[0]]
    for line in lines[1:]:
        if line.strip():
            result.append(indent + line)
        else:
            result.append("")
    return "\n".join(result)


def generate_columns_markdown(columns, aliases=None):
    """Generate markdown for a list of columns, with optional aliases."""
    lines = []
    for col in columns:
        name = col["name"]
        type_str = col["type"]
        comment = col.get("comment", "").strip()

        type_linked = format_type_link(type_str)

        if comment:
            formatted_comment = indent_continuation(comment)
            lines.append(f"- `{name}` ({type_linked}) — {formatted_comment}")
        else:
            lines.append(f"- `{name}` ({type_linked})")

    if aliases:
        lines.append("")
        lines.append("**Aliases:**")
        lines.append("")
        for alias in aliases:
            name = alias["name"]
            target = alias.get("default_expression", "")
            lines.append(f"- `{name}` — Alias for `{target}`.")

    return "\n".join(lines) + "\n" if lines else ""


def update_doc_file(filepath, columns_md):
    """Update a doc file by replacing content between AUTOGENERATED markers."""
    with open(filepath, "r") as f:
        content = f.read()

    if AUTOGENERATED_START in content and AUTOGENERATED_END in content:
        # Replace content between markers.
        pattern = re.compile(
            re.escape(AUTOGENERATED_START) + r".*?" + re.escape(AUTOGENERATED_END),
            re.DOTALL,
        )
        new_content = pattern.sub(
            AUTOGENERATED_START + "\n" + columns_md + AUTOGENERATED_END,
            content,
            count=1,
        )
    else:
        # No markers found. Look for "Columns:" header and add markers.
        new_content = add_markers_to_file(content, columns_md)

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


def add_markers_to_file(content, columns_md):
    """Add AUTOGENERATED markers to a file that doesn't have them.

    Looks for a "Columns:" line and replaces the column list that follows it
    (up to the next major section or end of file) with autogenerated markers.
    """
    # Find "Columns:" or "## Columns" or "## Columns {#columns}" header.
    columns_pattern = re.compile(r"^(#{0,3}\s*Columns:?(?:\s*\{#[\w-]+\})?\s*)$", re.MULTILINE)
    match = columns_pattern.search(content)

    if match:
        after_header = match.end()
        rest = content[after_header:]

        # Find the next major section boundary after the column list.
        # Major sections are: headings (## or **See Also**, **Example**),
        # or specific patterns that indicate end of column documentation.
        section_pattern = re.compile(
            r"^(?:#{1,3}\s|\*\*(?:See Also|Example|See also)\*\*)",
            re.MULTILINE,
        )
        section_match = section_pattern.search(rest)

        if section_match:
            # Keep content from the next section onward.
            after_columns = rest[section_match.start():]
        else:
            after_columns = ""

        # `\s*$` in the columns header pattern can consume the trailing
        # newline, so `before` may or may not end with `\n`. Normalize it
        # before appending the blank line that separates the header from
        # the autogenerated block.
        before = content[:after_header].rstrip("\n")
        new_content = (
            before
            + "\n\n"
            + AUTOGENERATED_START
            + "\n"
            + columns_md
            + AUTOGENERATED_END
            + "\n"
        )
        if after_columns:
            new_content += "\n" + after_columns
    else:
        # No Columns header found; append it.
        new_content = (
            content.rstrip()
            + "\n\nColumns:\n\n"
            + AUTOGENERATED_START
            + "\n"
            + columns_md
            + AUTOGENERATED_END
            + "\n"
        )

    return new_content


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

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

{table_comment}

Columns:

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


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/en/operations/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)

    # Query table metadata.
    # Filter out rotated log tables (e.g., query_log_0, trace_log_1) and
    # user-created tables by requiring a non-empty comment (all real system
    # tables have comments set during registration).
    print("Querying system tables...")
    tables = run_query(
        binary,
        config,
        "SELECT name, comment FROM system.tables"
        " WHERE database = 'system'"
        " AND comment != ''"
        " AND NOT match(name, '_\\\\d+$')"
        " ORDER BY name",
        use_client,
        extra_args,
    )

    # Build set of known table names to filter columns.
    known_tables = {t["name"] for t in tables}

    # Query regular column metadata.
    print("Querying system columns...")
    columns = run_query(
        binary,
        config,
        "SELECT table, name, type, comment, default_kind, default_expression"
        " FROM system.columns WHERE database = 'system'"
        " AND default_kind NOT IN ('ALIAS')"
        " ORDER BY table, position",
        use_client,
        extra_args,
    )
    columns = [c for c in columns if c["table"] in known_tables]

    # Query alias columns separately.
    print("Querying alias columns...")
    aliases = run_query(
        binary,
        config,
        "SELECT table, name, type, default_expression"
        " FROM system.columns WHERE database = 'system'"
        " AND default_kind = 'ALIAS'"
        " ORDER BY table, position",
        use_client,
        extra_args,
    )
    aliases = [a for a in aliases if a["table"] in known_tables]

    # Group columns and aliases by table, preserving order.
    table_columns = {}
    for col in columns:
        table_columns.setdefault(col["table"], []).append(col)

    table_aliases = {}
    for alias in aliases:
        table_aliases.setdefault(alias["table"], []).append(alias)

    # Build table comment map.
    table_comments = {t["name"]: t["comment"] for t in tables}

    updated = 0
    created = 0
    skipped = 0

    for table_name in sorted(table_columns.keys()):
        if table_name in SKIP_TABLES:
            continue

        cols = table_columns[table_name]
        table_als = table_aliases.get(table_name, [])
        columns_md = generate_columns_markdown(cols, table_als)
        comment = table_comments.get(table_name, "")

        doc_filename = f"{table_name}.md"
        doc_filepath = os.path.join(docs_dir, doc_filename)

        if os.path.isfile(doc_filepath):
            if update_doc_file(doc_filepath, columns_md):
                updated += 1
                print(f"  Updated: {doc_filename}")
            else:
                skipped += 1
        else:
            if comment:
                create_doc_file(doc_filepath, table_name, comment, columns_md)
                created += 1
                print(f"  Created: {doc_filename}")
            else:
                print(f"  Skipped (no comment): {table_name}", file=sys.stderr)

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


if __name__ == "__main__":
    main()
