import struct
import zlib
import sys
import os

# ===========================================
# ============= Tree Flattening =============
# ===========================================

def flattenName(parent, child):
  if(parent == ""): return child
  else: return "%s.%s" % (parent, child)

def flattenRecurse(vals, name = ""):
  global flattenReturn
  if(type(vals) == type([])):
    if(type(vals[0]) == type("")):
      if(type(vals[1]) != type([])):
        result = [flattenName(name, vals[0]), vals[1]]
        flattenReturn.append(result)
      else:
        for val in vals[1:]:
          flattenRecurse(val, flattenName(name, vals[0]))
    else:
      for val in vals:
        flattenRecurse(val, name)

def flatten(vals):
  global flattenReturn
  flattenReturn = []
  flattenRecurse(vals)
  return flattenReturn

# ===========================================
# ============== Final Output ===============
# ===========================================

# findImage should work for image*'s, but the image*'s
#  aren't always valid, or possibly point to a parent
#  class's image...

#def findImage(name, vals):
#  name = name.replace("/", ".")
#  name = "%s.image" % name
#  for val in vals:
#    if(val[0] == name): return val[1]
#
#  name = "%s*" % name
#  for val in vals:
#    if(val[0] == name): return findImage(val[1], vals)
#
#  print "Could not find image %s" % name
#  sys.exit(-1)

def endMatches(str, values):
  for s in values:
    if(str[-len(s):] == s): return 1
  return 0

def textOut(vals):
  for val in vals:
    print "%s: %s" % (val[0], val[1])
  # Code used for testing findImage... which doesn't work (see above note)
  #for val in vals:
  #  if(endMatches(val[0], ["image*"])):
  #    print "Searching for %s: %s" % (val[0], val[1])
  #    print "%s: %s)" % (val[0][:-1], findImage(val[1], vals), val[1])
  #  else:
  #    print "%s: %s" % (val[0], val[1])

def htmlOut(vals):
  print "<table>"
  for val in vals:
    if(endMatches(val[0], ["image"])):
      print "<tr><td>%s</td><td><img src='%s'></td></tr>" % (val[0], val[1])
    elif(endMatches(val[0], ["mp3"])):
      print "<tr><td>%s</td><td><a href='%s'>mp3</a></td></tr>" % (val[0], val[1])
    else:
      print "<tr><td>%s</td><td>%s</td></tr>" % (val[0], val[1])
  print "</table>"

def safe_mkdir(dirName):
  try:
    os.mkdir(dirName)
  except OSError, (errno, strerror):
    if(errno != 17):
      print "%d: %s" % (errno, strerror)
      sys.exit(-1)

# ===========================================
# ============ PNG/Image Routines ===========
# ===========================================

def pack_chunk(tag,data):
  to_check = tag + data
  return struct.pack("!I",len(data)) + \
         to_check + \
         struct.pack("!I",zlib.crc32(to_check))

def dumpToPNG(width, height, data):
  raw_data = ""
  for y in range(height):
    raw_data += chr(0)
    raw_data += data[y*width*4:(y+1)*width*4]

  return struct.pack("8B", 137, 80, 78, 71, 13, 10, 26, 10) + \
         pack_chunk('IHDR',
           struct.pack("!IIBBBBB",width,height,8,6,0,0,0)) + \
         pack_chunk('IDAT', zlib.compress(raw_data,9)) + \
         pack_chunk('IEND', '')

def image16to32(width, height, data):
  datalen = len(data)
  buf = ""
  for i in range(width*height):
    # Not all images decompress correectly. This makes up for it.
    if(i > datalen/2 - 1):
      p = [0,0]
    elif(width*height*2 == len(data)/2):
     p = struct.unpack("BBBB", data[i*4:i*4+4])
    else:
     p = struct.unpack("BB", data[i*2:i*2+2])
    if(width*height*2 == len(data)/2):
      a = p[3]
      r = p[2]
      g = p[1]
      b = p[0]
    else:
      a = p[1] / 16 ; a = a*16 + a
      r = p[1] % 16 ; r = r*16 + r
      g = p[0] / 16 ; g = g*16 + g
      b = p[0] % 16 ; b = b*16 + b

    if(a >  255): a = 255
    if(r >  255): r = 255
    if(g >  255): g = 255
    if(b >  255): b = 255

    buf += struct.pack("BBBB", r, g, b, a)
  return buf

# ===========================================
# =========== Abort/Debug Routines ==========
# ===========================================

def printItem(name, value):
  print "%-30s: %s" % (name, value)

def check(flag, val, f, msg="Default Error"):
  if(flag != val):
    item  = "Unknown: %s" % repr(flag) ; value = "Expecting: %s" % repr(val)
    printItem(item, value)
    printItem("Error", msg)
    printItem("FileLocation", repr(f.tell()))
    sys.exit(-1)

def checkA(flag, vals, f, msg="Default Error"):
  if(not flag in vals):
    item  = "Unknown: %s" % repr(flag) ; value = "Expecting: %s" % repr(vals)
    printItem(item, value)
    printItem("Error", msg)
    printItem("FileLocation", repr(f.tell()))
    sys.exit(-1)

# ===========================================
# ============ Basic File Reading ===========
# ===========================================

def rU8 (f): return struct.unpack('B', f.read(1))[0]
def rS8 (f): return struct.unpack('b', f.read(1))[0]
def rU16(f): return struct.unpack('H', f.read(2))[0]
def rU32(f): return struct.unpack('I', f.read(4))[0]
def rS32(f): return struct.unpack('i', f.read(4))[0]
def rU64(f): return struct.unpack('Q', f.read(8))[0]
def rF32(f): return struct.unpack('f', f.read(4))[0]
def rF64(f): return struct.unpack('d', f.read(8))[0]

def rPackNum(f):
  size = rS8(f)
  if(size == -128): size = rS32(f)
  return size

# ===========================================
# ============ String Decryption ============
# ===========================================

def hexStr(s):
  str = ''
  for c in s:
    str += "\\x%02x" % ord(c)
  return str

def xor  (a, p): return  ((a & ~p) ^ (~a & p)) & 0xff
def xor16(a, p): return  ((a & ~p) ^ (~a & p)) & 0xffff

def transStr(str):
  s = ''
  p = 0xAA
  for cc in str:
    a = ord(cc)
    s += struct.pack("B", xor(a,p) )
    p += 1
  return s

def transStr16(encStr):
  s = ''
  p = 0xAAAA
  for i in range(len(encStr)/2):
    a = struct.unpack("h", encStr[i*2:i*2+2])[0]
    #s += struct.pack("h", xor16(a,p) )
    s += "&#" + str(xor16(a,p))
    p += 1
  return s

# ===========================================
# ========= External File Extraction ========
# ============= (MP3s and PNGs) =============
# ===========================================

def dumpImage(f, imageName):
  global dumpDir
  global dumpImages

  width  = rPackNum(f)
  height = rPackNum(f)
  val   = rPackNum(f)    ; checkA(val, [0, 1, 2, 10, 513], f)
  val2  = rU8(f)         ; checkA(val2, [0, 2, 4], f)
  nulls = f.read(4)      ; check(nulls, "\x00\x00\x00\x00", f, "Dump Image")
  bufSize = rU32(f) - 1
  flag = rU8(f)          ; check(flag,  0x00, f);

  if(not dumpImages):
    f.seek(bufSize, 1)
  else:
    obj = zlib.decompressobj()
    buf = obj.decompress(f.read(bufSize))
    g = open(dumpDir + imageName, "wb")
    g.write(dumpToPNG(width, height, image16to32(width, height, buf)))
    g.close()

def dumpMp3(f, size, endOfBlock, imageName):
  startOfMP3 = endOfBlock - size
  remaining = startOfMP3 - f.tell()
  unknown = f.read(remaining)
  g = open(dumpDir + imageName, "wb")
  g.write(f.read(size))
  g.close()
  #mp3count += 1
  f.seek(endOfBlock)

# ===========================================
# ================= Header ==================
# ===========================================

def extractHeader(f):
  v = {}
  v["FileIdentity"] = f.read(4)
  v["FileSize"]     = rU64(f)
  v["FileStart"]    = rU32(f)
  v["Copyright"]    = f.read(v["FileStart"] - f.tell())
  return v

# ===========================================
# ========== Directory / File Tree ==========
# ===========================================

def extractDirectory(f, baseName, blockOffset):
  global eFiles

  dirs  = []
  entryCount = rPackNum(f)
  for j in range(entryCount):
    filetype = rU8(f)    ; checkA(filetype, [0x02, 0x03, 0x04], f)
    if(filetype in [0x03, 0x04]):
      filename = rUStr(f)
    if(filetype in [0x02]):
      filenameloc  = rU32(f)
      filename = rUStrAt(f, filenameloc, blockOffset)
    filesize = rPackNum(f)
    unknown1 = rPackNum(f)
    unknown2 = rU32(f)

    v = {"Name": baseName + "/" + filename,
         "Size": filesize,
         "Unk1": unknown1,
         "Unk2": unknown2}

    if(filetype in [0x03]):
      dirs.append(v)
    if(filetype in [0x02, 0x04]):
      eFiles.append(v)

  for dir in dirs:
    extractDirectory(f, dir["Name"], blockOffset)

def computeFileOffsets(fileBase):
  global eFiles

  totalOffset = 0
  for file in eFiles:
    file["Offset"] = fileBase + totalOffset
    totalOffset += file["Size"]

def extractDirectories(f, blockOffset):
  global eFiles
  eFiles = []
  extractDirectory(f, "", blockOffset)
  computeFileOffsets(f.tell())
  return eFiles

def extractDirTree(f, blockOffset):
  v = {}
  v["Unknown"] = rU16(f)
  v["Files"]   = extractDirectories(f, blockOffset)
  return v

# ===========================================
# ============= Everything else =============
# ===========================================

def subBlock(f, count, endOfBlock = 0):
  result = []
  if(count == 0): return '-NONE-'
  for i in range(count):
    flag = rU8(f)  ; check(flag, 0x1b, f)
    type = rU32(f)
    val1 = rPackNum(f) ; val2 = 0
    if(f.tell() != endOfBlock):
      val2 = rPackNum(f)
    result.append([val1, val2])
  return result

def rUStr(f): return extract00type(f)
def rUStrAt(f, loc, baseOffset): return extract00typeAt(f, loc, baseOffset)

previousStrings = {}
parsingFileTree = 0
def extract00type(f):
  size = rS8(f)
  if(size == 0):
    return ''
  elif(size < 0):
    if(size == -128): size = rU32(f)
    else:             size = -size
    result = transStr(f.read(size))

    # Keep a list of previous translations
    global previousStrings
    previousStrings[result] = 0

    return result
  elif(size > 0):
    if(size == 127): size = rU32(f)

    global parsingFileTree
    if(parsingFileTree):
      subSize = rS8(f) ; subCount = 1
      if(subSize < 0): subSize = -subSize
      else:            subSize = rU32(f) ; subCount += 4
      subStr  = transStr(f.read(size*2 - subCount))

      for previousString in previousStrings.keys():
        if((len(previousString) == subSize) and 
           (previousString[0:size*2-1] == subStr)):
          return previousString
      print "Can't find String '%s' at %d" % (subStr, f.tell())
      sys.exit(-1)
    else:
      if(size == 1):
        val = rU16(f)
        val = xor16(val, 0x558a)
        if(val < 256):
          return chr(val)
        else:
          f.seek(-2, 1)
      
      result = f.read(size*2)
      
      result = transStr16(result)
      return result

def extract00typeAt(f, loc, baseOffset):
  pos = f.tell()
  f.seek(loc + baseOffset, 0)
  value = extract00type(f)
  f.seek(pos)
  return value

def extract01type(f, baseOffset):
  loc = rU32(f)
  return extract00typeAt(f, loc, baseOffset)

def extract02type(f):
  return rU16(f)

def extract03type(f):
  return rPackNum(f)

def extract04type(f):
  type = rU8(f) ; checkA(type, [0x00, 0x80], f, "0x04 type")
  if   (type == 0x80): return rF32(f)
  elif (type == 0x00): return 0

def extract05type(f):
  return rF64(f)

def extract08type(f, name, baseOffset):
  return dumpBlock(f, name, baseOffset)

def extract09type(f, name, baseOffset):
  size = rU32(f)
  loc = f.tell()
  endOfBlock = f.tell() + size
  result = dumpBlock(f, name, baseOffset, f.tell() + size)
  if(f.tell() != endOfBlock):
    print "Incomplete Block [%s] at offset %d, size: %d, loc: %d" % (name, loc, size, f.tell())
    f.seek(endOfBlock, 0)
   # sys.exit(-1)
  return result

def extract1btype(f, name, baseOffset, endOfBlock):
  flag  = rU16(f)
  flagb = rU16(f)

  locb  = f.tell()
  
  val1 = rPackNum(f) ; val2 = 0
  if(f.tell() != endOfBlock):
    val2 = rPackNum(f)

  if(f.tell() == endOfBlock):
    return [['x', val1], ['y', val2]]

  f.seek(locb, 0)
  val1 = rU8(f)
  if(val1 == 0x00):
    val2 = rU8(f) ; subt = 0
    if(val2 not in [0x00, 0x80] and flag != 0x01):
      nil = rU16(f)  ; check(nil,  0x00, f, "Inner 1b type1")

    if(val2 == 0x80):
      size = rU32(f)
      imgname = name.replace('/', '.') + ".mp3"
      dumpMp3(f, size, endOfBlock, imgname)
      return ["mp3", name]
    else:
      vals = []
      if(not (val1 == 0x00 and val2 == 0x00) or (flag == 0x01)):
        descCount = rPackNum(f)
        for j in range(descCount):
          qname = dumpBlock(f, name, baseOffset)
          val = dumpEntryValue(f, name + "/" + str(qname), baseOffset)
          vals.append([qname, val])

      if(f.tell() == endOfBlock):
        if(vals == []): return '-NONE-'
        else:           return vals
      else:
        imgname = name.replace('/', '.') + ".png"
        dumpImage(f, imgname)
        vals.append(['image', imgname])
        return vals
  else:
    blockval = val1
    return subBlock(f, val1, endOfBlock)

def computeName(image, name):
  name = name[0:name.rfind("/")+1]
  while(image[0:3] == "../"):
    name = name[0:-1]
    name = name[0:name.rfind("/")+1]
    image = image[3:]
  return name + image

def extract73type(f, name, baseOffset, endOfBlock):
  iname = rUStr(f)
  if(iname == "Shape2D#Convex2D"):
    count = rPackNum(f)
    return subBlock(f, count, endOfBlock)
  if(iname == "Shape2D#Vector2D"):
    val1 = rPackNum(f)
    val2 = rPackNum(f)
    return [["x", val1], ["y", val2]]
  if(iname == "Canvas"):
    space = rU16(f)  ; checkA(space, [0, 256], f)
    if(space == 256):
      nil = rU16(f)
      count = rU8(f)
      for i in range(count):
        qname = dumpBlock(f, name, baseOffset)
        val = dumpEntryValue(f, name + "/" + str(qname), baseOffset)
    imgname = name.replace('/', '.') + ".png"
    dumpImage(f, imgname)
    return ['image', imgname]
  if(iname == "UOL"):
    nil = rU8(f)
    image = dumpBlock(f, name, baseOffset)
    cName = computeName(image, name)
    return ['image*', cName]
  if(iname == "Sound_DX8"):
    nil = rU8(f)
    imgname = name.replace('/', '.') + ".mp3"
    size = rPackNum(f)
    dumpMp3(f, size, endOfBlock, imgname)
    return ["mp3", name]
# ---------------------------------------------

def dumpBlock(f, name, baseOffset, endOfBlock = 0):
  global vals
  type = rU8(f)
  if(name != "coconut/fontTime" and name != "town/cloud/0/0/foothold" and name != "Subway/wall/0/0/foothold"):
    if   (type == 0x00): return extract00type(f)
    elif (type == 0x01): return extract01type(f, baseOffset)
    elif (type == 0x1b): return extract1btype(f, name, baseOffset, endOfBlock)
    elif (type == 0x73): return extract73type(f, name, baseOffset, endOfBlock)

def dumpEntryValue(f, name, baseOffset):
  global vals
  type = rU8(f)

  if   (type == 0x00): return "-NIL-"
  elif (type == 0x02): return extract02type(f)
  elif (type == 0x03): return extract03type(f)
  elif (type == 0x04): return extract04type(f)
  elif (type == 0x05): return extract05type(f)
  elif (type == 0x08): return extract08type(f, name, baseOffset)
  elif (type == 0x09): return extract09type(f, name, baseOffset)

def dumpFile(f, sname, baseOffset):
  global vals
  vals = {}

  flag = rU8(f)      ; check(flag, 0x73, f)
  imgName = rUStr(f) ; check(imgName, "Property", f)
  space = rU8(f)     ; check(space, 0x00, f, "File Outer Space")
  space = rU8(f)     ; check(space, 0x00, f, "File Outer Space 2")
  entries = rPackNum(f)
  
  vals = []
  for i in range(entries):
    name  = dumpBlock(f, sname, baseOffset)
    value = dumpEntryValue(f, str(name), baseOffset)
    vals.append([name, value])
  return vals


def dump():
  global dumpImages, dumpHTML
  global dumpDir, dumpDirBase
  f = open("data.wz", "rb")

  if(dumpImages or dumpHTML):
    safe_mkdir(dumpDirBase)
    origstdout = sys.stdout

  header = extractHeader(f)
  DirTreeOffset = f.tell()

  global parsingFileTree
  parsingFileTree = 1
  dirTree = extractDirTree(f, DirTreeOffset)
  parsingFileTree = 0

  for i in range(len(dirTree["Files"])):
    file = dirTree["Files"][i]
    f.seek(file["Offset"], 0)

    if(dumpImages or dumpHTML):
      dumpDir = dumpDirBase + file["Name"].replace("/", ".") + "/"
      safe_mkdir(dumpDir)

    print "%4d: %s, %d, %d" % (i, repr(file["Name"])[1:-1], file["Size"], file["Offset"])

    vals = dumpFile(f, file["Name"], file["Offset"])
    vals = flatten(vals)
    if(dumpHTML):
      g = open(dumpDir + "index.html", "w")
      sys.stdout = g

      htmlOut(vals)

      g.close
      sys.stdout = origstdout
    else:
      textOut(vals)

dumpImages = 1
dumpHTML = 1

dumpDirBase = "dump/"
dumpDir     = dumpDirBase

dump()