#!/usr/bin/env python3
# Copyright lowRISC contributors.
# Licensed under the Apache License, Version 2.0, see LICENSE for details.
# SPDX-License-Identifier: Apache-2.0

'''Wrapper script to run the OTBN random instruction generator'''

import argparse
import json
import os
import random
import sys
from typing import Optional, cast

from shared.insn_yaml import InsnsFile, load_insns_yaml

# Ensure that the OTBN utils directory is on sys.path. This means that RIG code
# can import modules like "shared.foo" and get the OTBN shared code.
sys.path.append(os.path.dirname(__file__))

from rig.init_data import InitData  # noqa: E402
from rig.rig import gen_program, snippets_to_program  # noqa: E402
from rig.snippet import Snippet  # noqa: E402


def get_insns_file() -> Optional[InsnsFile]:
    '''Load up insns.yml'''
    try:
        return load_insns_yaml()
    except RuntimeError as err:
        print(err, file=sys.stderr)
        return None


def gen_main(args: argparse.Namespace) -> int:
    '''Entry point for the gen subcommand'''
    random.seed(args.seed)

    insns_file = get_insns_file()
    if insns_file is None:
        return 1

    # Run the generator
    init_data, snippets = gen_program(args.start_addr, args.size, insns_file)

    # Write out the data and snippets to a JSON file
    ser_data = init_data.as_json()
    ser_snippets = [snippet.to_json() for snippet in snippets]
    ser = [ser_data, ser_snippets]
    try:
        if args.output == '-':
            json.dump(ser, sys.stdout)
            # Add a newline at end of output: json.dump doesn't, and it makes a
            # bit of a mess of some consoles.
            sys.stdout.write('\n')
        else:
            with open(args.output, 'w') as out_file:
                json.dump(ser, out_file)
    except OSError as err:
        print('Failed to open json output file {!r}: {}.'
              .format(args.output, err),
              file=sys.stderr)
        return 1

    return 0


def asm_main(args: argparse.Namespace) -> int:
    '''Entry point for the asm subcommand'''

    insns_file = get_insns_file()
    if insns_file is None:
        return 1

    # Load JSON
    try:
        json_data = json.load(args.snippets)
    except json.JSONDecodeError as err:
        print('Snippets file at {!r} is not valid JSON: {}.'
              .format(args.snippets.name, err),
              file=sys.stderr)
        return 1

    # Parse these to proper init data and snippet objects.
    try:
        if not (isinstance(json_data, list) and len(json_data) == 2):
            raise ValueError('Top-level structure should be a length 2 list.')

        json_init_data, json_snippets = json_data
        init_data = InitData.read(json_init_data)
        snippets = [Snippet.from_json(insns_file, idx, x)
                    for idx, x in enumerate(json_snippets)]
    except ValueError as err:
        print('Failed to parse snippets from {!r}: {}'
              .format(args.snippets, err),
              file=sys.stderr)
        return 1

    program = snippets_to_program(snippets)
    dsegs = init_data.as_segs()

    # Dump the assembly output, and the linker script too if we're writing to
    # something other than stdout.
    if args.output is None or args.output == '-':
        program.dump_asm(sys.stdout, dsegs)
    else:
        try:
            asm_path = args.output + '.s'
            with open(asm_path, 'w') as out_file:
                program.dump_asm(out_file, dsegs)
        except OSError as err:
            print('Failed to open asm output file {!r}: {}.'
                  .format(args.output, err),
                  file=sys.stderr)
            return 1

        try:
            ld_path = args.output + '.ld'
            with open(ld_path, 'w') as out_file:
                program.dump_linker_script(out_file, dsegs)
        except OSError as err:
            print('Failed to open ld script output file {!r}: {}.'
                  .format(ld_path, err),
                  file=sys.stderr)
            return 1

    return 0


def main() -> int:
    parser = argparse.ArgumentParser()
    subparsers = parser.add_subparsers(dest='cmd')
    subparsers.required = True

    gen = subparsers.add_parser('gen', help='Generate a random program')
    asm = subparsers.add_parser('asm', help='Convert snippets to assembly')

    gen.add_argument('--output', '-o',
                     default='-',
                     help=("Path for JSON output of generated snippets. The "
                           "special path '-' (which is the default) means to "
                           "write to stdout"))
    gen.add_argument('--seed', type=int, default=0,
                     help='Random seed. Defaults to 0.')
    gen.add_argument('--size', type=int, default=100,
                     help=('Max number of instructions in stream. '
                           'Defaults to 100.'))
    gen.add_argument('--start-addr', type=int, default=0,
                     help='Reset address. Defaults to 0.')
    gen.set_defaults(func=gen_main)

    asm.add_argument('--output', '-o',
                     metavar='out',
                     help=('Base path for output filenames. Will generate '
                           'out.s with an assembly listing and out.ld with '
                           'a linker script. If this is not supplied, the '
                           'assembly (but no linker script) will be dumped '
                           'to stdout.'))
    asm.add_argument('snippets', metavar='path.json',
                     type=argparse.FileType('r'), nargs='?', default=sys.stdin,
                     help=('A JSON file of snippets, as generated by '
                           'otbn-rig gen.'))
    asm.set_defaults(func=asm_main)

    args = parser.parse_args()
    return cast(int, args.func(args))


if __name__ == '__main__':
    sys.exit(main())