# Version of mathdef.py for bikeshed build.
# Tweaked to generate single page links.
# TODO(bradnelson): Figure out a way to merge this back into
# mathdef.py controlled by buildername.

from sphinx.ext.mathbase import math
from sphinx.ext.mathbase import MathDirective
from sphinx.ext.mathjax import html_visit_math
from sphinx.ext.mathjax import html_visit_displaymath
from sphinx.writers.html5 import HTML5Translator
from docutils import nodes
from docutils.parsers.rst.directives.misc import Replace
import re


# Transform \xref in math nodes

xref_re = re.compile('\\\\xref\{([^}]*)\}\{([^}]*)\}', re.M)

def html_hyperlink(file, id):
  return '\\href{#%s}' % (id.replace('_', '-'))

def html_transform_math_xref(node):
  new_text = xref_re.sub(lambda m: html_hyperlink(m.group(1), m.group(2)), node.astext())
  node.children[0] = nodes.Text(new_text)

# Expand mathdef names in math roles and directives

def_re = re.compile('\\\\[A-Za-z][0-9A-Za-z]*', re.M)

auxcounter = 0

def lookup_mathdef(defs, name):
  if name in defs:
    [arity, s] = defs[name]
    if arity > 0:
      global auxcounter
      auxcounter = auxcounter + 1
      name = "\\mathdef%d" % auxcounter
      s = "\\def%s#%d{%s}%s" % (name, arity, s, name)
    return s
  return name

def replace_mathdefs(doc, s):
  if not hasattr(doc, 'mathdefs'):
    return s
  return def_re.sub(lambda m: lookup_mathdef(doc.mathdefs, m.group(0)), s)

def ext_math_role(role, raw, text, line, inliner, options = {}, content = []):
  text = replace_mathdefs(inliner.document, raw.split('`')[1])
  return [math(raw, text)], []

class ExtMathDirective(MathDirective):
  def run(self):
    doc = self.state.document
    for i, s in enumerate(self.content):
      self.content[i] = replace_mathdefs(doc, s)
    for i, s in enumerate(self.arguments):
      self.arguments[i] = replace_mathdefs(doc, s)
    return super(ExtMathDirective, self).run()

class MathdefDirective(Replace):
  def run(self):
    name = '\\' + self.state.parent.rawsource.split('|')[1]
    name = name.split('#')
    if len(name) > 1:
      arity = int(name[1])
    else:
      arity = 0
    name = name[0]
    doc = self.state.document
    if not hasattr(doc, 'mathdefs'):
      doc.mathdefs = {}
    for i, s in enumerate(self.content):
      self.content[i] = replace_mathdefs(doc, s)
    doc.mathdefs[name] = [arity, ''.join(self.content)]
    self.content[0] = ':math:`' + self.content[0]
    self.content[-1] = self.content[-1] + '`'
    return super(MathdefDirective, self).run()

class WebAssemblyHTML5Translator(HTML5Translator):
  """
  Customize HTML5Translator.
  Convert xref in math and math block nodes to hrefs.
  """
  def visit_math(self, node, math_env = ''):
    html_transform_math_xref(node)
    super().visit_math(node, math_env)

  def visit_math_block(self, node, math_env  = ''):
    html_transform_math_xref(node)
    super().visit_math_block(node, math_env)

# Setup

def setup(app):
  app.set_translator('singlehtml', WebAssemblyHTML5Translator)
  app.add_role('math', ext_math_role)
  app.add_directive('math', ExtMathDirective, override = True)
  app.add_directive('mathdef', MathdefDirective)
