#!/usr/bin/env python3
import re
from collections import OrderedDict

DOC_PATH = "docs/cql/built-in-tables.md"

PHYSICAL_KEYSPACES = {
    "system",
    "system_schema",
    "system_auth",
    "system_traces",
    "system_distributed",
    "system_cluster_metadata",
    "system_accord",
}

VIRTUAL_KEYSPACES = {
    "system_views",
    "system_metrics",
    "system_accord_debug",
    "system_accord_debug_remote",
}


def strip_quotes(name: str) -> str:
    name = name.strip()
    if len(name) >= 2 and name[0] == '"' and name[-1] == '"':
        return name[1:-1]
    return name


def split_top_level(s: str, sep: str = ','):
    parts = []
    depth = 0
    start = 0
    for i, ch in enumerate(s):
        if ch == '(':
            depth += 1
        elif ch == ')':
            depth -= 1
        elif ch == sep and depth == 0:
            parts.append(s[start:i].strip())
            start = i + 1
    parts.append(s[start:].strip())
    return [p for p in parts if p]


def parse_primary_key(line: str):
    # line like "PRIMARY KEY ((a, b), c, d)"
    start = line.find('(')
    end = line.rfind(')')
    if start == -1 or end == -1 or end <= start:
        return [], []
    inner = line[start + 1:end].strip()
    if not inner:
        return [], []

    if inner.startswith('('):
        depth = 0
        close_idx = None
        for i, ch in enumerate(inner):
            if ch == '(':
                depth += 1
            elif ch == ')':
                depth -= 1
                if depth == 0:
                    close_idx = i
                    break
        if close_idx is None:
            return [], []
        partition_part = inner[1:close_idx]
        rest = inner[close_idx + 1:].strip()
        if rest.startswith(','):
            rest = rest[1:].strip()
        partition_keys = [strip_quotes(p) for p in split_top_level(partition_part)]
        clustering_keys = [strip_quotes(c) for c in split_top_level(rest)] if rest else []
        return partition_keys, clustering_keys

    parts = split_top_level(inner)
    if not parts:
        return [], []
    partition_keys = [strip_quotes(parts[0])]
    clustering_keys = [strip_quotes(p) for p in parts[1:]]
    return partition_keys, clustering_keys


def parse_clustering_order(line: str):
    # line contains "CLUSTERING ORDER BY (col1 ASC, col2 DESC)"
    m = re.search(r'CLUSTERING ORDER BY\s*\(([^\)]+)\)', line)
    if not m:
        return None
    inner = m.group(1)
    order = {}
    for item in split_top_level(inner):
        bits = item.split()
        if len(bits) >= 2:
            order[strip_quotes(bits[0])] = bits[1].upper()
    return order


def parse_group_tables(doc_text: str):
    # Parse system_metrics group tables list
    group_tables = []
    m = re.search(r"Tables:\s+(.+)", doc_text)
    if not m:
        return group_tables
    line = m.group(1)
    group_tables = re.findall(r"`([^`]+)`", line)
    return group_tables


def parse_accord_debug_remote_tables(doc_text: str):
    m = re.search(r"system_accord_debug_remote.*?Tables:\s+(.+)", doc_text, re.S)
    if not m:
        return []
    line = m.group(1).strip()
    return re.findall(r"`([^`]+)`", line)


def parse_doc(path: str):
    with open(path, 'r', encoding='utf-8') as f:
        lines = f.readlines()
    doc_text = ''.join(lines)

    tables = []
    in_sql = False
    current = None
    last_table = None

    create_re = re.compile(r"^CREATE TABLE\s+(?:IF NOT EXISTS\s+)?([^\s(]+)")

    for raw in lines:
        line = raw.rstrip('\n')
        stripped = line.strip()
        if stripped.startswith('```sql'):
            in_sql = True
            current = None
            last_table = None
            continue
        if in_sql and stripped.startswith('```'):
            in_sql = False
            current = None
            last_table = None
            continue
        if not in_sql:
            continue

        m = create_re.match(stripped)
        if m:
            full_name = m.group(1)
            if '<' in full_name or '>' in full_name:
                current = None
                last_table = None
                continue
            if '.' not in full_name:
                current = None
                last_table = None
                continue
            keyspace, table = full_name.split('.', 1)
            keyspace = strip_quotes(keyspace)
            table = strip_quotes(table)
            current = {
                'keyspace': keyspace,
                'table': table,
                'columns': [],
                'partition_keys': [],
                'clustering_keys': [],
                'clustering_order': {},
                'inline_pk': None,
            }
            tables.append(current)
            last_table = current
            continue

        if current:
            if not stripped:
                continue
            if stripped.startswith('PRIMARY KEY'):
                pk, ck = parse_primary_key(stripped)
                current['partition_keys'] = pk
                current['clustering_keys'] = ck
                continue
            if stripped.startswith(')'):
                order = parse_clustering_order(stripped)
                if order:
                    current['clustering_order'].update(order)
                current = None
                continue

            if stripped.endswith(','):
                stripped = stripped[:-1]
            parts = stripped.split(None, 1)
            if len(parts) < 2:
                continue
            col = strip_quotes(parts[0])
            ctype = parts[1].strip()
            if 'PRIMARY KEY' in ctype:
                ctype = ctype.split('PRIMARY KEY')[0].strip()
                current['inline_pk'] = [col]
            current['columns'].append((col, ctype))
            continue

        if last_table:
            order = parse_clustering_order(stripped)
            if order:
                last_table['clustering_order'].update(order)

    for table in tables:
        if not table['partition_keys'] and table['inline_pk']:
            table['partition_keys'] = table['inline_pk']
            table['clustering_keys'] = []

    group_tables = parse_group_tables(doc_text)
    for table_name in group_tables:
        tables.append({
            'keyspace': 'system_metrics',
            'table': table_name,
            'columns': [
                ('name', 'text'),
                ('scope', 'text'),
                ('type', 'text'),
                ('value', 'text'),
            ],
            'partition_keys': ['name'],
            'clustering_keys': [],
            'clustering_order': {},
            'inline_pk': None,
        })

    remote_tables = parse_accord_debug_remote_tables(doc_text)
    local_defs = {t['table']: t for t in tables if t['keyspace'] == 'system_accord_debug'}
    for table_name in remote_tables:
        local = local_defs.get(table_name)
        if not local:
            continue
        local_pk = local['partition_keys']
        local_ck = local['clustering_keys']
        pk_set = set(local_pk)
        ck_set = set(local_ck)
        regular_cols = [c for c, _ in local['columns'] if c not in pk_set and c not in ck_set]

        columns = [('node_id', 'int')]
        for c in local_pk:
            ctype = next(ct for name, ct in local['columns'] if name == c)
            columns.append((c, ctype))
        for c in local_ck:
            ctype = next(ct for name, ct in local['columns'] if name == c)
            columns.append((c, ctype))
        for c in regular_cols:
            ctype = next(ct for name, ct in local['columns'] if name == c)
            columns.append((c, ctype))

        tables.append({
            'keyspace': 'system_accord_debug_remote',
            'table': table_name,
            'columns': columns,
            'partition_keys': ['node_id'],
            'clustering_keys': list(local_pk) + list(local_ck),
            'clustering_order': local.get('clustering_order', {}).copy(),
            'inline_pk': None,
        })

    return tables


def emit_go(defs, package_name, var_name):
    lines = []
    lines.append("package %s" % package_name)
    lines.append("")
    lines.append("type columnDef struct {")
    lines.append("\tName string")
    lines.append("\tType string")
    lines.append("}")
    lines.append("")
    lines.append("type tableDef struct {")
    lines.append("\tKeyspace string")
    lines.append("\tTable string")
    lines.append("\tColumns []columnDef")
    lines.append("\tPartitionKeys []string")
    lines.append("\tClusteringKeys []string")
    lines.append("\tClusteringOrder map[string]string")
    lines.append("}")
    lines.append("")
    lines.append("var %s = []tableDef{" % var_name)

    for t in defs:
        lines.append("\t{")
        lines.append("\t\tKeyspace: \"%s\"," % t['keyspace'])
        lines.append("\t\tTable: \"%s\"," % t['table'])
        lines.append("\t\tColumns: []columnDef{")
        for name, ctype in t['columns']:
            lines.append("\t\t\t{Name: \"%s\", Type: \"%s\"}," % (name, ctype))
        lines.append("\t\t},")
        if t['partition_keys']:
            lines.append("\t\tPartitionKeys: []string{%s}," % ", ".join("\"%s\"" % k for k in t['partition_keys']))
        else:
            lines.append("\t\tPartitionKeys: nil,")
        if t['clustering_keys']:
            lines.append("\t\tClusteringKeys: []string{%s}," % ", ".join("\"%s\"" % k for k in t['clustering_keys']))
        else:
            lines.append("\t\tClusteringKeys: nil,")
        if t['clustering_order']:
            lines.append("\t\tClusteringOrder: map[string]string{")
            for k, v in t['clustering_order'].items():
                lines.append("\t\t\t\"%s\": \"%s\"," % (k, v))
            lines.append("\t\t},")
        else:
            lines.append("\t\tClusteringOrder: nil,")
        lines.append("\t},")
    lines.append("}")
    return "\n".join(lines) + "\n"


def main():
    tables = parse_doc(DOC_PATH)
    physical = [t for t in tables if t['keyspace'] in PHYSICAL_KEYSPACES]
    virtual = [t for t in tables if t['keyspace'] in VIRTUAL_KEYSPACES and t['keyspace'] != 'system_virtual_schema']

    physical_go = emit_go(physical, "systemtables", "builtinSystemTableDefs")
    virtual_go = emit_go(virtual, "virtualtables", "builtinVirtualTableDefs")

    with open("pkg/cql/systemtables/builtins_gen.go", "w", encoding="utf-8") as f:
        f.write(physical_go)
    with open("pkg/cql/virtualtables/builtins_gen.go", "w", encoding="utf-8") as f:
        f.write(virtual_go)


if __name__ == '__main__':
    main()