#!/usr/bin/env python3
import platform
import argparse
import os
import sys

import criu_coredump

PLATFORMS = ["aarch64", "armv7l", "x86_64"]


def coredump(opts):
    generator = criu_coredump.coredump_generator()
    cores = generator(os.path.realpath(opts['in']))
    for pid in cores:
        if opts['pid'] and pid != opts['pid']:
            continue
        with open(os.path.realpath(opts['out']) + "/core." + str(pid), 'wb+') as f:
            cores[pid].write(f)


def main():
    desc = 'CRIU core dump'
    parser = argparse.ArgumentParser(description=desc,
                                     formatter_class=argparse.RawTextHelpFormatter)

    parser.add_argument('-i',
                        '--in',
                        default='.',
                        help='directory where to get images from')
    parser.add_argument('-p',
                        '--pid',
                        type=int,
                        help='generate coredump for specific pid(all pids py default)')
    parser.add_argument('-o',
                        '--out',
                        default='.',
                        help='directory to write coredumps to')

    opts = vars(parser.parse_args())

    if platform.machine() not in PLATFORMS:
        print("ERROR: %s is only supported on: %s" % (sys.argv[0], ', '.join(PLATFORMS)))
        sys.exit(1)

    try:
        coredump(opts)
    except SystemExit as error:
        print('ERROR: %s' % error)
        print('Exiting')
        sys.exit(1)


if __name__ == '__main__':
    main()
