#!/usr/bin/env python3
 
 
 ## Help
 """
 glregistry 1.0
 
 Cache and query the OpenGL registry locally.
 
 Usage:
   glregistry xml
   glregistry xml-path
   glregistry ext            <extension>
   glregistry ext-path       <extension>
   glregistry exts
   glregistry exts-download
   glregistry exts-all       <name>

   glregistry vendors

   glregistry type           <type>

   glregistry aliases        <enum>

   glregistry value          <enum>
   glregistry enum           <value>
   glregistry supports       <name>
   glregistry names          [<support>]
   glregistry groups         [<enum>]
   glregistry enums          [<group>]
   glregistry enums-tree     [<group>]
   glregistry params         [<group>]
   glregistry params-tree    [<group>]
   glregistry audit          [<path>]
   glregistry audit-tree     [<path>]

   glregistry refs           <name>
   glregistry refs-all       <name>

   glregistry -h|--help
 
 Commands:
   xml
     Download the registry XML and open it with an editor.
   xml-path
     Download the registry XML and print its local path.
   ext <extension>
     Download the <extension> spec and open it with an editor.
   ext-path <extension>
     Download the <extension> spec and print its local path.
   exts
     Print the names of all extension specs.
   exts-download
     Download all extension specs.
   exts-all <name>
     Print all downloaded extensions that mention <name>.

   vendors
     Print all vendor abbreviations.

   type <type>
     Print the definition of <type>.

   aliases <enum>
     Print the KHR, ARB, and EXT aliases of <enum>.

   value <enum>
     Print the value of <enum>.
   enum <value>

     Print the enum(s) that has the given <value>.

   supports <name>
     Print the OpenGL version or extension required to use <name>.
   names [<support>]
     Print the names introduced by the OpenGL version or extension <support> if
     given, or all names if omitted. The special values VERSION and EXTENSION
     print the names introduced by all versions or all extensions respectively.
   groups [<enum>]
     Print the groups of <enum> if given, or all groups if omitted.
   enums [<group>]
     Print the enums in <group> if given, or all enums if omitted.
   enums-tree [<group>]
     Print the enums in <group> if given, or all enums if omitted, sorted on
     support, in a tree.
   params [<group>]
     Print the parameter names of <group> if given, or all parameter names if
     omitted.
   params-tree [<group>]
     Print the parameter names of <group> if given, or all parameter names if
     omitted, sorted on count, together with the commands sorted on support, in
     a tree.
   audit [<path>]
     Search files in <path> if given, or the current directory if omitted,
     recursively for OpenGL API names and print them sorted on location,
     support, and name, in a list.
   audit-tree [<path>]
     Search files in <path> if given, or the current directory if omitted,
     recursively for OpenGL API names and print them sorted on support, name,
     and location, in a tree.

   refs <name>
     Print all URLs of all reference pages with name <name>.
   refs-all <name>
     Print all URLs of all reference pages that mention <name>, sorted on
     support, in a tree.

 
 Environment variables:
   GLREGISTRY_CACHE
     The directory to cache files in. Defaults to `$XDG_CACHE_HOME/glregistry`
     or, if `$XDG_CACHE_HOME` is not defined, `$HOME/.cache/glregistry`.
   GLREGISTRY_EDITOR
     The editor to use when opening files. Defaults to `$EDITOR` or, if
     `$EDITOR` is not defined, `editor` if it exists in `$PATH`, else `vi`. The
     value is interpreted by the shell.
   GLREGISTRY_PAGER
     The pager to use when viewing output. Defaults to `$PAGER` or, if `$PAGER`
     is not defined, `pager` if it exists in `$PATH`, else `less` . The value is
     interpreted by the shell. If the `$LESS` environment variable is unset, it
     is set to `FR`.

   GLREGISTRY_COLORS
     If standard out is a terminal, the colors used in output of the `enums`,
     `enums-tree`, `params`, `params-tree`, `audit,` and `audit-tree` commands.
     It uses the same format (and defaults) as GREP_COLORS, i.e. a
     colon-separated list of capabilties: `ms` (matching selected), `fn` (file
     name), `ln` (line number), `se` (separators). Added custom capabilities
     are: `ve` (version), `ex` (extension), `un` (unsupported). Defaults to
     `ms=01;31:fn=35:ln=32:se=36:ve=01;34:ex=34:un=01;33`.

 """
 
 
 ## Imports
 import os
 import sys
 import collections
 import re

 import functools

 import subprocess
 import shutil
 import shlex
 import urllib.request
 import docopt
 from lxml import etree
 
 
 ## Constants

 REFPAGES_URL  = 'https://registry.khronos.org/OpenGL-Refpages/'

 REGISTRY_URL  = 'https://registry.khronos.org/OpenGL/'

 REFPAGES_GIT  = 'https://github.com/KhronosGroup/OpenGL-Refpages'

 XML_PATH      = 'xml/gl.xml'

 USER_AGENT    = 'Mozilla/5.0'

 REGEX         = r'\b(gl|GL_)[0-9A-Z][0-9A-Za-z_]+\b'
 EXCLUDE_DIRS  = ['.?*', '_*']
 EXCLUDE_FILES = ['README*', 'TODO*']
 INDENT        = 2
 ENV_XDG = lambda var, default: (
     os.environ.get(f'GLREGISTRY_{var}') or
     os.path.join(
         (
             os.environ.get(f'XDG_{var}_HOME') or
             os.path.expanduser(default)
         ),
         'glregistry',
     )
 )
 ENV_PRG = lambda var, default: (
     os.environ.get(f'GLREGISTRY_{var}') or
     os.environ.get(var) or
     shutil.which(var.lower()) or
     default
 )
 CACHE  = ENV_XDG('CACHE',  os.path.join('~', '.cache'))
 EDITOR = ENV_PRG('EDITOR', 'vi')
 PAGER  = ENV_PRG('PAGER',  'less')
 LESS   = os.environ.get('LESS') or 'FR'

 COLORS = collections.defaultdict(str, [
     (color.split('=') + [''])[:2]
     for color in
     filter(None, (
         os.environ.get('GLREGISTRY_COLORS') or
         (lambda x, y: ':'.join([x, x and y]))(
             (
                 os.environ.get('GREP_COLORS') or
                 'ms=01;31:fn=35:ln=32:se=36'
             ),
             've=01;34:ex=34:un=01;33',
         )
     ).split(':'))
 ])

 IN    = lambda a, v, s: f"contains(concat('{s}',@{a},'{s}'),'{s}{v}{s}')"
 MAYBE = lambda a, v:    f"(@{a}='{v}' or not(@{a}))"
 TYPES    = "/registry/types"
 ENUMS    = "/registry/enums"
 COMMANDS = "/registry/commands"
 CATEGORY_ATTRIBS = {
     'VERSION': [
         "/registry/feature[@api='gl']",
         'number',
     ],
     'EXTENSION': [
         f"/registry/extensions/extension[{IN('supported','gl','|')}]",
         'name',
     ],
 }
 REQUIRE = f"require[{MAYBE('api','gl')} and {MAYBE('profile','core')}]"

 REMOVE  = f"remove[{ MAYBE('api','gl')} and {MAYBE('profile','core')}]"
 CHANGE_PREFIXES = [
     [REQUIRE, '' ],
     [REMOVE,  '<'],
 ]

 VENDORS = ['KHR', 'ARB', 'EXT']

 KEY_SUBS = [

     *[[f'^{   prefix}',  f''       ] for _, prefix in CHANGE_PREFIXES],
     *[[f'{    vendor}$', f'{   i}' ] for i, vendor in enumerate(VENDORS)],
     *[[f'^GL_{vendor}_', f'GL_{i}_'] for i, vendor in enumerate(VENDORS)],

 ]

 
 
 ## Helpers
 
 
 ### `log`
 def log(*args, **kwargs):
     print(*args, file=sys.stderr, flush=True, **kwargs)
 
 
 ### `edit`
 def edit(paths):
     if EDITOR and sys.stdout.isatty():
         args = ' '.join([EDITOR, *map(shlex.quote, paths)])
         subprocess.run(args, shell=True)
     else:
         for path in paths:
             with open(path) as f:
                 shutil.copyfileobj(f, sys.stdout)
 
 
 ### `page`
 def page(lines):
     lines = ''.join(f'{line}\n' for line in lines)
     if lines and PAGER and sys.stdout.isatty():
         args = f'LESS={shlex.quote(LESS)} {PAGER}'
         subprocess.run(args, shell=True, text=True, input=lines)
     else:
         sys.stdout.write(lines)
 
 

 ### `color`
 def color(capability, string):
     if not sys.stdout.isatty():
         return string
     return f'\x1b[{COLORS[capability]}m{string}\x1b[m'
 
 
 ### `color_supports`
 def color_supports(supports):
     for support in supports:
         if support == 'UNSUPPORTED' or support.startswith('<'):
             yield color('un', support)
         elif support.startswith('GL_'):
             yield color('ex', support)
         else:
             yield color('ve', support)
 
 

 ### `indentjoin`
 def indentjoin(indent, sep, parts):

     return ' ' * INDENT * indent + color('se', sep).join(map(str, parts))

 
 
 ### `removeprefix`
 def removeprefix(prefix, string):
     if string.startswith(prefix):
         return string[len(prefix):]
     return string
 
 

 ### `key`
 def key(item):
     for sub in KEY_SUBS:
         item = re.sub(*sub, item)
     return item
 
 

 ### `download`
 def download(path, exit_on_failure=True):
     remote = urllib.parse.urljoin(REGISTRY_URL, path)
     local  = os.path.join        (CACHE,        path)
     if not os.path.exists(local):
         try:
             log(f"Downloading '{path}' ... ", end='')
             with urllib.request.urlopen(remote) as response:
                 os.makedirs(os.path.dirname(local), exist_ok=True)
                 with open(local, 'wb') as f:
                     shutil.copyfileobj(response, f)
         except urllib.error.URLError as error:
             log(error.reason)
             if exit_on_failure:
                 exit(1)
         else:
             log(response.reason)
     return local
 
 
 ### `grep`
 def grep(
     path=None,
     regex=REGEX,
     exclude_dirs=EXCLUDE_DIRS,
     exclude_files=EXCLUDE_FILES,
     silent=False,
 ):
     path          = path if path else '.'
     cmd           = ['grep', '-EIrno']
     exclude_dirs  = [f'--exclude-dir={exclude}' for exclude in exclude_dirs]
     exclude_files = [f'--exclude={exclude}'     for exclude in exclude_files]
     process = subprocess.run(
         [*cmd, *exclude_dirs, *exclude_files, regex, path],
         stdout=subprocess.PIPE,
         stderr=subprocess.DEVNULL if silent else None,
         text=True,
     )
     for string in process.stdout.splitlines():
         string           = removeprefix(f'.{os.path.sep}', string)
         file, line, name = string.split(':', 2)
         line             = int(line)
         yield file, line, name
 
 
 ## Commands
 
 
 ### `xml_`
 def xml_():
     return etree.parse(download(XML_PATH))
 
 
 ### `xml`
 def xml():
     return [download(XML_PATH)]
 
 
 ### `ext_`
 def ext_(extension):
     prefix, vendor, name = extension.split('_', 2)
     if prefix != 'GL':
         log("Extension names must start with 'GL_'.")
         exit(1)
     return f'extensions/{vendor}/{vendor}_{name}.txt'
 
 
 ### `ext`
 def ext(extension):
     return [download(ext_(extension))]
 
 
 ### `exts`
 def exts(xml):
     category, attrib = CATEGORY_ATTRIBS['EXTENSION']
     exts = xml.xpath(f"{category}/@{attrib}")

     return sorted(exts, key=key)

 
 
 ### `exts_download`
 def exts_download(xml):
     for ext in exts(xml):
         download(ext_(ext), exit_on_failure=False)
     return []
 
 
 ### `exts_all`
 def exts_all(xml, name):
     # exts_download(xml)
     exts_all = set(
         'GL_' + os.path.splitext(os.path.basename(file))[0]
         for name in set([
             name,
             removeprefix('gl',  name),
             removeprefix('GL_', name),
         ])
         for file, *_ in
         grep(os.path.join(CACHE, 'extensions'), rf'\b{name}\b', [], [])
     )

     return sorted(exts_all, key=key)
 
 
 ### `vendors`
 def vendors(xml):
     vendors = set(extension.split('_')[1] for extension in exts(xml))
     return sorted(vendors, key=key)

 
 
 ### `type`
 def type(xml, type):
     return [xml.xpath(f"string({TYPES}/type/name[text()='{type}']/..)")]
 
 

 ### `aliases`
 def aliases(xml, name, supports_=None, vendors_=[]):
     if not vendors_:
         vendors_[:] = vendors(xml)
     for vendor in vendors_:
         if name.endswith(f'_{vendor}'):
             return []
     value_  = value(xml, name)
     if not value_:
         return []
     if not supports_:
         supports_ = supports(xml, name, False)
     if not supports_:
         return []
     aliases = []
     # for vendor in vendors_:
     for vendor in VENDORS:
         alias = f'{name}_{vendor}'
         if supports(xml, alias, False) and value(xml, alias) == value_:
             aliases.append(alias)
     return sorted(aliases, key=key)
 
 

 ### `value`
 def value(xml, enum):
     return xml.xpath(f"{ENUMS}/enum[@name='{enum}']/@value")
 
 
 ### `enum`
 def enum(xml, value):

     def conv(s):
         return int(s, 16 if s.startswith('0x') else 10)
     value = conv(value)
     enum = (
         enum.get('name')
         for enum in xml.xpath(f"{ENUMS}/enum")
         if conv(enum.get('value')) == value
     )

     return sorted(enum, key=key)

 
 
 ### `supports`

 @functools.cache

 def supports(xml, name, use_aliases=True):

     category, attrib = CATEGORY_ATTRIBS['EXTENSION']
     if xml.xpath(f"{category}[@{attrib}='{name}']"):
         return ['EXTENSION']
     supports_ = [

         f'{prefix}{support}'

         for category, attrib in CATEGORY_ATTRIBS.values()

         for change, prefix   in CHANGE_PREFIXES

         for support in

         xml.xpath(f"{category}/{change}/*[@name='{name}']/../../@{attrib}")

     ]

     if supports_ and use_aliases:
         supports_.extend(
             support
             for alias   in aliases (xml, name, supports_)
             for support in supports(xml, alias, False)
         )

     return sorted(supports_, key=key)

 
 
 ### `names`
 def names(xml, support=None):
     if support in CATEGORY_ATTRIBS.keys():
         category_attribs = [
             [category, ""]
             for category, _ in [CATEGORY_ATTRIBS[support]]
         ]
     elif support:
         category_attribs = [
             [category, f"[@{attrib}='{support}']"]
             for category, attrib in CATEGORY_ATTRIBS.values()
         ]
     else:
         category_attribs = [
             [category, ""]
             for category, _ in CATEGORY_ATTRIBS.values()
         ]
     names = set(
         name
         for category, attrib in category_attribs
         for name in
         xml.xpath(f"{category}{attrib}/{REQUIRE}/*/@name")
     )

     return sorted(names, key=key)

 
 
 ### `groups`
 def groups(xml, enum=None):
     name = f"[@name='{enum}']" if enum else ""
     return sorted(set(
         group
         for groups in xml.xpath(f"{ENUMS}/enum{name}/@group")
         for group  in groups.split(',')
     ))
 
 
 ### `enums_`
 def enums_(xml, group=None):
     group = f"[{IN('group',group,',')}]" if group else ""
     enums_ = collections.defaultdict(list)
     for enum in xml.xpath(f"{ENUMS}/enum{group}/@name"):
         supports_ = supports(xml, enum)
         if supports_:
             enums_[tuple(supports_)].append(enum)
     return enums_
 
 
 ### `enums`
 def enums(xml, group=None):
     enums = [
         enum
         for _, enums in enums_(xml, group).items()
         for enum     in enums
     ]

     return sorted(enums, key=key)

 
 
 ### `enums_tree`
 def enums_tree(xml, group=None):
     for supports, enums in sorted(enums_(xml, group).items()):

         yield indentjoin(0, ',', color_supports(supports))

         for enum in sorted(enums):

             yield indentjoin(1, '', [color('ms', enum)])

 
 
 ### `params_`
 def params_(xml, group=None):
     group   = f"[@group='{group}']" if group else ""
     counts  = collections.defaultdict(int)
     params_ = collections.defaultdict(lambda: collections.defaultdict(list))
     for xmlcommand in xml.xpath(f"{COMMANDS}/command/param{group}/.."):
         command   = xmlcommand.xpath(f"string(proto/name)")
         supports_ = supports(xml, command)
         if supports_:
             for xmlparam in xmlcommand.xpath(f"param{group}"):
                 param = xmlparam.xpath(f"string(name)")
                 params_[param][tuple(supports_)].append(command)
                 counts[param] -= 1
     return {
         (count, param): params_[param]
         for param, count in counts.items()
     }
 
 
 ### `params`
 def params(xml, group=None):
     params = [param for (_, param), _ in params_(xml, group).items()]
     return sorted(params)
 
 
 ### `params_tree`
 def params_tree(xml, group=None):
     for (count, param), occurences in sorted(params_(xml, group).items()):
         yield indentjoin(0, ':', [

             color('ms', param),
             color('ln', -count),

         ])
         for supports_, commands in sorted(occurences.items()):

             yield indentjoin(1, ',', color_supports(supports_))

             for command in sorted(commands):

                 yield indentjoin(2, '', [color('fn', command)])

 
 
 ### `audit_`
 def audit_(xml, path=None):
     audit_ = collections.defaultdict(lambda: collections.defaultdict(list))
     for file, line, name in grep(path):
         supports_ = supports(xml, name)
         if not supports_:
             supports_ = ['UNSUPPORTED']
         audit_[tuple(supports_)][name].append([file, line])
     return audit_
 
 
 ### `audit`
 def audit(xml, path=None):
     for file, line, supports, name in sorted(
         [file, line, supports, name]
         for supports, names  in audit_(xml, path).items()
         for name, locations  in names.items()
         for file, line       in locations
     ):
         yield indentjoin(0, ':', [

             color('fn', file),
             color('ln', line),
             indentjoin(0, ',', color_supports(supports)),
             color('ms', name),

         ])
 
 
 ### `audit_tree`
 def audit_tree(xml, path=None):
     for supports, names in sorted(audit_(xml, path).items()):

         yield indentjoin(0, ',', color_supports(supports))

         for name, locations in sorted(names.items()):

             yield indentjoin(1, '', [color('ms', name)])

             for file, line in sorted(locations):
                 yield indentjoin(2, ':', [

                     color('fn', file),
                     color('ln', line),

                 ])
 
 

 ### `refs_`
 def refs_(name):
     local = os.path.join(CACHE, os.path.basename(REFPAGES_GIT))
     if not os.path.exists(local):
         os.makedirs(os.path.dirname(local), exist_ok=True)
         subprocess.run(['git', 'clone', REFPAGES_GIT, local])
     refs_ = collections.defaultdict(set)
     for file, *_ in grep(local, rf'\b{name}\b', [], [], True):
         file = removeprefix(f'{local}{os.path.sep}', file)
         try:
             support, *_, dir, base = os.path.normpath(file).split(os.path.sep)
         except:
             continue
         if support.startswith('gl') and dir.endswith('html'):
             support   = removeprefix('gl', support)
             name, ext = os.path.splitext(base)
             url       = urllib.parse.urljoin(REFPAGES_URL, file)
             if ext in ['.xml', '.xhtml']:
                 refs_[support].add((name, url))
     return refs_
 
 
 ### `refs`
 def refs(name):
     return sorted(
         url
         for support, locations in refs_(name).items()
         for name_, url         in locations
         if name_ == name
     )
 
 
 ### `refs_all`
 def refs_all(name):
     for support, locations in sorted(refs_(name).items()):

         yield indentjoin(0, ',', color_supports([support]))

         for name_, url in sorted(locations):
             yield indentjoin(1, ':', [

                 color('ms', name_),
                 color('fn', url),

             ])
 
 

 ## Main
 def main():

     opener = urllib.request.build_opener()
     opener.addheaders = [('User-Agent', USER_AGENT)]
     urllib.request.install_opener(opener)

     args = docopt.docopt(__doc__)
     if args['xml']:           edit(xml          ())
     if args['xml-path']:      page(xml          ())
     if args['ext']:           edit(ext          (args['<extension>']))
     if args['ext-path']:      page(ext          (args['<extension>']))
     if args['exts']:          page(exts         (xml_()))
     if args['exts-download']: page(exts_download(xml_()))
     if args['exts-all']:      page(exts_all     (xml_(), args['<name>']))

     if args['vendors']:       page(vendors      (xml_()))

     if args['type']:          page(type         (xml_(), args['<type>']))

     if args['aliases']:       page(aliases      (xml_(), args['<enum>']))

     if args['value']:         page(value        (xml_(), args['<enum>']))
     if args['enum']:          page(enum         (xml_(), args['<value>']))
     if args['supports']:      page(supports     (xml_(), args['<name>']))
     if args['names']:         page(names        (xml_(), args['<support>']))
     if args['groups']:        page(groups       (xml_(), args['<enum>']))
     if args['enums']:         page(enums        (xml_(), args['<group>']))
     if args['enums-tree']:    page(enums_tree   (xml_(), args['<group>']))
     if args['params']:        page(params       (xml_(), args['<group>']))
     if args['params-tree']:   page(params_tree  (xml_(), args['<group>']))
     if args['audit']:         page(audit        (xml_(), args['<path>']))
     if args['audit-tree']:    page(audit_tree   (xml_(), args['<path>']))

     if args['refs']:          page(refs         (args['<name>']))
     if args['refs-all']:      page(refs_all     (args['<name>']))

 
 
 if __name__ == '__main__':
     main()