import plib
from argparse import ArgumentParser
import sys


def parse_int(str):
    base = 10
    if len(str) >= 2 and str[0] == '0' and (str[1] == 'x' or str[1] == 'X'):
        base = 16

    return int(str, base)


def hexdump(words):
    for word in words:
        print("%04x " % word)


def select_default_port(brom: bool):
    import config
    try:
        if brom:
            return config.BR_DEV_PATH
        else:
            return config.PL_DEV_PATH
    except AttributeError:
        if brom:
            print('Define BR_DEV_PATH or specify on command line.')
        else:
            print('Define PL_DEV_PATH or specify on command line.')
        sys.exit(1)


def main():
    parser = ArgumentParser()
    parser.add_argument('-w', '--wait', action='store_true',
                        help='Wait for device')
    parser.add_argument("-v", "--verbose", action="store_true")
    parser.add_argument("--skip-handshake", action="store_true")
    parser.add_argument("-b", '--brom', action='store_true',
                        help='Use BROM mode')
    parser.add_argument("-p", "--port", type=str)
    parser.add_argument("--da-base-address", type=int, default=0x00200000)
    parser.add_argument("--da", type=str)
    parser.add_argument("--meta", action="store_true", help="Use META Mode")
    parser.add_argument("--identify", action="store_true")
    parser.add_argument("--unlock", action="store_true",
                        help="Try to disable signature verification (temporary) must be run from bootrom")
    parser.add_argument("--read32", type=str,
                        help="Read memory eg, 0x0,1")
    parser.add_argument("--write32", type=str,
                        help="Write memory eg, 0x0:0xff00ff3f,0x00000004,...")
    parser.add_argument('--bootcda', action='store_true')
    parser.add_argument('--test2', action='store_true')
    parser.add_argument('--crash-preloader', action='store_true', help="On nougat this can be used to enter BootROM mode, on oreo it enters demigod crash handler")
    args = parser.parse_args()

    if args.port:
        port = args.port
    else:
        port = select_default_port(args.brom)

    device = plib.Device(port, args.wait, brom_mode=args.brom)
    if not args.skip_handshake:
        if args.meta:
            device.meta_mode()
        device.handshake()

    if args.crash_preloader:
        device.crash_preloader()
        sys.exit(0)

    if args.unlock:
        device.unlock()

    if args.identify:
        device.identify()

    if args.read32:
        s = args.read32
        n = s.find(',')
        if n == -1:
            try:
                addr = parse_int(s)
                hexdump(device.read32(addr, 1))
            except ValueError:
                print("Invalid argument: %s" % s)
        else:
            s1 = s[:n]
            s2 = s[n+1:]
            try:
                addr = parse_int(s1)
                size = parse_int(s2)
                hexdump(device.read32(addr, size))
            except ValueError:
                print("Invalid argument: %s" % s)

    if args.write32:
        s = args.write32
        n = s.find(':')
        if n == -1:
            print("Invalid argument: %s" % s)
        else:
            addr = parse_int(s[:n])
            wordlist_str = s[n+1:].split(",")
            if len(wordlist_str) == 0 or (len(wordlist_str) == 1 and wordlist_str[0] == ''):
                print("Invalid argument: %s" % s)
            else:
                wordlist = []
                try:
                    for word_str in wordlist_str:
                        word = parse_int(word_str)
                        if word > 0xffffffff:
                            raise RuntimeError()
                        wordlist.append(word)
                    device.write32(addr, wordlist)
                except (ValueError, RuntimeError):
                    print("Invalid word: %s" % word_str)

    if args.da:
        with open(args.da, "rb") as f:
            payload = f.read()
            device.upload_da(args.da_base_address, payload)
            f.close()
    if args.bootcda:
        device.bootcda()
    if args.test2:
        device.test2()


if __name__ == "__main__":
    main()
