from trgformat import *
import struct

# arguments
executing_player = Player2
str_section_length = 2408
desired_unit_count = 5000
counter_unit = 0
tmp_unit = 1


# impl

counter_memory = PlayerUnit2Memory(executing_player, counter_unit)
tmp_memory = PlayerUnit2Memory(executing_player, tmp_unit)

str_section_pointer = 0x005993D4
mrgn_section_start = 0x0058DC60 + 20 * 6

offset_unitnode_empty_start = 0x00628438
offset_triggernode_start = 0x0051a280 + 12 * executing_player + 8

unit_extendn = desired_unit_count - 1700

# offset for payloads
offset_setter    = mrgn_section_start + 0
offset_advancer  = mrgn_section_start + 2408
offset_finalizer = mrgn_section_start + 300

# offset for setter
"""
setter code

*x = z
*y = w
x += 336
y += 336
z += 336
w += 336
counter++
"""

offset_x = offset_setter + 8 + 320 + 32*0 + 16
offset_y = offset_setter + 8 + 320 + 32*1 + 16
offset_z = offset_setter + 8 + 320 + 32*0 + 20
offset_w = offset_setter + 8 + 320 + 32*1 + 20

print(offset_x, offset_y, offset_z, offset_w)

# offset for advancer
"""
advancer code

if counter == unit_extendn:
	advancer.next = finalizer
"""

offset_advancer_next = offset_advancer + 4

# offset for finalizer
"""
finalizer code

finalizer.next                  = trg_start_old
*trgstart                       = trg_start_old

*empty_start                    = myunit[0]
myunit[0].prev                  = 0
myunit[unit_extendn - 1].next   = empty_start_old
empty_start_old.prev            = myunit[unit_extendn - 1]

"""

offset_finalizer_next = offset_finalizer + 4

offset_death1 = offset_finalizer + 8 + 320
offset_death2 = offset_finalizer + 8 + 320 + 32 * 1
offset_death3 = offset_finalizer + 8 + 320 + 32 * 2
offset_death4 = offset_finalizer + 8 + 320 + 32 * 3
offset_death5 = offset_finalizer + 8 + 320 + 32 * 4
offset_death6 = offset_finalizer + 8 + 320 + 32 * 5

offset_death1_player = offset_death1 + 16
offset_death2_player = offset_death2 + 16
offset_death3_player = offset_death3 + 16
offset_death4_player = offset_death4 + 16
offset_death5_player = offset_death5 + 16
offset_death6_player = offset_death6 + 16

offset_death1_number = offset_death1 + 20
offset_death2_number = offset_death2 + 20
offset_death3_number = offset_death3 + 20
offset_death4_number = offset_death4 + 20
offset_death5_number = offset_death5 + 20
offset_death6_number = offset_death6 + 20


# setter
triggers = []

triggers.append(Trigger(
	players = [executing_player],
	conditions = [],
	actions = [
		SetMemory(0, SetTo, -336),
		SetMemory(4, SetTo, 336),
		SetMemory(offset_x, Add, 336 // 4),
		SetMemory(offset_y, Add, 336 // 4),
		SetMemory(offset_z, Add, 336),
		SetMemory(offset_w, Add, 336),
		SetMemory(counter_memory, Add, 1)
	],
	preservetrigger = True
))

f = open('setter.trg', 'wb')
f.write(struct.pack('L', offset_advancer))
f.write(struct.pack('L', offset_advancer))
WriteTrg(f, triggers)
f.close()

# advancer
triggers = []

print(offset_advancer_next)
triggers.append(Trigger(
	players = [executing_player],
	conditions = [
		Memory(counter_memory, Exactly, unit_extendn)
	],
	actions = [
		SetMemory(offset_advancer_next, SetTo, offset_finalizer)
	]
))

f = open('advancer.trg', 'wb')
f.write(struct.pack('L', offset_setter))
f.write(struct.pack('L', offset_setter))
WriteTrg(f, triggers)
f.close()

# finalizer
triggers = []
triggers.append(Trigger(
	players = [executing_player],
	conditions = [],
	actions = [
		SetMemory(offset_finalizer_next, SetTo, 0),
		SetMemory(offset_triggernode_start, SetTo, 0),
		SetMemory(offset_unitnode_empty_start, SetTo, 0),
		SetMemory(0, SetTo, 0),
		SetMemory(0, SetTo, 0),
		SetMemory(0, SetTo, 0)
	]
))

f = open('finalizer.trg', 'wb')
f.write(struct.pack('L', offset_advancer))
f.write(struct.pack('L', 0))
WriteTrg(f, triggers)
f.close()

# payload is created. now do some work.

# what vector does:
"""
# do things with setter
x += (stroffset + str_section_length) // 4
y += (stroffset + str_section_length) // 4
z += (stroffset + str_section_length) // 4
w += (stroffset + str_section_length) // 4


# do things with finalizer
/trgstart_old    = trgstart
/firstunit       = stroffset + str_section_length
/firstunit_prev  = (stroffset + str_section_length + 4) // 4
/finalunit_next  = (stroffset + str_section_length + (336 * unit_extendn - 1) + 4) // 4
/empty_start_old = empty_start
/empty_old_prev  = empty_start
/
/death1_number = trgstart_old
/death2_number = trgstart_old
/death3_number = firstunit
/death4_player = finalunit_next
/death4_number = empty_start_old
/death5_player = firstunit.prev
/death6_player = empty_old_prev
/death6_number = firstunit

death1_number = trgstart
death2_number = trgstart
death3_number = stroffset + str_section_length
death4_player = (stroffset + str_section_length + (336 * unit_extendn - 1) + 4) // 4
death4_number = empty_start
death5_player = (stroffset + str_section_length) // 4
death6_player = empty_start // 4
death6_number = stroffset + str_section_length + (336 * unit_extendn - 1)
"""

triggers = []

# step 1
"""
# do things with setter & finalizer
x += (stroffset + str_section_length) // 4
y += (stroffset + str_section_length) // 4
z += (stroffset + str_section_length) // 4
w += (stroffset + str_section_length) // 4

death3_number = stroffset + str_section_length
death4_player = (stroffset + str_section_length + (336 * (unit_extendn - 1)) + 4) // 4
death5_player = (stroffset + str_section_length) // 4
death6_number = stroffset + str_section_length
"""

for i in range(23, 1, -1):
	triggers.append(Trigger(
		players = [executing_player],
		conditions = [
			Memory(counter_memory, Exactly, 0),
			Memory(str_section_pointer, AtLeast, 336 * 2**i)
		],

		actions = [
			SetMemory(str_section_pointer, Subtract, 336 * 2**i),
			SetMemory(tmp_memory, Add, 336 * 2**i),

			SetMemory(offset_x, Add, 336 * 2**(i-2)),
			SetMemory(offset_y, Add, 336 * 2**(i-2)),
			SetMemory(offset_z, Add, 336 * 2**i),
			SetMemory(offset_w, Add, 336 * 2**i),

			
			SetMemory(offset_death3_number, Add, 336 * 2**i),
			SetMemory(offset_death4_player, Add, 336 * 2**(i-2)),
			SetMemory(offset_death5_player, Add, 336 * 2**(i-2)),
			SetMemory(offset_death6_number, Add, 336 * 2**i)
		]
	))


for i in range(23, -1, -1):
	triggers.append(Trigger(
		players = [executing_player],
		conditions = [
			Memory(counter_memory, Exactly, 0),
			Memory(tmp_memory, AtLeast, 336 * 2**i)
		],

		actions = [
			SetMemory(tmp_memory, Subtract, 336 * 2**i),
			SetMemory(str_section_pointer, Add, 336 * 2**i)
		]
	))

triggers.append(Trigger(
	players = [executing_player],
	conditions = [
		Memory(counter_memory, Exactly, 0),
	],

	actions = [
		SetMemory(offset_x, Add, str_section_length // 4 + 336 // 4),
		SetMemory(offset_y, Add, str_section_length // 4 + 336 // 4),
		SetMemory(offset_z, Add, str_section_length + 336),
		SetMemory(offset_w, Add, str_section_length + 336),
		
		SetMemory(offset_death3_number, Add, 336 + str_section_length),
		SetMemory(offset_death4_player, Add, 336 // 4 + str_section_length // 4 + (84 * (unit_extendn - 1)) + 1),
		SetMemory(offset_death5_player, Add, 336 // 4 + str_section_length // 4),
		SetMemory(offset_death6_number, Add, 336 + str_section_length + (336 * (unit_extendn - 1)))
	]
))

"""
death1_number = trgstart
death2_number = trgstart
"""

for i in range(31, 1, -1):
	triggers.append(Trigger(
		players = [executing_player],
		conditions = [
			Memory(counter_memory, Exactly, 0),
			Memory(offset_triggernode_start, AtLeast, 2**i)
		],

		actions = [
			SetMemory(offset_triggernode_start, Subtract, 2**i),
			SetMemory(tmp_memory, Add, 2**i),

			SetMemory(offset_death1_number, Add, 2**i),
			SetMemory(offset_death2_number, Add, 2**i)
		]
	))

for i in range(31, 1, -1):
	triggers.append(Trigger(
		players = [executing_player],
		conditions = [
			Memory(counter_memory, Exactly, 0),
			Memory(tmp_memory, AtLeast, 2**i)
		],

		actions = [
			SetMemory(tmp_memory, Subtract, 2**i),
			SetMemory(offset_triggernode_start, Add, 2**i)
		]
	))


"""
death4_number = empty_start
death6_player = empty_start // 4
"""

for i in range(31, 1, -1):
	triggers.append(Trigger(
		players = [executing_player],
		conditions = [
			Memory(counter_memory, Exactly, 0),
			Memory(offset_unitnode_empty_start, AtLeast, 2**i)
		],

		actions = [
			SetMemory(offset_unitnode_empty_start, Subtract, 2**i),
			SetMemory(tmp_memory, Add, 2**i),

			SetMemory(offset_death4_number, Add, 2**i),
			SetMemory(offset_death6_player, Add, 2**(i-2))
		]
	))

for i in range(31, 1, -1):
	triggers.append(Trigger(
		players = [executing_player],
		conditions = [
			Memory(counter_memory, Exactly, 0),
			Memory(tmp_memory, AtLeast, 2**i)
		],

		actions = [
			SetMemory(tmp_memory, Subtract, 2**i),
			SetMemory(offset_unitnode_empty_start, Add, 2**i)
		]
	))

# last stage
triggers.append(Trigger(
	players = [executing_player],
	conditions = [
		Memory(counter_memory, Exactly, 0)
	],
	actions = [
		SetMemory(offset_triggernode_start, SetTo, offset_setter),
		SetMemory(0x006283F0, Add, -unit_extendn),
		Wait(0)
	]
))


print(len(triggers))
WriteTrg('vector.trg', triggers)


# make payload
data_setter = open('setter.trg', 'rb').read()
data_advcer = open('advancer.trg', 'rb').read()
data_finalr = open('finalizer.trg', 'rb').read()

data_payload = bytearray(4816)

data_payload[0:2408] = data_setter[:]
data_payload[2408:4816] = data_advcer[:]
data_payload[300:308] = data_finalr[:8]
data_payload[308+320:308+320+6*32] = data_finalr[8+320:8+320+6*32]
data_payload[308+320+2048:308+320+2048+32] = data_finalr[8+320+2048:2408]

open('payload', 'wb').write(data_payload)

data_payload = open('payload', 'rb').read()
data_vector = open('vector.trg', 'rb').read()
data_scenario = bytearray(open('scenario.chk', 'rb').read())

print(len(data_vector))
data_scenario[0x136252:0x136252+4816] = data_payload[:]
data_scenario[0x13B5D0:0x13B5D0+403200] = data_vector[:]
open('scenario.chk', 'wb').write(data_scenario)

p = Memory2Player(0x006283F0)
u = p // 12
p = p % 12
print(u, p)