from inc_noesis import *
import noesis
import rapi

def registerNoesisTypes():
	handle = noesis.register("Black Rock Shooter", ".mdl")
	noesis.setHandlerTypeCheck(handle, brsCheckType)
	noesis.setHandlerLoadModel(handle, brsLoadModel)
	#noesis.logPopup()
	return 1

#check if it's this type based on the data

def brsCheckType(data):
	bs = NoeBitStream(data)
	idMagic = bs.readInt()
	if idMagic != 0x4D534E49:
		return 0
	return 1       

#load the model
def brsLoadModel(data, mdlList):
	ctx = rapi.rpgCreateContext()
	bs = NoeBitStream(data)
	idMagic = bs.readInt()
	unk00, texCount, matCount, unk01, boneCount, count2, count3, unk02, meshCount = bs.read("8BI")
	matOffset, meshOffset, boneOffset, unk00Offset, meshGOffset, texOffset, vertOffset, size00, size01, texOffset2 = bs.read("10I")

	matInfo = []
	bs.seek(matOffset)
	for i in range(0, matCount):
		matInfo.append(bs.read("4B8I"))

	bs.seek(meshOffset)
	meshInfo = []
	meshInfoEnd = []
	for i in range(0, meshCount):
		info1 = bs.read("4BI")
		meshInfo.append(info1)
		meshInfoEnd.append(info1[4])
	meshInfoEnd.append(matInfo[0][4])
	#print("meshInfo", meshInfo)
	#print("meshInfoEnd", meshInfoEnd)
	
	bs.seek(boneOffset)
	boneSize = (unk00Offset - boneOffset) // boneCount
	#print(boneSize)
	boneList = []
	for i in range(0, boneCount):
		boneName = str(i)
		if boneSize == 0x20:
			info = bs.read("4b")
			pos = NoeVec3.fromBytes(bs.readBytes(12))
			quat = NoeQuat.fromBytes(bs.readBytes(16))
			quat[3] = -quat[3]
			boneMtx = quat.toMat43()
			boneMtx[3] = pos
			newBone = NoeBone(i, boneName, boneMtx, None, info[2])
			boneList.append(newBone)
		elif boneSize == 0x34:
			info = bs.read("4b")
			boneMtx = NoeMat43.fromBytes(bs.readBytes(48))
			newBone = NoeBone(i, boneName, boneMtx, None, info[2])
			boneList.append(newBone)
	if boneSize == 0x20:
		boneList = rapi.multiplyBones(boneList)

	bs.seek(unk00Offset)
	info3 = []
	for i in range(0, count2):
		info3.append(bs.read("2H2I"))
	#print("info3", info3)
	info31 = []
	for i in range(0, count2):
		bs.seek(info3[i][2], NOESEEK_ABS)
		info31.append(bs.read(str(info3[i][0]) + "I"))
	#print(info31)

	bs.seek(meshGOffset)

	texList = []
	matList = []
	bs.seek(texOffset)
	tdata = bs.readBytes(len(data) - texOffset)
	td = NoeBitStream(tdata)
	palArray = []
	for i in range(0, len(tdata) // 4):
		if td.readUInt() == 0x444D5450:
			palArray.append(td.tell() - 4)
	palArray.append(len(tdata))
	for i in range(0, texCount):
		td.seek(palArray[i])
		texInfo = td.read("4H4BI")
		texData = td.readBytes(texInfo[8] - 16)
		palData = []
		palSize = palArray[i + 1] - td.tell()
		palData = td.readBytes(palSize)
		pixWidth = 2 ** ((texInfo[3] & 0xF000) // 0x1000)
		pixHeight = (texInfo[3] & 0xFFF)
		#print(texInfo[5])
		if texInfo[4] == 5:
			pix = rapi.imageUntwiddlePSP(texData, pixWidth, pixHeight, 8)
			if texInfo[5] == 0x30:
				pix = rapi.imageDecodeRawPal(pix, palData, pixWidth, pixHeight, 8, "r8g8b8a8")
			elif texInfo[5] == 0:
				pix = rapi.imageDecodeRawPal(pix, palData, pixWidth, pixHeight, 8, "r5g6b5")
			elif texInfo[5] == 16:
				pix = rapi.imageDecodeRawPal(pix, palData, pixWidth, pixHeight, 8, "r5g5b5a1")
		elif texInfo[4] == 4:
			pix = rapi.imageUntwiddlePSP(texData, pixWidth, pixHeight, 4)
			if texInfo[5] == 0x30:
				pix = rapi.imageDecodeRawPal(pix, palData, pixWidth, pixHeight, 4, "r8g8b8a8")
			elif texInfo[5] == 0:
				pix = rapi.imageDecodeRawPal(pix, palData, pixWidth, pixHeight, 4, "r5g6b5")
			elif texInfo[5] == 16:
				pix = rapi.imageDecodeRawPal(pix, palData, pixWidth, pixHeight, 4, "r5g5b5a1")
		elif texInfo[4] == 3:
			pix = rapi.imageUntwiddlePSP(texData, pixWidth, pixHeight, 4)
		texList.append(NoeTexture(str(i), pixWidth, pixHeight, pix, noesis.NOESISTEX_RGBA32))

	matList = []
	for i in range(0, matCount):
		matName = str(i)
		material = NoeMaterial(matName, "")
		material.setTexture(str(matInfo[i][1]))
		material.setFlags(0, 1)
		matList.append(material)

	bs.seek(vertOffset)
	for a in range(0, meshCount):
		#rapi.rpgSetName(str(a))
		rapi.rpgSetName(str(meshInfo[a][2]))
		rapi.rpgSetMaterial(str(meshInfo[a][2]))
		bs.seek(meshInfo[a][4], NOESEEK_ABS)
		stripCount = (((meshInfoEnd[a + 1] - meshInfo[a][4]) - 12) // 16)
		#print(stripCount)
		#print(meshInfo[a][4])
		stripInfo = []
		tag1, tag2 = bs.read("2I")
		for b in range(0, stripCount):
			s = bs.read("I4BI4B"); s = [s[0], s[4], (s[5] & 0xFFFFFF), s[6], s[8], s[9]]
			stripInfo.append(s)
			#print(s)
		for b in range(0, stripCount):
			#print(stripInfo[b], str(a))
			bs.seek(meshInfo[a][4] + stripInfo[b][2], NOESEEK_ABS)
			pspVert = rapi.decodePSPVert(stripInfo[b][0])
			#print(bs.tell())
			#print(pspVert.weightType, pspVert.numWeights, pspVert.indexType)
			#print(pspVert.uvOfs, pspVert.uvType, pspVert.colorOfs, pspVert.colorType, pspVert.normalOfs, pspVert.normalType)
			if pspVert.uvType == 2:
				rapi.rpgSetUVScaleBias(NoeVec3 ((2.0, 2.0, 1.0)), NoeVec3 ((0.0, 0.0, 0.0)))
			elif pspVert.uvType == 3:
				rapi.rpgSetUVScaleBias(NoeVec3 ((1.0, 1.0, 1.0)), NoeVec3 ((0.0, 0.0, 0.0)))
			vertData = bs.readBytes(pspVert.vertexSize * stripInfo[b][3])
			if vertData is not None and len(vertData) > 0:
				pspVert.bindBuffers(vertData)
				rapi.rpgCommitTriangles(None, noesis.RPGEODATA_SHORT, stripInfo[b][3], noesis.RPGEO_TRIANGLE_STRIP, 1)
	rapi.rpgClearBufferBinds()
	try:
		mdl = rapi.rpgConstructModel()
	except:
		mdl = NoeModel()
	mdl.setModelMaterials(NoeModelMaterials(texList, matList))
	mdlList.append(mdl)
	mdl.setBones(boneList)
	return 1