import sys

# Address where the Forth dictionary is stored
baseAddr = 0x3C51
# set-context routine in the ROM - the actionfor  words created with the
# VOCABULARY word
SETCONTEXT = 0x11B5
# address of code field for word "FORTH"
FORTHWORD = 0x3C4A

wordsROM = {
  0x3C4A: 'FORTH',
  0x1D59: 'UFLOAT',
  0x1D22: 'INT',
  0x1D0F: 'FNEGATE',
  0x1C7B: 'F/',
  0x1C4B: 'F*',
  0x1BB1: 'F+',
  0x1BA4: 'F-',
  0x198A: 'LOAD',
  0x1979: 'BVERIFY',
  0x1967: 'VERIFY',
  0x1954: 'BLOAD',
  0x1944: 'BSAVE',
  0x1934: 'SAVE',
  0x1670: 'LIST',
  0x165E: 'EDIT',
  0x1638: 'FORGET',
  0x13FD: 'REDEFINE',
  0x13F0: 'EXIT',
  0x1396: '."',
  0x1379: '(',
  0x13D5: '[',
  0x133C: '+LOOP',
  0x1332: 'LOOP',
  0x1323: 'DO',
  0x128D: 'UNTIL',
  0x1276: 'REPEAT',
  0x129F: 'BEGIN',
  0x12A4: 'THEN',
  0x1271: 'ELSE',
  0x1288: 'WHILE',
  0x1283: 'IF',
  0x13E1: ']',
  0x1316: 'LEAVE',
  0x1302: 'J',
  0x12F7: 'I\'',
  0x12E9: 'I',
  0x11AB: 'DEFINITIONS',
  0x117D: 'VOCABULARY',
  0x1160: 'IMMEDIATE',
  0x1140: 'RUNS>',
  0x10E8: 'DOES>',
  0x10F5: 'COMPILER',
  0x10A7: 'CALL',
  0x1074: 'DEFINER',
  0x104B: 'ASCII',
#  0x1006: 'LITERAL',  # should never appear
  0x0FE2: 'CONSTANT',
  0x0FCF: 'VARIABLE',
  0x0F76: 'ALLOT',
  0x0F5F: 'C,',
  0x0F4E: ',',
  0x0ED0: 'CREATE',
  0x0EAF: ':',
  0x0EA3: 'DECIMAL',
  0x0E87: 'MIN',
  0x0E75: 'MAX',
  0x0E60: 'XOR',
  0x0E4B: 'AND',
  0x0E36: 'OR',
  0x0E29: '2-',
  0x0E1F: '1-',
  0x0E13: '2+',
  0x0E09: '1+',
  0x0DEE: 'D+',
  0x0DE1: '-',
  0x0DD2: '+',
  0x0DBA: 'DNEGATE',
  0x0DA9: 'NEGATE',
  0x0D8C: 'U/MOD',
  0x0D7A: '*/',
  0x0D6D: '*',
  0x0D61: 'MOD',
  0x0D51: '/',
  0x0D31: '*/MOD',
  0x0D00: '/MOD',
  0x0CA8: 'U*',
  0x0C83: 'D<',
  0x0C72: 'U<',
  0x0C65: '<',
  0x0C56: '>',
  0x0C4A: '=',
  0x0C3A: '0>',
  0x0C2E: '0<',
  0x0C1A: '0=',
  0x0C0D: 'ABS',
  0x0BFD: 'OUT',
  0x0BEB: 'IN',
  0x0BDB: 'INKEY',
  0x0B98: 'BEEP',
  0x0B4A: 'PLOT',
  0x0B19: 'AT',
  0x0AAF: 'F.',
  0x0AA3: 'EMIT',
  0x0A95: 'CR',
  0x0A83: 'SPACES',
  0x0A73: 'SPACE',
  0x0A5C: 'HOLD',
  0x0A1D: 'CLS',
  0x09F7: '#',
  0x09E1: '#S',
  0x09D0: 'U.',
  0x09B3: '.',
  0x0A4A: 'SIGN',
  0x099C: '#>',
  0x098D: '<#',
  0x096E: 'TYPE',
  0x0933: 'ROLL',
  0x0925: 'PICK',
  0x0912: 'OVER',
  0x08FF: 'ROT',
  0x08EE: '?DUP',
  0x08DF: 'R>',
  0x08D2: '>R',
  0x08C1: '!',
  0x08B3: '@',
  0x08A5: 'C!',
  0x0896: 'C@',
  0x0885: 'SWAP',
  0x0879: 'DROP',
  0x086B: 'DUP',
  0x0846: 'SLOW',
  0x0837: 'FAST',
  0x0828: 'INVIS',
  0x0818: 'VIS',
  0x078A: 'CONVERT',
  0x06A9: 'NUMBER',
  0x069A: 'EXECUTE',
  0x063D: 'FIND',
  0x062D: 'VLIST',
  0x05AB: 'WORD',
  0x0578: 'RETYPE',
  0x058C: 'QUERY',
  0x0506: 'LINE',
  0x04B6: ';',
  0x0499: 'PAD',
  0x048A: 'BASE',
  0x0480: 'CURRENT',
  0x0473: 'CONTEXT',
  0x0460: 'HERE',
  0x00AB: 'ABORT',
  0x0099: 'QUIT',
}

def error(ret, msg):
  sys.stderr.write(msg + "\n")
  return ret

def wordAt(data, ofs):
  return data[ofs] + 256 * data[ofs + 1]

def peek(data, ofs):
  return data[ofs - baseAddr]

def wpeek(data, ofs):
  return wordAt(data, ofs - baseAddr)

def getName(stream, addr):
  name = stream[addr - baseAddr - 4 - (peek(stream, addr) & 0x3F) :
    addr - baseAddr - 4]
  name[-1] &= 0x7F
  return name.decode('latin1')

def pathToWord(word, context, parents):
  path = []
  target = word
  while target != context and target != FORTHWORD:
    target = parents[target]
    path.append(target)
  path.reverse()
  return path

def endWord(stream, addr):
  if peek(stream, addr - 1) & 0x40:
    sys.stdout.write(' IMMEDIATE\n')
  else:
    sys.stdout.write('\n')

def toFloat(bcd):
  s = ''
  if bcd & 0x80000000:
    s = '-'
  exp = (bcd & 0x7F000000) >> 24
  if exp == 0:
    exp = 65
  if 61 <= exp <= 73:
    # Use standard notation
    mant = ''
    bcd &= 0x00FFFFFF
    for i in range(6):
      if bcd == 0:
        break
      mant += chr(((bcd >> 20) & 0xF) + ord('0'))
      bcd = (bcd << 4) & 0x00FFFFFF
    mant = '0' * (64 - exp) + mant
    if mant == '':
      mant = '0' * (exp - 64)
    mant = mant[:max(0, exp - 64)] + '.' + mant[max(0, exp - 64):]
    s += mant + ' '
    return s

  # Scientific notation
  for i in range(6):
    s += chr(((bcd >> 20) & 0xF) + ord('0'))
    bcd = (bcd << 4) & 0xFFFFFF
    if i == 0:
      s += '.'
    if bcd == 0:
      break
  s += 'E'
  s += str(exp - 65)
  return s


def listThread(stream, insptr, names, voc, parents):
  ncol = 2
  indent = 2
  newLine = True
  while True:
    word = wpeek(stream, insptr)
    insptr += 2

    if word in {0x04B6, 0x10E8, 0x1140, 0x133C, 0x1332, 0x128D, 0x1276,
        0x12A4, 0x1271, 0x1288}:
      # ; DOES> RUNS> +LOOP LOOP UNTIL REPEAT THEN ELSE WHILE
      # These words cause an unindent.
      indent -= 2
      newLine = True
    if word in {0x1283, 0x1323, 0x129F, 0x1288, 0x1271}:
      # IF DO BEGIN WHILE ELSE
      # These words initiate a new line
      newLine = True

    if newLine:
      sys.stdout.write('\n')
      sys.stdout.write(' ' * indent)
      ncol = indent
      newLine = False
    else:
      sys.stdout.write(' ')
      ncol += 1


    if word == 0x1064:
      # decompile float
      s = toFloat(wpeek(stream, insptr) + (wpeek(stream, insptr + 2) << 16))
      sys.stdout.write(s)
      ncol += len(s)
      insptr += 4
    elif word == 0x1011:
      # decompile int
      i = wpeek(stream, insptr)
      s = '%d' % (i if i < 32768 else i - 65536)
      sys.stdout.write(s)
      ncol += len(s)
      insptr += 2
    elif word == 0x104B:
      # ASCII
      sys.stdout.write('ASCII %c' % peek(stream, insptr))
      ncol += 7
      insptr += 1
    elif word in {0x1140, 0x10E8}:
      # RUNS> DOES>
      # these words indent
      sys.stdout.write(wordsROM[word])
      ncol += len(wordsROM[word])
      insptr += 5
      indent += 2
      newLine = True
    elif word in {0x1283, 0x1323, 0x129F, 0x1288, 0x1271}:
      # IF DO BEGIN WHILE ELSE
      # these words indent
      sys.stdout.write(wordsROM[word])
      ncol += len(wordsROM[word])
      indent += 2
      newLine = True
    elif word in {0x1379, 0x1396}:
      # ( ."
      sys.stdout.write(wordsROM[word] + ' ')
      ncol += len(wordsROM[word]) + 1
      Len = wpeek(stream, insptr)
      for i in range(insptr + 2, insptr + 2 + Len):
        sys.stdout.write(chr(peek(stream, i)))
        ncol += 1
      sys.stdout.write('"' if word == 0x1396 else ')')
      ncol += 1
      newLine = True
      insptr = insptr + 2 + Len
    elif word in wordsROM:
      sys.stdout.write(wordsROM[word])
      ncol += len(wordsROM[word])
    elif word in names:
      if parents[word] != voc:
        s = pathToWord(word, voc, parents)
        s = ' '.join(names[i] for i in s)
        ncol += len(s) + 5
        sys.stdout.write('[ ' + s + ' ] ')
        voc = parents[word]
      sys.stdout.write(names[word])
      ncol += len(names[word])
    elif (word >= baseAddr and word < baseAddr + len(stream) - 1
          and wpeek(stream, word) == 0x1142):
      # word defined with COMPILER ... RUNS>
      if parents[word] != voc:
        s = pathToWord(word, voc, parents)
        s = ' '.join(names[i] for i in s)
        ncol += len(s) + 5
        sys.stdout.write('[ ' + s + ' ] ')
        voc = parents[word]
      defName = word + wpeek(stream, word - 2) - 0x10001
      # defName points to the name field; skip it and point to the code field
      while peek(stream, defName) < 128:
        defName += 1
      sys.stdout.write(names[defName + 6])
      Len = peek(stream, word - 3)
      p = insptr
      if Len != 255:
        insptr += Len
      else:
        insptr += wpeek(stream, insptr) + 2

      if insptr != p:
        sys.stdout.write(' {operands:')
        for i in range(p, insptr):
          sys.stdout.write(' %d' % peek(stream, i))
        sys.stdout.write('}')

      ncol += len(names[defName + 6])
      newLine = True
    else:
      s = '{unknown word %04X}' % word
      sys.stdout.write(s)
      ncol += len(s)
    if word == 0x04B6:
      return voc

    # Skip operand of brfalse, branch
    if word in {0x133C, 0x1332, 0x128D, 0x1276, 0x1271, 0x1288,
        0x1283}:
      insptr += 2
      newLine = True
    if word == 0x12A4:
      newLine = True

    if ncol >= 20:
      newLine = True

def Decompile(startAddr, stream):
  # Gather all words we can find
  actions = {}
  parents = {}  # their vocabularies
  names = {FORTHWORD: 'FORTH'}
  # Fetch all chains of words in the Forth vocabulary
  addr = startAddr
  while addr >= baseAddr:
    actions[addr + 1] = wpeek(stream, addr + 1)
    parents[addr + 1] = FORTHWORD
    names[addr + 1] = getName(stream, addr)
    addr = wpeek(stream, addr - 2)

  # Now search every other vocabulary
  while True:
    vocWords = set(j for j in actions if actions[j] == SETCONTEXT)
    newWords = {}
    for i in vocWords:
      addr = wpeek(stream, i + 2)  # grab link from the vocabulary word
      newWords.clear()
      while addr >= baseAddr and addr + 1 not in actions:
        if peek(stream, addr) != 0:
          newWords[addr + 1] = wpeek(stream, addr + 1)
          assert addr + 1 not in parents
          parents[addr + 1] = i
          names[addr + 1] = getName(stream, addr)
        addr = wpeek(stream, addr - 2)
      actions.update(newWords)
    if not newWords:
      break

  del vocWords
  wordList = sorted(actions)

  voc = FORTHWORD
  firstTime = True
  for nextWord in wordList:
    if not firstTime:
      sys.stdout.write('\n')
    firstTime = False
    needed = parents[nextWord]
    if needed != voc:
      # Find chain of contexts
      ctx = [needed]
      while needed != FORTHWORD:
        needed = parents[needed]
        ctx.append(needed)
      ctx.reverse()

#      try:
#        pos = ctx.index(voc)
#        ctx = ctx[pos + 1 : ]
#      except:
#        pass

      sys.stdout.write(' '.join(names[i] for i in ctx))
      voc = parents[nextWord]

      sys.stdout.write(' CONTEXT\n\n')

    action = actions[nextWord]
    if action == SETCONTEXT:
      sys.stdout.write('VOCABULARY %s' % names[nextWord])
    elif action == 0x0FF5:
      sys.stdout.write('%d CONSTANT %s' % (wpeek(stream, nextWord + 2),
        names[nextWord]))
    elif action == 0x0FF0:
      sys.stdout.write('%d VARIABLE %s' % (wpeek(stream, nextWord + 2),
        names[nextWord]))
    elif action == 0x0FEC:
      sys.stdout.write('CREATE %s' % names[nextWord])
      Len = wpeek(stream, nextWord - 5) - 7
      ncol = 25
      for i in range(nextWord + 2, nextWord + 2 + Len):
        if ncol >= 25:
          sys.stdout.write('\n')
          ncol = 0
        s = ' %d C,' % peek(stream, i)
        ncol += len(s)
        sys.stdout.write(s)
    elif action == 0x0EC3:
      sys.stdout.write(': %s' % names[nextWord])
      voc = listThread(stream, nextWord + 2, names, voc, parents)
    elif action == 0x1085:
      sys.stdout.write('DEFINER %s' % names[nextWord])
      voc = listThread(stream, nextWord + 4, names, voc, parents)
    elif action == 0x1108:
      operandLen = peek(stream, wpeek(stream, nextWord + 2) - 3)
      sys.stdout.write('%d COMPILER %s' % (-1 if operandLen == 255 else
        operandLen, names[nextWord]))
      voc = listThread(stream, nextWord + 4, names, voc, parents)
    elif action >= baseAddr and peek(stream, action) == 0xCD and wpeek(stream,
        action + 1) == 0x0FF0:
      # Find beginning of Definer word
      defName = action + wpeek(stream, action - 2) - 0x10001
      # defName points to the name field; skip it and point to the code field
      while peek(stream, defName) < 128:
        defName += 1
      sys.stdout.write('%s %s' % (names[defName + 6], names[nextWord]))
      Len = wpeek(stream, nextWord - 5) - 7
      ncol = 25
      for i in range(nextWord + 2, nextWord + 2 + Len):
        if ncol >= 25:
          sys.stdout.write('\n')
          ncol = 0
        s = ' %d C,' % peek(stream, i)
        ncol += len(s)
        sys.stdout.write(s)
    elif wpeek(stream, nextWord - 5) == 7 and wpeek(stream, nextWord
        ) < baseAddr:
      # Word with no parameter field - hack the code field
      sys.stdout.write('CREATE %s %d HERE 2- !' % (names[nextWord],
        wpeek(stream, nextWord)))
    else:
      # TODO: hack the whole word, including code and parameter field?
      sys.stdout.write("Can't list %s" % names[nextWord])

    if action != 0x1108:  # COMPILER is implicitly immediate
      endWord(stream, nextWord)
    else:
      sys.stdout.write('\n')


def main(argv):
  if len(argv) < 2:
    sys.stdout.write("No input .TAP file specified\n")
    return 1
  try:
    f = open(argv[1], 'rb')
  except:
    return error(2, "Error opening input file %s" % argv[1])
  kind = 'u'  # for unknown - otherwise 'j' for jupiter, 's' for spectrum
  try:
    while True:
      taphdr = bytearray(f.read(28))
      if not taphdr:
        break
      # Verification of header and autodetection of format
      if len(taphdr) >= 28 and taphdr[0] == 0x1B and taphdr[1:3] == b'\0\0':
        if kind == 'u':
          sys.stderr.write("Detected as a Spectrum-format .TAP file\n")
          kind = 's'
        elif kind != 's':
          return error(2, "Mixed Jupiter/Spectrum format .TAP file")
        taphdr = taphdr[3:] + f.read(1)  # we only read 28 bytes, fix that
      elif len(taphdr) >= 28 and taphdr[0] == 0x1A and taphdr[1] == 0:
        if kind == 'u':
          sys.stderr.write("Detected as a Jupiter-format .TAP file\n")
          kind = 'j'
        elif kind != 'j':
          return error(2, "Mixed Jupiter/Spectrum format .TAP file")
        taphdr = taphdr[2:]
      else:
        return error(2, "Not a .TAP file with Jupiter ACE data")

      # Read data part of the block
      blocklen = bytearray(f.read(2))
      if not blocklen:
        break
      if len(blocklen) == 1:
        return error(2, "Malformed .TAP file")
      blocklen = wordAt(blocklen, 0)
      data = bytearray(f.read(blocklen))
      if len(data) != blocklen:
        return error(2, "Truncated .TAP file")
      if taphdr[0] != 0:
        sys.stderr.write("Skipping Bytes block: %s\n"
          % taphdr[1:10].decode('latin1'))
        continue
      if kind == 's':
        if data[0] != 255:
          sys.stderr.write("Skipped non-standard block in Spectrum-format"
            " .TAP\n")
          continue
        data = data[1:]
      if wordAt(taphdr, 11) != blocklen - 1:
        return error(2, "Saved length inconsistent with actual block length"
          "\n")

      # Finally, we have the header and data
      LastWord = wordAt(taphdr, 15)
      CURRENT  = wordAt(taphdr, 17)
      CONTEXT  = wordAt(taphdr, 19)
      VOCLNK   = wordAt(taphdr, 21)
      STKBOT   = wordAt(taphdr, 23)

      Decompile(LastWord, data)
      sys.stdout.write('\nSAVE %s\n\n\n' % taphdr[1:11].decode('latin1').rstrip(' '))

  finally:
    f.close()

if __name__ == '__main__':
  sys.exit(main(sys.argv))
