#!/usr/bin/env python3
"""Extract the wS interpreter dispatch tree from stage2.lua and resolve the
implementation block for every opcode value."""
import re, json

SRC = '/tmp/opencode/stage2.luau'
KW = {'if','then','else','elseif','end','while','do','for','repeat','until',
      'function','local','return','break','continue','and','or','not','in','nil','true','false'}

def tokenize(s):
    toks = []
    i, n = 0, len(s)
    while i < n:
        c = s[i]
        if c in ' \t\r\n':
            i += 1; continue
        if s.startswith('--', i):
            j = s.find('\n', i)
            i = n if j < 0 else j
            continue
        if c in '"\'':
            q = c; j = i + 1
            while j < n:
                if s[j] == '\\': j += 2; continue
                if s[j] == q: j += 1; break
                j += 1
            toks.append(('str', s[i:j])); i = j; continue
        if c.isdigit() or (c == '.' and i+1 < n and s[i+1].isdigit()):
            j = i
            while j < n and (s[j].isalnum() or s[j] in '._'):
                if s[j] == '.' and j+1 < n and s[j+1] == '.':
                    break
                j += 1
            toks.append(('num', s[i:j])); i = j; continue
        if c.isalpha() or c == '_':
            j = i
            while j < n and (s[j].isalnum() or s[j] == '_'):
                j += 1
            w = s[i:j]
            toks.append(('kw' if w in KW else 'id', w)); i = j; continue
        for op in ('...', '..', '==', '~=', '<=', '>=', '+=', '-=', '*=', '/=', '%=', '^=', '//'):
            if s.startswith(op, i):
                toks.append(('op', op)); i += len(op); break
        else:
            toks.append(('op', c)); i += 1
    return toks

def header_end(toks, i):
    v = toks[i][1]
    if v == 'if':
        while toks[i][1] != 'then': i += 1
        return i + 1
    if v in ('while', 'for'):
        while toks[i][1] != 'do': i += 1
        return i + 1
    if v == 'do':
        return i + 1
    if v == 'function':
        while toks[i][1] != '(': i += 1
        d = 0
        while True:
            if toks[i][1] == '(': d += 1
            elif toks[i][1] == ')':
                d -= 1
                if d == 0:
                    return i + 1
            i += 1
    if v == 'repeat':
        return i + 1
    raise ValueError(v)

def match_end(toks, i):
    """i = index of opening keyword; return index of terminating end/until."""
    v = toks[i][1]
    if v == 'repeat':
        j = header_end(toks, i)
        while True:
            t, w = toks[j]
            if t == 'kw':
                if w in ('if','while','for','function','do','repeat'):
                    j = match_end(toks, j) + 1
                    continue
                if w == 'until':
                    return j
            j += 1
    j = header_end(toks, i)
    while True:
        t, w = toks[j]
        if t == 'kw':
            if w in ('if','while','for','function','do','repeat'):
                j = match_end(toks, j) + 1
                continue
            if w == 'end':
                return j
        j += 1

def find_zif(toks, a, b):
    depth = 0; i = a
    while i < b:
        t, v = toks[i]
        if t == 'kw':
            if v == 'if' and depth == 0:
                j = i + 1; hasz = False
                while j < b and not (toks[j][0] == 'kw' and toks[j][1] == 'then'):
                    if toks[j][1] == 'Z': hasz = True
                    j += 1
                if hasz:
                    return i
            if v in ('if','while','for','function','do','repeat'):
                depth += 1
            elif v in ('end','until'):
                depth -= 1
        i += 1
    return None

def parse_if(toks, i):
    """returns cond, then_span, else_span|None, next_i"""
    i = header_end(toks, i)
    ta = i
    e = match_end(toks, ta - 1) if False else None
    # find else/elseif/end at this level
    j = ta
    while True:
        t, v = toks[j]
        if t == 'kw':
            if v in ('if','while','for','function','do','repeat'):
                j = match_end(toks, j) + 1
                continue
            if v in ('else','elseif','end'):
                break
        j += 1
    header = ' '.join(x for _, x in toks[i - 0:ta])
    # recompute cond text: from after 'if' to 'then'
    k = i
    return None

def parse_if2(toks, i, elseif=False):
    j = i + 1
    if not elseif:
        assert toks[i][1] == 'if'
    ca = j
    while not (toks[j][0] == 'kw' and toks[j][1] == 'then'):
        j += 1
    cond = ' '.join(v for _, v in toks[ca:j])
    body_start = j + 1
    j = body_start
    while True:
        t, v = toks[j]
        if t == 'kw':
            if v in ('if','while','for','function','do','repeat'):
                j = match_end(toks, j) + 1
                continue
            if v in ('else','elseif','end'):
                break
        j += 1
    tspan = (body_start, j)
    if toks[j][1] == 'end':
        return cond, tspan, None, j + 1
    if toks[j][1] == 'elseif':
        c2, t2, e2, nxt = parse_if2(toks, j, True)
        # represent elseif as an else branch holding a node
        return cond, tspan, ('node', c2, t2, e2), nxt
    ea = j + 1
    k = ea
    while True:
        t, v = toks[k]
        if t == 'kw':
            if v in ('if','while','for','function','do','repeat'):
                k = match_end(toks, k) + 1
                continue
            if v == 'end':
                return cond, tspan, ('else', (ea, k)), k + 1
        k += 1

def find_zif_any(toks, a, b):
    i = a
    while i < b:
        t, v = toks[i]
        if t == 'kw' and v in ('if', 'elseif'):
            j = i + 1
            while j < b and not (toks[j][0] == 'kw' and toks[j][1] == 'then'):
                if toks[j][1] == 'Z':
                    return i
                j += 1
        i += 1
    return None

def build(toks, a, b):
    i = find_zif_any(toks, a, b)
    if i is None:
        return {'type': 'leaf', 'span': (a, b)}
    cond, tb, eb, nxt = parse_if2(toks, i, toks[i][1] == 'elseif')
    node = {'type': 'node', 'cond': cond, 'tb': build(toks, *tb)}
    if eb is None:
        node['eb'] = {'type': 'leaf', 'span': (nxt, nxt)}
    elif eb[0] == 'else':
        node['eb'] = build(toks, *eb[1])
    else:
        _, c2, t2, e2 = eb
        n2 = {'type': 'node', 'cond': c2, 'tb': build(toks, *t2)}
        n2['eb'] = build(toks, *e2[1]) if e2 and e2[0] == 'else' else {'type': 'leaf', 'span': (nxt, nxt)}
        node['eb'] = n2
    return node

def body(toks, span):
    return ' '.join(v for _, v in toks[span[0]:span[1]])

def translate(c):
    return c.replace('~=', '!=')

def resolve(tree, v):
    n = tree
    while n['type'] == 'node':
        n = n['tb'] if eval(translate(n['cond']), {'Z': v}) else n['eb']
    return n

def main():
    src = open(SRC, encoding='latin-1').read()
    a = src.index('local Z = o[Q]')
    b = src.index('Q += 1.0', a)
    toks = tokenize(src[a:b])
    i = 0
    while not (toks[i][0] == 'kw' and toks[i][1] == 'if'):
        i += 1
    tree = build(toks, i, len(toks))
    out = {}
    for v in range(0, 256):
        out[v] = body(toks, resolve(tree, v)['span'])
    json.dump(out, open('/tmp/opencode/isa.json', 'w'), indent=0)
    print('ops', len(out), 'distinct bodies', len(set(out.values())))
    for v in list(range(0, 16)) + [112, 173, 185, 200, 255]:
        print(v, '|', out[v][:160])

if __name__ == '__main__':
    main()