#!/usr/bin/env python3
"""Vendor xkeyboard-config and compile it to a native Zig keymap library.

danos does not ship an X11 runtime, but it wants X11's keyboard *data*: the layout
tables (US, UK, German, ...) that turn a physical key + modifiers into a character.
So we compile xkeyboard-config down to Zig at build time, the same way the initial
ramdisk is packed by a host-side Python tool.

Two subcommands:

  fetch     Download the pinned xkeyboard-config release, verify its sha256, and copy
            the transitively-needed data (the symbols files for the configured layouts
            plus everything they `include`) into library/xkeyboard-config/vendor/,
            alongside a vendored keysymdef.h, the upstream COPYING, and a PROVENANCE.md.
            This is the one step that needs the network; run it when bumping the version.

  generate  Parse the vendored data and emit library/xkeyboard-config/generated/layouts.zig
            (deterministic, offline). Run whenever the layout list or emitter changes.

The pipeline per key: our events carry a USB HID usage; HID_TO_NAME maps it to an XKB
key name (<AC01>); the layout's symbols give the keysyms per level; keysymdef.h resolves
each keysym name to its value and Unicode scalar. Level *selection* (which modifier picks
which level) is deliberately left to the Zig side (xkeyboard-config.zig) — we only emit the
per-key type and its up-to-4 keysyms here.
"""

import argparse
import hashlib
import io
import os
import re
import sys
import tarfile
import urllib.request

# --- pinned upstream -------------------------------------------------------

VERSION = "2.44"
TARBALL_URL = (
    "https://gitlab.freedesktop.org/xkeyboard-config/xkeyboard-config/-/archive/"
    f"xkeyboard-config-{VERSION}/xkeyboard-config-{VERSION}.tar.gz"
)
TARBALL_SHA256 = "35e34edeaf4e8da8d0696ff6b241ee11ddb1b8c6730bac7252d4d0a88ea5f05b"

# keysymdef.h is xorgproto, not xkeyboard-config; fetch copies it from the build host.
KEYSYMDEF_CANDIDATES = [
    "/opt/homebrew/include/X11/keysymdef.h",
    "/usr/include/X11/keysymdef.h",
    "/usr/X11/include/X11/keysymdef.h",
    "/usr/X11R6/include/X11/keysymdef.h",
]

# The layouts we generate: (zig name, symbols file, section or None=default).
TARGETS = [
    ("us", "us", None),
    ("gb", "gb", None),
    ("de", "de", None),
    ("fr", "fr", None),
    ("es", "es", None),
    ("dvorak", "us", "dvorak"),
]

# USB HID usage (keyboard page 0x07) -> XKB key name. The physical keys we can turn into
# characters; positions are the ANSI standard shared by HID usages and XKB names.
HID_TO_NAME = {
    0x04: "AC01", 0x05: "AB05", 0x06: "AB03", 0x07: "AC03", 0x08: "AD03",
    0x09: "AC04", 0x0A: "AC05", 0x0B: "AC06", 0x0C: "AD08", 0x0D: "AC07",
    0x0E: "AC08", 0x0F: "AC09", 0x10: "AB07", 0x11: "AB06", 0x12: "AD09",
    0x13: "AD10", 0x14: "AD01", 0x15: "AD04", 0x16: "AC02", 0x17: "AD05",
    0x18: "AD07", 0x19: "AB04", 0x1A: "AD02", 0x1B: "AB02", 0x1C: "AD06",
    0x1D: "AB01",
    0x1E: "AE01", 0x1F: "AE02", 0x20: "AE03", 0x21: "AE04", 0x22: "AE05",
    0x23: "AE06", 0x24: "AE07", 0x25: "AE08", 0x26: "AE09", 0x27: "AE10",
    0x2C: "SPCE",
    0x2D: "AE11", 0x2E: "AE12", 0x2F: "AD11", 0x30: "AD12", 0x31: "BKSL",
    0x33: "AC10", 0x34: "AC11", 0x35: "TLDE",
    0x36: "AB08", 0x37: "AB09", 0x38: "AB10",
    0x64: "LSGT",  # the extra key on ISO keyboards (102nd key)
}

# XKB type name -> the KeyType enum tag emitted for the Zig side.
TYPE_TO_TAG = {
    "ONE_LEVEL": "one_level",
    "TWO_LEVEL": "two_level",
    "ALPHABETIC": "alphabetic",
    "FOUR_LEVEL": "four_level",
    "FOUR_LEVEL_ALPHABETIC": "four_level_alphabetic",
    "FOUR_LEVEL_SEMIALPHABETIC": "four_level_semialphabetic",
    "FOUR_LEVEL_MIXED_KEYPAD": "keypad",
    "FOUR_LEVEL_KEYPAD": "keypad",
    "KEYPAD": "keypad",
    "PC_SYSRQ": "other",
}

REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
LIB = os.path.join(REPO, "library", "xkeyboard-config")
VENDOR = os.path.join(LIB, "vendor")


# --- keysymdef.h: keysym name -> (value, unicode scalar or None) ------------

def parse_keysymdef(path):
    table = {}
    line_re = re.compile(r"#define\s+XK_(\w+)\s+0x([0-9a-fA-F]+)\s*(?:/\*\s*U\+([0-9A-Fa-f]+)|/\*<U\+([0-9A-Fa-f]+))?")
    with open(path, encoding="latin-1") as f:
        for line in f:
            m = line_re.match(line)
            if not m:
                continue
            name = m.group(1)
            value = int(m.group(2), 16)
            uni = m.group(3) or m.group(4)
            unicode_scalar = int(uni, 16) if uni else None
            # First definition wins (keysymdef lists the canonical one first).
            table.setdefault(name, (value, unicode_scalar))
    return table


def resolve_keysym(name, keysymdef):
    """A keysym token from a symbols file -> (keysym value, unicode scalar or 0)."""
    if name in ("NoSymbol", "VoidSymbol", ""):
        return (0, 0)
    # Unicode escape forms: U00E9 / u00E9.
    m = re.fullmatch(r"[Uu]([0-9A-Fa-f]{4,6})", name)
    if m:
        cp = int(m.group(1), 16)
        return (0x01000000 + cp, cp)
    # Raw hex keysym value: 0x1000161 (a Unicode keysym) or a plain keysym number.
    m = re.fullmatch(r"0x([0-9A-Fa-f]+)", name)
    if m:
        value = int(m.group(1), 16)
        if 0x01000100 <= value <= 0x0110FFFF:
            return (value, value - 0x01000000)
        if 0x20 <= value <= 0x7E or 0xA0 <= value <= 0xFF:
            return (value, value)
        return (value, 0)
    entry = keysymdef.get(name)
    if entry is None:
        return (0, 0)  # unknown named keysym (dead_*, ISO_*, rare) -> no character
    value, unicode_scalar = entry
    return (value, unicode_scalar or 0)


# --- symbols parsing --------------------------------------------------------

def strip_comments(text):
    text = re.sub(r"/\*.*?\*/", "", text, flags=re.S)
    text = re.sub(r"//[^\n]*", "", text)
    return text


class KeyDef:
    __slots__ = ("levels", "type_name")

    def __init__(self, levels, type_name):
        self.levels = levels          # list of keysym-name strings (group 1)
        self.type_name = type_name    # explicit XKB type name, or None


class Symbols:
    """Loads and resolves xkb_symbols sections from a directory of symbols files."""

    def __init__(self, symbols_dir):
        self.dir = symbols_dir
        self._files = {}       # filename -> {section_name: body, "__default__": name}
        self.visited_files = set()

    def _load(self, fname):
        if fname in self._files:
            return self._files[fname]
        path = os.path.join(self.dir, fname)
        self.visited_files.add(fname)
        with open(path, encoding="latin-1") as f:
            text = strip_comments(f.read())
        sections = {}
        default = None
        # Find each `[flags] xkb_symbols "NAME" {` and brace-match its body.
        for m in re.finditer(r'(\w[\w\s]*?)?\bxkb_symbols\s+"([^"]+)"\s*\{', text):
            flags = m.group(1) or ""
            name = m.group(2)
            body, _ = self._brace_match(text, m.end() - 1)
            sections[name] = body
            if "default" in flags.split() and default is None:
                default = name
        if default is None and sections:
            default = next(iter(sections))
        sections["__default__"] = default
        self._files[fname] = sections
        return sections

    @staticmethod
    def _brace_match(text, open_index):
        depth = 0
        for i in range(open_index, len(text)):
            c = text[i]
            if c == "{":
                depth += 1
            elif c == "}":
                depth -= 1
                if depth == 0:
                    return text[open_index + 1:i], i
        raise ValueError("unbalanced braces")

    def resolve(self, fname, section=None, stack=()):
        """Merged {key name -> KeyDef} for (fname, section), following includes."""
        sections = self._load(fname)
        if section is None:
            section = sections["__default__"]
        key = (fname, section)
        if key in stack:
            return {}
        body = sections.get(section)
        if body is None:
            return {}
        keys = {}
        default_type = None
        for stmt in self._statements(body):
            kind = stmt[0]
            if kind == "include":
                augment, inc_file, inc_section = stmt[1], stmt[2], stmt[3]
                inc = self.resolve(inc_file, inc_section, stack + (key,))
                for n, kd in inc.items():
                    if augment and n in keys:
                        continue
                    keys[n] = kd
            elif kind == "default_type":
                default_type = stmt[1]
            elif kind == "key":
                augment, name, kd = stmt[1], stmt[2], stmt[3]
                if augment and name in keys:
                    continue
                keys[name] = kd
        if default_type is not None:
            for kd in keys.values():
                if kd.type_name is None:
                    kd.type_name = default_type
        return keys

    def _statements(self, body):
        """Yield ('include', augment, file, section) / ('default_type', name) /
        ('key', augment, name, KeyDef) in source order."""
        # Recognise the three statement heads and step through the body in order.
        head = re.compile(
            r'(?P<inc>(?:(?P<augi>augment|override|replace)\s+)?include\s+"(?P<incarg>[^"]+)")'
            r'|(?P<dtype>key\.type(?:\[[^\]]*\])?\s*=\s*"(?P<dtypearg>[^"]+)")'
            r'|(?P<key>(?:(?P<augk>augment|override|replace)\s+)?key\s+<(?P<kname>[^>]+)>\s*\{)'
        )
        i = 0
        while i < len(body):
            m = head.search(body, i)
            if not m:
                break
            if m.group("inc"):
                inc_file, inc_section = self._parse_include(m.group("incarg"))
                yield ("include", m.group("augi") == "augment", inc_file, inc_section)
                i = m.end()
            elif m.group("dtype"):
                yield ("default_type", m.group("dtypearg"))
                i = m.end()
            else:  # a key definition; brace-match its body
                kbody, end = self._brace_match(body, m.end() - 1)
                kd = self._parse_key_body(kbody)
                if kd is not None:
                    yield ("key", m.group("augk") == "augment", m.group("kname"), kd)
                i = end + 1

    @staticmethod
    def _parse_include(arg):
        m = re.fullmatch(r"([^()]+)(?:\(([^)]+)\))?", arg.strip())
        return (m.group(1), m.group(2))

    @staticmethod
    def _parse_key_body(kbody):
        type_name = None
        mt = re.search(r'type(?:\[[^\]]*\])?\s*=\s*"([^"]+)"', kbody)
        if mt:
            type_name = mt.group(1)
        groups = re.findall(r"\[([^\]]*)\]", kbody)
        if not groups:
            return None
        levels = [tok.strip() for tok in groups[0].split(",")]
        levels = [t for t in levels if t != ""]
        if not levels:
            return None
        return KeyDef(levels, type_name)


# --- type inference + emission ---------------------------------------------

def is_case_pair(a, b):
    return len(a) == 1 and len(b) == 1 and a.isalpha() and a.islower() and b == a.upper()


def key_type_tag(kd):
    if kd.type_name is not None:
        return TYPE_TO_TAG.get(kd.type_name, "other")
    n = len(kd.levels)
    if n <= 1:
        return "one_level"
    if n == 2:
        return "alphabetic" if is_case_pair(kd.levels[0], kd.levels[1]) else "two_level"
    return "four_level_alphabetic" if is_case_pair(kd.levels[0], kd.levels[1]) else "four_level"


def build_layout(symbols, keysymdef, fname, section):
    keys = symbols.resolve(fname, section)
    table = []  # 256 entries: (tag, [(keysym, unicode) x4])
    for hid in range(256):
        name = HID_TO_NAME.get(hid)
        kd = keys.get(name) if name else None
        if kd is None:
            table.append(("one_level", [(0, 0)] * 4))
            continue
        levels = [resolve_keysym(kd.levels[i], keysymdef) if i < len(kd.levels) else (0, 0)
                  for i in range(4)]
        table.append((key_type_tag(kd), levels))
    return table


def emit_zig(layouts):
    out = io.StringIO()
    out.write("// GENERATED by tools/make-xkeyboard-config.py from xkeyboard-config "
              f"{VERSION}. Do not edit.\n")
    out.write("// Keyboard layout tables compiled from the X11 xkeyboard-config database\n")
    out.write("// (MIT/X11 licensed; see ../vendor/COPYING and ../vendor/PROVENANCE.md).\n\n")
    out.write("pub const Level = struct { keysym: u32 = 0, unicode: u21 = 0 };\n\n")
    out.write("pub const KeyType = enum {\n")
    out.write("    one_level,\n    two_level,\n    alphabetic,\n    four_level,\n"
              "    four_level_alphabetic,\n    four_level_semialphabetic,\n    keypad,\n    other,\n};\n\n")
    out.write("pub const Key = struct { kind: KeyType = .one_level, levels: [4]Level = [_]Level{.{}} ** 4 };\n\n")
    out.write("pub const Layout = struct { name: []const u8, keys: [256]Key };\n\n")

    names = []
    for zig_name, fname, section in TARGETS:
        table = layouts[zig_name]
        names.append(zig_name)
        out.write(f"pub const {zig_name}: Layout = .{{\n")
        out.write(f'    .name = "{zig_name}",\n')
        out.write("    .keys = .{\n")
        for hid, (tag, levels) in enumerate(table):
            if tag == "one_level" and all(k == 0 and u == 0 for k, u in levels):
                out.write("        .{},\n")
                continue
            parts = ", ".join(f".{{ .keysym = {k}, .unicode = {u} }}" for k, u in levels)
            out.write(f"        .{{ .kind = .{tag}, .levels = .{{ {parts} }} }},\n")
        out.write("    },\n};\n\n")

    out.write("pub const all = [_]*const Layout{ " + ", ".join("&" + n for n in names) + " };\n")
    return out.getvalue()


# --- fetch ------------------------------------------------------------------

def collect_needed(symbols):
    for _, fname, section in TARGETS:
        symbols.resolve(fname, section)
    return set(symbols.visited_files)


def find_keysymdef():
    env = os.environ.get("KEYSYMDEF")
    if env and os.path.isfile(env):
        return env
    for p in KEYSYMDEF_CANDIDATES:
        if os.path.isfile(p):
            return p
    sys.exit("error: keysymdef.h not found; install xorgproto or set KEYSYMDEF=/path/to/keysymdef.h")


def cmd_fetch(_args):
    local = os.environ.get("XKB_TARBALL")
    if local:
        data = open(local, "rb").read()
    else:
        print(f"downloading {TARBALL_URL}")
        data = urllib.request.urlopen(TARBALL_URL).read()
    digest = hashlib.sha256(data).hexdigest()
    if digest != TARBALL_SHA256:
        sys.exit(f"error: sha256 mismatch\n  expected {TARBALL_SHA256}\n  got      {digest}")

    tar = tarfile.open(fileobj=io.BytesIO(data), mode="r:gz")
    members = tar.getnames()
    root = members[0].split("/")[0]

    # Extract symbols/ + COPYING to a temp view, resolve includes, keep only what's needed.
    tmp = os.path.join(VENDOR, ".upstream")
    for m in tar.getmembers():
        if m.name.startswith(f"{root}/symbols/") or m.name == f"{root}/COPYING":
            m.name = m.name[len(root) + 1:]
            tar.extract(m, tmp)
    tar.close()

    needed = collect_needed(Symbols(os.path.join(tmp, "symbols")))

    os.makedirs(os.path.join(VENDOR, "symbols"), exist_ok=True)
    for fname in sorted(needed):
        src = os.path.join(tmp, "symbols", fname)
        dst = os.path.join(VENDOR, "symbols", fname)
        os.makedirs(os.path.dirname(dst), exist_ok=True)
        with open(src, "rb") as s, open(dst, "wb") as d:
            d.write(s.read())
    _copy(os.path.join(tmp, "COPYING"), os.path.join(VENDOR, "COPYING"))

    keysymdef = find_keysymdef()
    _copy(keysymdef, os.path.join(VENDOR, "keysymdef.h"))

    _rmtree(tmp)

    with open(os.path.join(VENDOR, "PROVENANCE.md"), "w") as f:
        f.write("# Vendored xkeyboard-config subset\n\n")
        f.write(f"- **Package**: xkeyboard-config {VERSION}\n")
        f.write(f"- **Source**: {TARBALL_URL}\n")
        f.write(f"- **sha256**: `{TARBALL_SHA256}`\n")
        f.write(f"- **keysymdef.h**: xorgproto, copied from `{keysymdef}`\n")
        f.write("- **License**: MIT/X11 (see COPYING)\n\n")
        f.write("Only the symbols files reachable from the generated layouts "
                "(tools/make-xkeyboard-config.py `TARGETS`) are vendored; regenerate with\n"
                "`python3 tools/make-xkeyboard-config.py fetch` then `... generate`.\n\n")
        f.write("Vendored symbols files:\n\n")
        for fname in sorted(needed):
            f.write(f"- `symbols/{fname}`\n")
    print(f"vendored {len(needed)} symbols files + keysymdef.h + COPYING into {VENDOR}")


def cmd_generate(args):
    symbols_dir = args.symbols or os.path.join(VENDOR, "symbols")
    keysymdef_path = args.keysymdef or os.path.join(VENDOR, "keysymdef.h")
    keysymdef = parse_keysymdef(keysymdef_path)
    symbols = Symbols(symbols_dir)
    layouts = {zig_name: build_layout(symbols, keysymdef, fname, section)
               for zig_name, fname, section in TARGETS}
    out_dir = os.path.join(LIB, "generated")
    os.makedirs(out_dir, exist_ok=True)
    out_path = os.path.join(out_dir, "layouts.zig")
    with open(out_path, "w") as f:
        f.write(emit_zig(layouts))
    print(f"wrote {out_path} ({len(TARGETS)} layouts)")


def _copy(src, dst):
    os.makedirs(os.path.dirname(dst), exist_ok=True)
    with open(src, "rb") as s, open(dst, "wb") as d:
        d.write(s.read())


def _rmtree(path):
    for root, dirs, files in os.walk(path, topdown=False):
        for name in files:
            os.remove(os.path.join(root, name))
        for name in dirs:
            os.rmdir(os.path.join(root, name))
    if os.path.isdir(path):
        os.rmdir(path)


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    sub = parser.add_subparsers(dest="command", required=True)
    sub.add_parser("fetch", help="download + vendor the needed xkeyboard-config data")
    g = sub.add_parser("generate", help="emit generated/layouts.zig from the vendored data")
    g.add_argument("--symbols", help="override the vendored symbols/ dir (for development)")
    g.add_argument("--keysymdef", help="override the vendored keysymdef.h (for development)")
    args = parser.parse_args()
    {"fetch": cmd_fetch, "generate": cmd_generate}[args.command](args)


if __name__ == "__main__":
    main()