import struct
import re
import tempfile
from collections import OrderedDict
from collections import namedtuple
import os
import shutil

def enum(name, values):
  enums = dict(zip(values, range(len(values))))
  enums['values'] = values
  return type(name, (), enums)

TypeId = enum('TypeId', [
  'RKSERIAL_TYPE_INTEGER',
  'RKSERIAL_TYPE_BOOLEAN',
  'RKSERIAL_TYPE_STRING',
  'RKSERIAL_TYPE_FLOAT',
  'RKSERIAL_TYPE_VECTOR4',
  'RKSERIAL_TYPE_MATRIX44',
  'RKSERIAL_TYPE_POINTER',
  'RKSERIAL_TYPE_ARRAY',
  'RKSERIAL_TYPE_VECTOR2',
  'RKSERIAL_TYPE_CUSTOMDATA',
  'RKSERIAL_TYPE_STRINGPAIR'
])

RKSERIAL_RKAT_ANIMATION   = 0
RKSERIAL_RKAT_FONT        = 1
RKSERIAL_RKAT_LOCALIZATION = 2
RKSERIAL_RKAT_MODEL       = 3
RKSERIAL_RKAT_SCRIPT      = 4
RKSERIAL_RKAT_SHADER      = 5
RKSERIAL_RKAT_SOUND       = 6
RKSERIAL_RKAT_SWF         = 7
RKSERIAL_RKAT_TEXTURE     = 8
RKSERIAL_RKAT_VIDEO       = 9
RKSERIAL_RKAT_XML         = 10
RKSERIAL_RKAT_ANIM_EVENT  = 11
RKSERIAL_RKAT_GLOBAL_EVENT = 12
RKSERIAL_RKAT_CSHARP      = 13
RKSERIAL_RKAT_FBX         = 14
RKSERIAL_RKAT_JSON        = 15


RKSERIAL_FLAG_GLOBAL_LINK = 1<<1
RKSERIAL_FLAG_OBJECT_LINK = 1<<2
RKSERIAL_FLAG_ASSET_LINK  = 1<<3
RKSERIAL_FLAG_TEMPLATE    = 1<<10
RKSERIAL_FLAG_RENDERABLE  = 1<<27

def _objectLink(flags):
  return ((flags & RKSERIAL_FLAG_TEMPLATE) != 0
      or (flags & RKSERIAL_FLAG_GLOBAL_LINK) != 0
      or (flags & RKSERIAL_FLAG_OBJECT_LINK) != 0)

def _assetLink(flags):
  return ((flags & RKSERIAL_FLAG_ASSET_LINK) and (flags >> 32) != RKSERIAL_RKAT_SOUND)

def _assetType(flags):
  if (flags & RKSERIAL_FLAG_ASSET_LINK) == 0: return ''
  t = flags >> 32
  if t == RKSERIAL_RKAT_ANIMATION:    return 'Animations'
  if t == RKSERIAL_RKAT_FONT:         return 'Fonts'
  if t == RKSERIAL_RKAT_LOCALIZATION: return 'Localization'
  if t == RKSERIAL_RKAT_MODEL:        return 'Models'
  if t == RKSERIAL_RKAT_SCRIPT:       return 'Scripts'
  if t == RKSERIAL_RKAT_SHADER:       return 'Shaders'
  if t == RKSERIAL_RKAT_SOUND:        return 'Sounds'
  if t == RKSERIAL_RKAT_SWF:          return 'Swf'
  if t == RKSERIAL_RKAT_TEXTURE:      return 'Textures'
  if t == RKSERIAL_RKAT_VIDEO:        return 'Video'
  if t == RKSERIAL_RKAT_XML:          return 'Xml'
  if t == RKSERIAL_RKAT_ANIM_EVENT:   return 'AnimationEvents'
  if t == RKSERIAL_RKAT_GLOBAL_EVENT: return 'GlobalEvents'
  if t == RKSERIAL_RKAT_CSHARP:       return 'CSharp'
  if t == RKSERIAL_RKAT_FBX:          return 'Fbx'
  if t == RKSERIAL_RKAT_JSON:         return 'Json'
  return 'Assets'

def _extractStringFromOffset(data, offset):
  for i in range(offset, len(data)):
    if data[i] == chr(0):
      return data[offset:i]
  return data[offset:]

def _setStringInData(value, data, header, offsetData):
  memberSize, typeId, flags = struct.unpack('<IIQ', header)
  valueStringOffset, memberNameOffset, editorNameOffset, editorDescrOffset, editorCategoryOffset, classNameOffset = struct.unpack('<HHHHHH', offsetData)
  oldValue = _extractStringFromOffset(data, valueStringOffset)
  dataLength = len(oldValue)
  sizeDifference = len(value) - dataLength

  if memberNameOffset > valueStringOffset: memberNameOffset += sizeDifference
  if editorNameOffset > valueStringOffset: editorNameOffset += sizeDifference
  if editorDescrOffset > valueStringOffset: editorDescrOffset += sizeDifference
  if editorCategoryOffset > valueStringOffset: editorCategoryOffset += sizeDifference
  if classNameOffset > valueStringOffset: classNameOffset += sizeDifference
  memberSize += sizeDifference

  extraData = data[dataLength:]
  data = value + extraData
    
  endData = max(valueStringOffset, memberNameOffset, editorNameOffset, editorDescrOffset, editorCategoryOffset, classNameOffset)
  endData += len(_extractStringFromOffset(data, endData)) + 1
  headerSize = len(header)
  offsetSize = len(offsetData)
  chunkSize = headerSize + offsetSize + endData
  paddedSize = ((chunkSize - 1) / 8 + 1) * 8
  paddedDifference = paddedSize - (len(data) + headerSize + offsetSize)
  memberSize += paddedDifference
  if paddedDifference > 0:
    for i in range (0, paddedDifference):
      data += '\0'
  elif paddedDifference < 0:
    data = data[:paddedDifference]
      
  header = struct.pack('<IIQ', memberSize, typeId, flags)
  offsetData = struct.pack('<HHHHHH', valueStringOffset, memberNameOffset, editorNameOffset, editorDescrOffset, editorCategoryOffset, classNameOffset)

  return data, header, offsetData


class EOFException(Exception):
  pass

class PointerRef():
  def __init__(self, ref):
    self.ref = ref

  def __str__(self):
    return 'Pointer("%s")' % self.ref
    
class AssetRef():
  def __init__(self, ref, kind):
    self.ref = ref
    self.kind = kind

  def __str__(self):
    return '%s("%s")' % (self.kind, self.ref)

class StringPair():
  def __init__(self, arg1, arg2 = None):
    if arg2 == None:
      self.first = _extractStringFromOffset(arg1, 0)
      self.second = _extractStringFromOffset(arg1, len(self.first)+1)
    else:
      self.first = arg1
      self.second = arg2

  def __str__(self):
    return '"%s" : "%s"' % (self.first, self.second)

def _changeChunkValue(f, value, header, offsetData, data):
  memberSize, typeId, flags = struct.unpack('<IIQ', header)
  valueStringOffset, memberNameOffset, editorNameOffset, editorDescrOffset, editorCategoryOffset, classNameOffset = struct.unpack('<HHHHHH', offsetData)

  if   typeId == TypeId.RKSERIAL_TYPE_INTEGER:
    data = struct.pack('<i', value) + data[4:]
  elif typeId == TypeId.RKSERIAL_TYPE_BOOLEAN:
    data = struct.pack('<?', value) + data[1:]
  elif typeId == TypeId.RKSERIAL_TYPE_FLOAT:
    data = struct.pack('<f', value) + data[4:]
  elif typeId == TypeId.RKSERIAL_TYPE_VECTOR4:
    dataValue = struct.pack('<ffff', value[0], value[1], value[2], value[3])
    data = dataValue + data[16:]
  elif typeId == TypeId.RKSERIAL_TYPE_MATRIX44:
    dataValue = struct.pack('<ffffffffffffffff', data, valueStringOffset, value[0], value[1], value[2], value[3],
                value[4], value[5], value[6], value[7],
                value[8], value[9], value[10], value[11],
                value[12], value[13], value[14], value[15])
    data = dataValue + data[64:]
  elif typeId == TypeId.RKSERIAL_TYPE_STRING:
     (data, header, offsetData) =_setStringInData(value, data, header, offsetData)
  elif typeId == TypeId.RKSERIAL_TYPE_POINTER:
     (data, header, offsetData) =_setStringInData(value.ref, data, header, offsetData)
  elif typeId == TypeId.RKSERIAL_TYPE_STRINGPAIR:
     (data, header, offsetData) =_setStringInData(value.first + '\0' + value.second, data, header, offsetData)
  else:
    assert False, 'This data type is not implemented yet: typeId == %s, file == %s' % (typeId, f)

  _writeChunkInternal(f, header, offsetData, data)

def _writeChunkInternal(f, header, offsetData, data):
  f.write(header)
  f.write(offsetData)
  f.write(data)


def _parseChunkInternal(f):

  (header, offsetData, data) = _parseChunkRaw(f)
  
  memberSize, typeId, flags = struct.unpack('<IIQ', header)
  valueStringOffset, memberNameOffset, editorNameOffset, editorDescrOffset, editorCategoryOffset, classNameOffset = struct.unpack('<HHHHHH', offsetData)

  memberName = _extractStringFromOffset(data, memberNameOffset)
  className = _extractStringFromOffset(data, classNameOffset)

  if   typeId == TypeId.RKSERIAL_TYPE_INTEGER:
    value = struct.unpack('<i', data[:4])[0]
  elif typeId == TypeId.RKSERIAL_TYPE_BOOLEAN:
    value = data[0] != chr(0)
  elif typeId == TypeId.RKSERIAL_TYPE_STRING:
    if _objectLink(flags):
      value = PointerRef(_extractStringFromOffset(data, 0))
    elif _assetLink(flags):
      value = AssetRef(_extractStringFromOffset(data, 0), _assetType(flags)) 
    else:
      value = _extractStringFromOffset(data, 0)
  elif typeId == TypeId.RKSERIAL_TYPE_FLOAT:
    value = struct.unpack('<f', data[:4])[0]
  elif typeId == TypeId.RKSERIAL_TYPE_VECTOR4:
    value = struct.unpack('<ffff', data[:4*4])
  elif typeId == TypeId.RKSERIAL_TYPE_MATRIX44:
    value = struct.unpack('<ffffffffffffffff', data[:4*16])
  elif typeId == TypeId.RKSERIAL_TYPE_POINTER:
    value = PointerRef(_extractStringFromOffset(data, 0))
  elif typeId == TypeId.RKSERIAL_TYPE_STRINGPAIR:
    value = StringPair(data)
  else:
    assert False, 'This data type is not implemented yet: typeId == %s, file == %s' % (typeId, f)

  # print 'class:%s, member:%s, value:%s, flags:%x, type:%s' % (className, memberName, value, flags, TypeId.values[typeId])

  return className, memberName, value, flags

def _parseChunkRaw(f):
  headerSize = 4 + 4 + 8
  offsetsSize = 2 * 6
  
  header = f.read(headerSize)
  offsetData = f.read(offsetsSize)
  if len(header) == 0: # eof happened
    raise EOFException()
  
  memberSize, typeId, flags = struct.unpack('<IIQ', header)
  remainingSize = memberSize - headerSize - offsetsSize
  data = f.read(remainingSize)

  return header, offsetData, data

def _parseChunk(f):
  className, memberName, value, flags = _parseChunkInternal(f)
  return className, memberName, value

def _quoteString(s):
  if s.__class__ == str:
    return "'%s'" % s
  else:
    return str(s)

class SerializableDef(object):
  def __init__(self):
    self.typeName = None
    self.attrs = OrderedDict()

  def __str__(self):
    return '%s {%s}' % (self.typeName, self.attrs)

castRE  = re.compile(r'\([^)]*\)(.*)$')
arrayRE = re.compile(r'([^\[]*)\[([^\]]*)\]')

def _extractArrayIdx(name):
  arrayMatch = arrayRE.match(name)
  if arrayMatch:
    return arrayMatch.group(1), arrayMatch.group(2)
  else:
    return name, None

def _addMember(o, className, memberName, value):
  # strip casts like (int&)member
  
  castMatch = castRE.match(memberName)
  if castMatch:
    memberName = castMatch.group(1)

  if className == 'RegionConditional' or className == 'RegionEnd' : return

  #print className, memberName, value
  memberParts = memberName.split('.')
  if len(memberParts) == 1:
    if o.typeName == None:
      o.typeName = className

    baseName, baseIdx = _extractArrayIdx(memberParts[0])
    if baseIdx != None:
      if not baseName in o.attrs or type(o.attrs[baseName]) == int: # arrays have it's size stored under array's name
        o.attrs[baseName] = OrderedDict()
      else:
        assert o.attrs[baseName].__class__ == OrderedDict
      o.attrs[baseName][baseIdx] = value
    else:
      o.attrs[memberName] = value
  else:
    subName = '.'.join(memberParts[1:])
    baseName, baseIdx = _extractArrayIdx(memberParts[0])
    if baseIdx != None:
      if not baseName in o.attrs or type(o.attrs[baseName]) == int: # arrays have it's size stored under array's name
        o.attrs[baseName] = OrderedDict()
      else:
        assert o.attrs[baseName].__class__ == OrderedDict

      if not baseIdx in o.attrs[baseName]:
        o.attrs[baseName][baseIdx] = SerializableDef()
      else:
        assert o.attrs[baseName][baseIdx].__class__ == SerializableDef

      _addMember(o.attrs[baseName][baseIdx], className, subName, value)
    else:
      if not baseName in o.attrs:
        o.attrs[baseName] = SerializableDef()
      else:
        assert o.attrs[baseName].__class__ == SerializableDef
      _addMember(o.attrs[baseName], className, subName, value)

def _convertDefIntoInstance(o):
  if o.__class__ == SerializableDef:
    Type = namedtuple(o.typeName or 'UnknownType', o.attrs.keys())

    convertedAttrs = OrderedDict()
    for k, v in o.attrs.items():
      convertedAttrs[k] = _convertDefIntoInstance(v)

    return Type(**convertedAttrs)

  elif o.__class__ == OrderedDict:
    try:
      res = []
      for i in range(len(o)):
        assert str(i) in o, "Invalid dictionary, doesn't containt %s key: %s" % (i, o)
        value = _convertDefIntoInstance(o[str(i)])
        res.append(value)
      return res
    except AssertionError:
      return o

  else:
    return o

def GetReferencedTemplates(f):
  if isinstance(f, basestring):
    #print 'Parsing %s' % f
    f = open(f, 'rb')

  templates = []
  try:
    while True:
      className, memberName, value, valueFlags = _parseChunkInternal(f)
      if value and ((valueFlags & RKSERIAL_FLAG_TEMPLATE) != 0 or memberName == 'm_RKObjectParentName'):
        templates.append(value)
  except EOFException:
    pass
  except AssertionError as e:
    print e

  return templates

def GetValuesMatchingFlag(f, flags):
  if isinstance(f, basestring):
    #print 'Parsing %s' % f
    f = open(f, 'rb')

  values = []
  try:
    while True:
      className, memberName, value, valueFlags = _parseChunkInternal(f)
      if value and (valueFlags & flags) != 0:
        values.append(value)
  except EOFException:
    pass
  except AssertionError as e:
    print e

  return values

# accepts either stream or string
def parseObject(f):
  if isinstance(f, basestring):
    #print 'Parsing %s' % f
    f = open(f, 'rb')

  o = SerializableDef()

  try:
    while True:
      className, memberName, value = _parseChunk(f)
      #print className, memberName, value
      _addMember(o, className, memberName, value)
  except EOFException:
    pass
  
  return _convertDefIntoInstance(o)
  
def parseObjectFriendlyName(f):
	friendlyName = f
	if isinstance(f, basestring):
		#print 'Parsing %s' % f
		f = open(f, 'rb')
		
	o = SerializableDef()
	
	try:
		className, memberName, friendlyName = _parseChunk(f)
		#print className, memberName, friendlyName
	except EOFException:
		print 'Failed to parse friendly name for rki %s' % friendlyName
		pass
		
	return friendlyName

def getFieldFromObject(f, field):
  if isinstance(f, basestring):
    f = open(f, 'rb')

  try:
    while True:
      className, memberName, value = _parseChunk(f)
      if memberName == field:
        return value
  except EOFException:
    pass

  return None

def setFieldInObject(f, field, value):
  friendlyName = f
  fOutLocation = f + ".temp"
  fOut = fOutLocation

  if isinstance(f, basestring):
    f = open(f, 'rb')
    fOut = open(fOut, 'wb')

  o = SerializableDef()

  try:
    while True:
      header, offsetData, data = _parseChunkRaw(f)
    
      valueStringOffset, memberNameOffset, editorNameOffset, editorDescrOffset, editorCategoryOffset, classNameOffset = struct.unpack('<HHHHHH', offsetData)

      memberName = _extractStringFromOffset(data, memberNameOffset)
      castIndex = memberName.rfind(')')
      if castIndex != -1: memberName = memberName[castIndex + 1:]
      if memberName == field:
        memberSize, typeId, flags = struct.unpack('<IIQ', header)
        
        convertedValue = 0
        if typeId == TypeId.RKSERIAL_TYPE_INTEGER:
          convertedValue = int(value[0])
        elif typeId == TypeId.RKSERIAL_TYPE_BOOLEAN:
          convertedValue = (value[0].lower() == 'true')
        elif typeId == TypeId.RKSERIAL_TYPE_FLOAT:
          convertedValue = float(value[0])
        elif typeId == TypeId.RKSERIAL_TYPE_VECTOR4:
          convertedValue = [float(value[0]), float(value[1]), float(value[2]), float(value[3])]
        elif typeId == TypeId.RKSERIAL_TYPE_MATRIX44:
          convertedValue = [float(value[0]), float(value[1]), float(value[2]), float(value[3]),
            float(value[4]), float(value[5]), float(value[6]), float(value[7]),
            float(value[8]), float(value[9]), float(value[10]), float(value[11]),
            float(value[12]), float(value[13]), float(value[14]), float(value[15])]
        elif typeId == TypeId.RKSERIAL_TYPE_STRING:
          if _objectLink(flags):
            convertedValue = PointerRef(value[0])
          elif _assetLink(flags):
            convertedValue = AssetRef(value[0], _assetType(flags)) 
          else:
            convertedValue = value[0]
        elif typeId == TypeId.RKSERIAL_TYPE_POINTER:
          convertedValue = PointerRef(value[0])
        elif typeId == TypeId.RKSERIAL_TYPE_STRINGPAIR:
          convertedValue = StringPair(value[0], value[1])
        else:
          assert False, 'This data type is not implemented yet: typeId == %s, file == %s' % (typeId, f)

        _changeChunkValue(fOut, convertedValue, header, offsetData, data)
      else:
        _writeChunkInternal(fOut, header, offsetData, data)
  except EOFException:
    pass

  f.close()
  fOut.close()
  shutil.copyfile(fOutLocation, friendlyName)  
  os.remove(fOutLocation)

def _tabLines(lines):
  return ['\t' + x for x in lines]

def _formatObjectLines(obj):
  if isinstance(obj, basestring):
    return ['"%s"' % obj]

  if isinstance(obj, list):
    resLines = ['[']
    for item in obj:
      resLines += _tabLines(_formatObjectLines(item))
    resLines.append(']')
    return resLines

  if isinstance(obj, dict):
    resLines = ['{']
    for key, value in obj.items():
      valueLines = _formatObjectLines(value)
      resLines.append('\t%s = %s' % (key, valueLines[0]))
      resLines += _tabLines(valueLines[1:])
    resLines.append('}')
    return resLines

  if not isinstance(obj, tuple) or not hasattr(obj, '__dict__'):
    return [str(obj).replace(',', ',\r\n')]

  resLines = ['%s {' % type(obj).__name__]
  for attrName, attrValue in obj.__dict__.items():
    attrLines = _formatObjectLines(attrValue)
    resLines.append('\t%s = %s' % (attrName, attrLines[0]))
    resLines += _tabLines(attrLines[1:])
  resLines.append('}')

  return resLines

def formatObject(obj):
  return '\n'.join(_formatObjectLines(obj))

def diffObjects(obj1, obj2, diffCmd):
  f1h, f1Path = tempfile.mkstemp('_%s' % type(obj1).__name__)
  with os.fdopen(f1h, 'wb') as f1:
    f1.write(formatObject(obj1))

  f2h, f2Path = tempfile.mkstemp('_%s' % type(obj2).__name__)
  with os.fdopen(f2h, 'wb') as f2:
    f2.write(formatObject(obj2))

  os.system('"%s" %s %s' % (diffCmd, f1Path, f2Path))

if __name__ == '__main__':
  import sys

  if len(sys.argv[1:]) == 1:
    obj = parseObject(sys.argv[1])
    print formatObject(obj)
  elif len(sys.argv[1:]) == 2:
    obj1 = parseObject(sys.argv[1])
    obj2 = parseObject(sys.argv[2])
    diffObjects(obj1, obj2, 'diff')
  elif len(sys.argv[1:]) == 3:
    obj1 = parseObject(sys.argv[1])
    obj2 = parseObject(sys.argv[2])
    diffCmd = sys.argv[3]
    diffObjects(obj1, obj2, diffCmd)
  else:
    print 'Invalid argument count, expected 1, 2 or 3 args'
    sys.exit(1)