#!/usr/bin/env python3
import sys
import os
p = os.path.join
pUp = os.path.dirname
s = False
if getattr(sys, 'frozen', False) and hasattr(sys, '_MEIPASS'):
	s = os.path.realpath(sys.executable)
else:
	s = os.path.realpath(__file__)
sp = pUp(s)

## Script start
import shutil
import platform
ops = platform.system().lower()

import subprocess
def procConfirm(proc,cmd = None,success = None):
	if cmd == None: cmd = lastCmd
	if success == None: success = [0]
	rtn = proc.wait()
	if not rtn in success: raise subprocess.CalledProcessError(rtn,cmd,None)

lastCmd = []
def popen(cmd,*args,**kwargs):
	global lastCmd
	lastCmd = cmd
	return subprocess.Popen(cmd,*args,**kwargs)

def call(cmd,*args,**kwargs):
	proc = popen(cmd,*args,**kwargs)
	procConfirm(proc)

def callTool(cmd,*args,**kwargs):
	if ops != "windows":
		if os.path.isfile(p(sp,"bin",ops,cmd[0])):
			cmd[0] = p(sp,"bin",ops,cmd[0])
		else:	
			cmd[0] = p(sp,"bin","windows",cmd[0] + ".exe")
			cmd.insert(0,"wine")
	else:
		cmd[0] = p(sp,"bin","windows",cmd[0] + ".exe")
	
	call(cmd,*args,**kwargs)

def wp(*args):
	pth = p(*args)
	if ops == "windows": return pth
	pth = pth.replace(os.path.sep,"\\")
	if pth[0] == "\\":
		pth = "Z:" +pth
		return pth
	return pth

def fchunkUnpack(fch,fchOut,pf = "pc"):
	fchName = fch.rsplit(os.path.sep,1)[-1]
	
	if os.path.isdir(fchOut): shutil.rmtree(fchOut)
	if pf == "x360":
		fchDecomp = p(pUp(fchOut),fchName + ".decomp")
		callTool([p("xbcompression","xbdecompress"),"/F",wp(fch),wp(fchDecomp)])
		callTool([p("LuaSF","LuaSF"),"--big",wp(fchDecomp),wp(fchOut)])
		os.remove(fchDecomp)
		return
	
	if pf == "ps3":
		callTool([p("LuaSF","LuaSF"),"--big",wp(fch),wp(fchOut)])
		return
	
	callTool([p("LuaSF","LuaSF"),wp(fch),wp(fchOut)])

def fchunkPack(fch,fchOut,pf = "pc"):
	fchName = fch.rsplit(os.path.sep,1)[-1] + ".fchunk"
	fchOutDir = pUp(fchOut)
	fchOutLsfName = p(fchOutDir,fchName.rsplit(".fchunk",1)[0])
	
	if not os.path.isdir(fchOutDir): os.makedirs(fchOutDir)
	if pf == "x360":
		callTool([p("LuaSF","LuaSF"),"--big","--compile",wp(fch),wp(fchOutDir)],cwd=p(sp,"bin","windows","LuaSF")) # ugly hack
		os.replace(fchOutLsfName,fchOut + ".decomp")
		fchunkFixBig(fchOut + ".decomp")
		callTool([p("xbcompression","xbcompress"),"/F",wp(fchOut + ".decomp"),wp(fchOut)])
		os.remove(fchOut + ".decomp")
		return
	
	if pf == "ps3":
		callTool([p("LuaSF","LuaSF"),"--big","--compile",wp(fch),wp(fchOutDir)])
		fchunkFixBig(fchOutLsfName)
	
	if pf == "pc":
		callTool([p("LuaSF","LuaSF"),"--compile",wp(fch),wp(fchOutDir)])
	
	if fchOutLsfName == fchOut: return
	os.replace(fchOutLsfName,fchOut)
	return

def fchunkFixBig(fchunk): # Fix SCRH (Script Package Header) and SCRC (Script Package Lookup) -- Thanks miru97 for the help!
	fh = open(fchunk,"r+b")
	
	# Find pointers
	pointers = {}
	fh.seek(0x810)
	headerName = fh.read(4).decode("ascii")
	pointers[headerName] = [0,0,0]
	pointers[headerName][0] = int.from_bytes(fh.read(4),"little") # address
	pointers[headerName][1] = int.from_bytes(fh.read(4),"little") # unknown
	pointers[headerName][2] = int.from_bytes(fh.read(4),"little") # size
	
	headerName = fh.read(4).decode("ascii")
	pointers[headerName] = [0,0,0]
	pointers[headerName][0] = int.from_bytes(fh.read(4),"little") # address
	pointers[headerName][1] = int.from_bytes(fh.read(4),"little") # unknown
	pointers[headerName][2] = int.from_bytes(fh.read(4),"little") # size
	
	#headerName = fh.read(4).decode("ascii")
	#pointers[headerName] = [0,0,0]
	#pointers[headerName][0] = int.from_bytes(fh.read(4),"big") # address
	#pointers[headerName][1] = int.from_bytes(fh.read(4),"big") # unknown
	#pointers[headerName][2] = int.from_bytes(fh.read(4),"big") # size
	
	# normalize addresses
	pointers["SCRC"][0] = pointers["SCRH"][0] + pointers["SCRC"][0]
	#pointers["SCRS"][0] = pointers["SCRC"][0] + pointers["SCRS"][0] # TODO: check if correct
	
	for space in pointers:
		cursor = pointers[space][0]
		end = cursor + pointers[space][2]
		fh.seek(cursor)
		while cursor < end:
			# invert bytes
			data = fh.read(1)
			data = fh.read(1) + data
			data = fh.read(1) + data
			data = fh.read(1) + data
			fh.seek(cursor)
			fh.write(data)
			cursor += 4
	
	fh.close()

def createPatchFromDir(dir1,dir2,output):
	proc = popen(["diff","-crB",dir1,dir2],stdout=subprocess.PIPE)
	with open(output,"w",encoding="utf-8") as fh:
		text = proc.stdout.read().decode("utf-8")
		text = text.replace(dir1,"a")
		text = text.replace(dir2,"b")
		fh.write(text)
	procConfirm(proc,success=[0,1])

def copytree(dir1,dir2):
	for root,dirs,files in os.walk(dir1):
		for file in dirs:
			ffile = p(root,file)
			lfile = ffile.replace(dir1 + os.path.sep,"",1)
			nfile = p(dir2,lfile)
			os.makedirs(nfile,exist_ok=True)
		
		for file in files:
			ffile = p(root,file)
			lfile = ffile.replace(dir1 + os.path.sep,"",1)
			nfile = p(dir2,lfile)
			if os.path.isfile(nfile): os.remove(nfile)
			shutil.copyfile(ffile,nfile)

def main():
	if os.path.isdir(p(sp,"tmp")): shutil.rmtree(p(sp,"tmp"))
	os.makedirs(p(sp,"tmp"))
	
	fchunksMod = [
		"ScriptsActionActivities.fchunk",
		"ScriptsChallenges.fchunk",
		"ScriptsChapter1.fchunk",
		"ScriptsChapter2.fchunk",
		"ScriptsChapter3.fchunk",
		"ScriptsChapter4.fchunk",
		"ScriptsChapter5.fchunk",
		"ScriptsChapter6.fchunk",
		"ScriptsChapter7.fchunk"
	]
	
	fchunksExtra = [
		"ScriptsFelonyReworkAudio.fchunk",
		"ScriptsStimulusPackage.fchunk"
	]
	
	if sys.argv[1] == "makepatch":
		if os.path.isdir(p(sp,"patches")): shutil.rmtree(p(sp,"patches"))
		os.makedirs(p(sp,"patches"))
		
		for fchunk in fchunksMod:
			print(">> " +fchunk+ " ...")
			fchunkUnpack(p(sp,"game",fchunk),p(sp,"tmp","game",fchunk))
			fchunkUnpack(p(sp,"felony_rework","media",fchunk),p(sp,"tmp","felony_rework",fchunk))
			createPatchFromDir(
				p(sp,"tmp","game",fchunk),
				p(sp,"tmp","felony_rework",fchunk),
				p(sp,"patches",fchunk + ".patch")
			)
			
		print(">> ScriptsLuaScripts.fchunk ...")
		fchunkUnpack(p(sp,"game","ScriptsLuaScripts.fchunk"),p(sp,"tmp","game","ScriptsLuaScripts.fchunk"))
		createPatchFromDir(
			p(sp,"tmp","game","ScriptsLuaScripts.fchunk"),
			p(sp,"felony_rework","Common","LuaScripts"),
			p(sp,"patches","ScriptsLuaScripts.fchunk.patch")
		)
	
	if sys.argv[1] == "apply":
		pf = sys.argv[2]
		game = sys.argv[3]
		os.makedirs(p(sp,"tmp","game"))
		if os.path.isdir(p(sp,"out")): shutil.rmtree(p(sp,"out"))
		
		fchunksMod.append("ScriptsLuaScripts.fchunk")
		
		for fchunk in fchunksMod:
			print(">> " +fchunk+ " ...")
			fchunkUnpack(p(game,"media",fchunk),p(sp,"tmp","game",fchunk),pf)
			proc = popen(["patch","-p1"],stdin=subprocess.PIPE,cwd=p(sp,"tmp","game",fchunk))
			with open(p(sp,"patches",fchunk + ".patch"),"rb") as fh:
				proc.stdin.write(fh.read())
			proc.stdin.close()
			procConfirm(proc)
			if fchunk == "ScriptsLuaScripts.fchunk":
				call(["patch",p(sp,"tmp","game",fchunk,"loadfiles.lua"),p(sp,"patches-extra","modloader.patch")])
			fchunkPack(p(sp,"tmp","game",fchunk),p(sp,"out","media",fchunk),pf)
		
		for fchunk in fchunksExtra:
			print(">> " +fchunk+ " ...")
			fchunkUnpack(p(sp,"felony_rework","media",fchunk),p(sp,"tmp","game",fchunk),"pc")
			fchunkPack(p(sp,"tmp","game",fchunk),p(sp,"out","media",fchunk),pf)
		
		copytree(p(sp,"files","xbox_modloader"),p(sp,"out"))
		copytree(p(sp,"felony_rework","Common","UserLuaScripts"),p(sp,"out","media","UserLuaScripts"))
		copytree(p(sp,"felony_rework","Common","fmv"),p(sp,"out","media","fmv"))
main()