#!/usr/bin/env python3
# ------------------------------------------------------------------------ #
#     (C) Copyright 2017-2025 Barcelona Supercomputing Center              #
#                             Centro Nacional de Supercomputacion          #
#                                                                          #
#     This file is part of OmpSs@FPGA toolchain.                           #
#                                                                          #
#     This code is free software; you can redistribute it and/or modify    #
#     it under the terms of the GNU Lesser General Public License as       #
#     published by the Free Software Foundation; either version 3 of       #
#     the License, or (at your option) any later version.                  #
#                                                                          #
#     OmpSs@FPGA toolchain is distributed in the hope that it will be      #
#     useful, but WITHOUT ANY WARRANTY; without even the implied           #
#     warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.     #
#     See the GNU Lesser General Public License for more details.          #
#                                                                          #
#     You should have received a copy of the GNU Lesser General Public     #
#     License along with this code. If not, see <www.gnu.org/licenses/>.   #
# ------------------------------------------------------------------------ #

import argparse
import importlib
import json
import math
import os
import re
import sys

from ait.frontend.config import LONG_VERSION, MIN_WRAPPER_VERSION
from ait.frontend.utils import backends, decimalFromHumanReadable, \
    decimalToHumanReadable, getNumJobs, msg

for backend in backends:
    importlib.import_module(f'ait.backend.{backend}.info')


# Custom argparse type representing a power of 2 int
class IntPowerType:
    def __init__(self, imin=None, imax=None):
        self.imin = imin
        self.imax = imax

    def __call__(self, arg):
        try:
            value = int(arg)
        except ValueError:
            raise self.exception()
        if (self.imin is not None and value < self.imin) or (self.imax is not None and value > self.imax):
            raise self.exception()
        elif value & (value - 1):
            raise self.exception(value)
        return value

    def exception(self, value=None):
        if value is not None and (value & (value - 1)):
            return argparse.ArgumentTypeError('must be power of 2')
        elif self.imin is not None and self.imax is not None:
            return argparse.ArgumentTypeError(f'must be an integer power of 2 in the range [{self.imin}, {self.imax}]')
        elif self.imin is not None:
            return argparse.ArgumentTypeError(f'must be an integer power of 2 >= {self.imin}')
        elif self.imax is not None:
            return argparse.ArgumentTypeError(f'must be an integer power of 2 <= {self.imax}')
        else:
            return argparse.ArgumentTypeError('must be an integer power of 2')


# Custom argparse type representing a bounded int
class IntRangeType:
    def __init__(self, imin=None, imax=None):
        self.imin = imin
        self.imax = imax

    def __call__(self, arg):
        try:
            value = int(arg)
        except ValueError:
            raise self.exception()
        if (self.imin is not None and value < self.imin) or (self.imax is not None and value > self.imax):
            raise self.exception()
        return value

    def exception(self):
        if self.imin is not None and self.imax is not None:
            return argparse.ArgumentTypeError(f'must be an integer in the range [{self.imin}, {self.imax}]')
        elif self.imin is not None:
            return argparse.ArgumentTypeError(f'must be an integer >= {self.imin}')
        elif self.imax is not None:
            return argparse.ArgumentTypeError(f'must be an integer <= {self.imax}')
        else:
            return argparse.ArgumentTypeError('must be an integer')


# Custom argparse type representing a human readable number
class HumanReadableType:
    def __init__(self, vmin=None, vmax=None):
        self.vmin = vmin
        self.vmax = vmax
        if vmin is not None:
            self.vmin = decimalFromHumanReadable(vmin)
        if vmax is not None:
            self.vmax = decimalFromHumanReadable(vmax)

    def __call__(self, arg):
        try:
            value = decimalFromHumanReadable(arg)
        except ValueError:
            raise self.valueException()
        except TypeError:
            raise self.typeException()
        if (self.vmin is not None and value < self.vmin) or (self.vmax is not None and value > self.vmax):
            raise self.rangeException()
        return value

    def typeException(self):
        return argparse.ArgumentTypeError('invalid unit')

    def valueException(self):
        return argparse.ArgumentTypeError('invalid value')

    def rangeException(self):
        if self.vmin is not None and self.vmax is not None:
            return argparse.ArgumentTypeError(f'value must be in the range [{decimalToHumanReadable(self.vmin)}, {decimalToHumanReadable(self.vmax)}]')
        elif self.vmin is not None:
            return argparse.ArgumentTypeError(f'value must be >= {decimalToHumanReadable(self.vmin)}')
        elif self.vmax is not None:
            return argparse.ArgumentTypeError(f'value must be <= {decimalToHumanReadable(self.vmax)}')


# Custom argparse type representing a path
class PathType:
    def __call__(self, arg):
        if os.path.isdir(arg):
            return os.path.realpath(arg)
        else:
            raise argparse.ArgumentTypeError('invalid path')


# Custom argparse type representing a path to a file
class FileType:
    def __call__(self, arg):
        if os.path.isfile(arg):
            return os.path.realpath(arg)
        else:
            raise argparse.ArgumentTypeError('invalid file')


class StoreChoiceValue(argparse.Action):
    def __call__(self, parser, namespace, values, option_string=None):
        setattr(namespace, self.dest, self.choices.index(values))


class ChangedArg(argparse.Action):
    def __init__(self, option_strings, new_arg, *args, **kwargs):
        self.new_arg = new_arg
        super(ChangedArg, self).__init__(option_strings=option_strings, *args, **kwargs)

    # Warn the user that the argument has changed, but try to set the new one with the received value
    def __call__(self, parser, namespace, values, option_string=None):
        msg.warning(f'Argument {"/".join(self.option_strings)} has changed to --{self.new_arg}. Check `ait -h` and fix your program call')
        setattr(namespace, self.new_arg, values)


class DeprecatedArg(argparse.Action):
    def __call__(self, parser, namespace, values, option_string=None):
        msg.warning(f'Argument {"/".join(self.option_strings)} is deprecated. Please remove it from the program call')


class CustomParser(argparse.ArgumentParser):
    def error(self, message):
        sys.stderr.write(f'usage: ait -b BOARD -n NAME\nait error: {message}\n')
        sys.exit(2)


class ArgParser:
    defaults = dict()
    for backend in backends:
        defaults[backend] = dict()

        # IP cache location defaults
        defaults[backend]['IP_cache_location'] = {
            'value': f'/var/tmp/ait/{backend}/IP_cache',
            'used': False
        }

        defaults[backend]['hwinst'] = {
            'value': False,
            'used': False
        }

    def __init__(self):
        # Load vendor-specific parsers
        self.backend_parser = dict()
        for backend in backends:
            backend_parser = f'ait.backend.{backend}.utils.parser'
            if importlib.util.find_spec(backend_parser):
                self.backend_parser[backend] = importlib.import_module(backend_parser)

        # Create main parser
        self.parser = CustomParser(formatter_class=argparse.RawTextHelpFormatter, epilog='\n'.join(self.backend_parser[backend].help_text for backend in self.backend_parser), prog='ait', add_help=False, usage='%(prog)s -b BOARD -n NAME\nThe Accelerator Integration Tool (AIT) automatically integrates OmpSs@FPGA accelerators into FPGA designs using different vendor backends.')

        # Required arguments
        required_args = self.parser.add_argument_group('Required')
        required_args.add_argument('-b', '--board',
                                   help='board model. Supported boards by vendor:\n' + '\n'.join(backend + ': ' + ', '.join(board for board in backends[backend]['boards']) for backend in backends),
                                   choices=[board for board in backends[backend]['boards'] for backend in backends],
                                   metavar='BOARD',
                                   type=str.lower,
                                   required=True)

        required_args.add_argument('-n', '--name',
                                   help='project name',
                                   metavar='NAME',
                                   required='--dump_board_info' not in sys.argv)

        # Generation flow arguments
        flow_args = self.parser.add_argument_group('Generation flow')
        flow_args.add_argument('-d', '--dir',
                               help='path where the project directory tree will be created\
                                    \n(def: .)',
                               type=PathType(),
                               default='.')

        flow_args.add_argument('--disable_IP_caching',
                               help='disable IP caching\
                                    \nSignificantly increases generation time',
                               action='store_true',
                               default=False)

        flow_args.add_argument('--disable_utilization_check',
                               help='disable resources utilization check during HLS generation',
                               action='store_true',
                               default=False)

        flow_args.add_argument('--disable_board_support_check',
                               help='disable board support check',
                               action='store_true',
                               default=False)

        flow_args.add_argument('--from_step',
                               help='initial generation step\
                                    \nGeneration steps by vendor:\n' + '\n'.join(backend + ': ' + ', '.join(step for step in backends[backend]['steps']) + "\n(def: '{}')".format(backends[backend]['initial_step']) for backend in backends),
                               choices=[step for step in backends[backend]['steps'] for backend in backends],
                               metavar='FROM_STEP',
                               default='HLS')

        flow_args.add_argument('--IP_cache_location',
                               help='path where the IP cache will be located\
                                    \n(def: /var/tmp/ait/<vendor>/IP_cache/)',
                               type=PathType())

        flow_args.add_argument('--to_step',
                               help='final generation step\
                                    \nGeneration steps by vendor:\n' + '\n'.join(backend + ': ' + ', '.join(step for step in backends[backend]['steps']) + "\n(def: '{}')".format(backends[backend]['final_step']) for backend in backends),
                               choices=[step for step in backends[backend]['steps'] for backend in backends],
                               metavar='TO_STEP',
                               default='bitstream')

        # Bitstream configuration arguments
        bitstream_args = self.parser.add_argument_group('Bitstream configuration')
        bitstream_args.add_argument('-c', '--clock',
                                    help='FPGA clock frequency in MHz\
                                         \n(def: 100)',
                                    type=int,
                                    default='100')

        bitstream_args.add_argument('--hwcounter',
                                    help='add a hardware counter to the bitstream',
                                    action='store_true',
                                    default=False)

        bitstream_args.add_argument('--bitinfo_note',
                                    help='custom note to add to the bitinfo',
                                    type=ascii,
                                    default='')

        bitstream_args.add_argument('--disable_static_constraints',
                                    help='disable static constraints \
                                         \nMay impact negatively on timing',
                                    action='store_true',
                                    default=False)

        # Data path arguments
        datapath_args = self.parser.add_argument_group('Data path')
        datapath_args.add_argument('--memory_interleaving_stride',
                                   help='size in bytes of the stride of the memory interleaving\
                                        \nMust be power of 2\
                                        \nIf set to 0 bytes, interleaving will not be enabled\
                                        \n(def: 0)',
                                   metavar='MEM_INTERLEAVING_STRIDE',
                                   type=HumanReadableType(vmin='4K'),
                                   default=0)

        # Hardware Runtime arguments
        hwruntime_args = self.parser.add_argument_group('Hardware Runtime')
        cmdin_len = hwruntime_args.add_mutually_exclusive_group()
        cmdin_len.add_argument('--cmdin_queue_len',
                               help='maximum length (64-bit words) of the queue for the hwruntime command in\
                                    \nThis argument is mutually exclusive with --cmdin_subqueue_len',
                               type=IntRangeType(imin=4))
        cmdin_len.add_argument('--cmdin_subqueue_len',
                               help='length (64-bit words) of each accelerator subqueue for the hwruntime command in\
                                    \nThis argument is mutually exclusive with --cmdin_queue_len\
                                    \nMust be power of 2\
                                    \n(def: max(64, 1024/num_instances))',
                               type=IntPowerType(imin=4))

        cmdout_len = hwruntime_args.add_mutually_exclusive_group()
        cmdout_len.add_argument('--cmdout_queue_len',
                                help='maximum length (64-bit words) of the queue for the hwruntime command out\
                                     \nThis argument is mutually exclusive with --cmdout_subqueue_len',
                                type=IntRangeType(imin=2))

        cmdout_len.add_argument('--cmdout_subqueue_len',
                                help='length (64-bit words) of each accelerator subqueue for the hwruntime command out\
                                     \nThis argument is mutually exclusive with --cmdout_queue_len\
                                     \nMust be power of 2\
                                     \n(def: max(64, 1024/num_instances))',
                                type=IntPowerType(imin=2))

        hwruntime_args.add_argument('--disable_spawn_queues',
                                    help='disable the hwruntime spawn in/out queues',
                                    action='store_true',
                                    default=False)

        hwruntime_args.add_argument('--spawnin_queue_len',
                                    help='length (64-bit words) of the hwruntime spawn in queue\
                                         \nMust be power of 2\
                                         \n(def: 1024)',
                                    type=IntPowerType(imin=4),
                                    default=1024)

        hwruntime_args.add_argument('--spawnout_queue_len',
                                    help='length (64-bit words) of the hwruntime spawn out queue\
                                         \nMust be power of 2\
                                         \n(def: 1024)',
                                    type=IntPowerType(imin=4),
                                    default=1024)

        hwruntime_args.add_argument('--hwruntime_interconnect',
                                    help='type of hardware runtime interconnection with accelerators\
                                         \ncentralized\
                                         \ndistributed\
                                         \n(def: centralized)',
                                    choices=['centralized', 'distributed'],
                                    metavar='HWR_INTERCONNECT',
                                    default='centralized')  # TODO: Explain what does each option do

        hwruntime_args.add_argument('--max_args_per_task',
                                    help='maximum number of arguments for any task in the bitstream\
                                         \n(def: 15)',
                                    type=IntRangeType(imin=1),
                                    default=15)

        hwruntime_args.add_argument('--max_deps_per_task',
                                    help='maximum number of dependencies for any task in the bitstream\
                                         \n(def: 8)',
                                    type=IntRangeType(imin=2),
                                    default=8)

        hwruntime_args.add_argument('--max_copies_per_task',
                                    help='maximum number of copies for any task in the bitstream\
                                         \n(def: 15)',
                                    type=IntRangeType(imin=1),
                                    default=15)

        hwruntime_args.add_argument('--enable_pom_axilite',
                                    help='enable the POM axilite interface with debug counters',
                                    action='store_true',
                                    default=False)

        # Picos arguments
        picos_args = self.parser.add_argument_group('Picos')
        picos_args.add_argument('--picos_num_dcts',
                                help='number of DCTs to instantiate\
                                     \n(def: 1)',
                                type=int,
                                choices=[1, 2, 4],
                                metavar='NUM_DCTS',
                                default=1)

        picos_args.add_argument('--picos_tm_size',
                                help='size of the TM memory\
                                     \n(def: 128)',
                                type=IntRangeType(imin=2),
                                default=128)

        picos_args.add_argument('--picos_dm_size',
                                help='size of the DM memory\
                                     \n(def: 512)',
                                type=IntRangeType(imin=2),
                                default=512)

        picos_args.add_argument('--picos_vm_size',
                                help='size of the VM memory\
                                     \n(def: 512)',
                                type=IntRangeType(imin=2),
                                default=512)

        picos_args.add_argument('--picos_dm_ds',
                                help='data structure of the DM memory\
                                     \nBINTREE: Binary search tree (not autobalanced)\
                                     \nLINKEDLIST: Linked list\
                                     \n(def: BINTREE)',
                                choices=['BINTREE', 'LINKEDLIST'],
                                metavar='DATA_STRUCT',
                                default='BINTREE')

        picos_args.add_argument('--picos_dm_hash',
                                help='hashing function applied to dependence addresses\
                                     \nP_PEARSON: Parallel Pearson function\
                                     \nXOR\
                                     \n(def: P_PEARSON)',
                                choices=['P_PEARSON', 'XOR'],
                                metavar='HASH_FUN',
                                default='P_PEARSON')  # TODO: Explain what XOR does

        picos_args.add_argument('--picos_hash_t_size',
                                help='DCT hash table size\
                                     \n(def: 64)',
                                type=IntRangeType(imin=2),
                                default=64)

        # User-defined files arguments
        user_args = self.parser.add_argument_group('User-defined files')
        user_args.add_argument('--user_config',
                               help='path to the JSON file containing user configuration',
                               type=FileType())

        user_args.add_argument('--user_constraints',
                               help='path to the user defined constraints file',
                               type=FileType())

        user_args.add_argument('--user_pre_design',
                               help='path to the user TCL script to be executed before the design step (not after the board base design)',
                               type=FileType())

        user_args.add_argument('--user_post_design',
                               help='path to the user TCL script to be executed after the design step',
                               type=FileType())

        # Miscellaneous arguments
        misc_args = self.parser.add_argument_group('Miscellaneous')
        misc_args.add_argument('-h', '--help',
                               action='help',
                               help='show this help message and exit')

        misc_args.add_argument('-i', '--verbose_info',
                               help='print extra information messages',
                               action='store_true',
                               default=False)

        misc_args.add_argument('--dump_board_info',
                               help='dump board info json for the specified board',
                               action='store_true',
                               default=False)

        misc_args.add_argument('-j', '--jobs',
                               help='specify the number of jobs to run simultaneously\
                                    \nBy default it will use as many jobs as cores with at least 5GB of dedicated free memory, or the value returned by `nproc`, whichever is less.',
                               type=IntRangeType(imin=1),
                               default=None)

        misc_args.add_argument('--mem_per_job',
                               help='specify the memory per core used to estimate the number of jobs to launch\
                                    \n(def: 5G)',
                               type=HumanReadableType(vmin='1G'),
                               default='5G')

        misc_args.add_argument('-k', '--keep_files',
                               help='keep files on error',
                               action='store_true',
                               default=False)

        misc_args.add_argument('-v', '--verbose',
                               help='print vendor backend messages',
                               action='store_true',
                               default=False)

        misc_args.add_argument('--version',
                               help='print AIT version and exits',
                               action='version',
                               version=str(LONG_VERSION))

        # Hidden arguments
        hidden_args = self.parser.add_argument_group('Hidden')
        hidden_args.add_argument('--wrapper_version',
                                 help=argparse.SUPPRESS,
                                 type=int,
                                 default=0)

        hidden_args.add_argument('--task_creation',
                                 help=argparse.SUPPRESS,
                                 action='store_true',
                                 default=False)

        hidden_args.add_argument('--deps_hwruntime',
                                 help=argparse.SUPPRESS,
                                 action='store_true',
                                 default=False)

        hidden_args.add_argument('--lock_hwruntime',
                                 help=argparse.SUPPRESS,
                                 action='store_true',
                                 default=False)

    def parse_args(self):
        # Check if configuration file exists and parse its values
        config_file_path = f'{os.path.expanduser("~")}/.ait/config.json'
        if os.path.exists(config_file_path):
            with open(config_file_path) as config_file:
                try:
                    config = json.load(config_file)
                    self.parser.set_defaults(**config)
                except ValueError as err:
                    print(f'Invalid configuration file ({config_file_path}). {err}')
                    print('================')
                    config_file.seek(0)
                    print(config_file.read())
                    print('================')
                    sys.exit(1)

        # Parse values from argv
        args, extras = self.parser.parse_known_args()

        # Get vendor backend and parse vendor-specific arguments
        args_vars = vars(args)
        args_vars['backend'] = [backend for backend in backends if args.board in backends[backend]['boards']][0]
        if (args.backend in self.backend_parser):
            args, extras = self.backend_parser[args.backend].parser.parse_known_args(extras, namespace=args)
        if len(extras):
            self.parser.error(f"unrecognized arguments: {','.join(extras)}. Try 'ait -h' for more information.")

        # Set default values for non-provided options
        for arg in self.defaults[args_vars['backend']]:
            if arg not in args_vars or args_vars[arg] is None:
                args_vars[arg] = self.defaults[args_vars['backend']][arg]['value']
                self.defaults[args_vars['backend']][arg]['used'] = True

        self.check_args(args)

        return args

    def check_args(self, args):
        # Validate arguments
        self.check_required_args(args)
        self.check_flow_args(args)
        self.check_bitstream_args(args)
        self.check_misc_args(args)

    def check_required_args(self, args):
        # Validate required args
        # Some arguments not required when --dump_board_info argument is present
        if not args.dump_board_info:
            if re.search('(^[^A-Za-z])|(\W|__)|([^A-Za-z0-9]$)', args.name):
                msg.error('Invalid project name. Must start with a letter, contain only letters, numbers or non-consecutive underscores, and end with a letter or number')

            if not os.path.exists(f'{args.dir}/{args.name}_ait'):
                os.mkdir(f'{args.dir}/{args.name}_ait')

            if args.wrapper_version and args.wrapper_version < MIN_WRAPPER_VERSION:
                msg.error(f'Unsupported wrapper version ({args.wrapper_version}). Minimum version is {MIN_WRAPPER_VERSION}')

    def check_flow_args(self, args):
        # Validate flow args
        if args.from_step not in backends[args.backend]['steps']:
            msg.error(f"Initial step '{args.from_step}' is not a valid generation step for '{args.backend}' backend. Set it correctly")

        if args.to_step not in backends[args.backend]['steps']:
            msg.error(f"Final step '{args.to_step}' is not a valid generation step for '{args.backend}' backend. Set it correctly")

        if backends[args.backend]['steps'].index(args.from_step) > backends[args.backend]['steps'].index(args.to_step):
            msg.error(f"Initial step '{args.from_step}' is posterior to the final step '{args.to_step}'. Set them correctly")

        if not args.disable_IP_caching and not os.path.isdir(args.IP_cache_location):
            if self.is_default('IP_cache_location', args.backend):
                # Create cache folder and set perms to allow all users writing there
                os.makedirs(args.IP_cache_location)
                os.chmod(args.IP_cache_location, 0o777)
            else:
                msg.error(f'Cache location ({args.IP_cache_location}) does not exist or is not a folder')

    def check_bitstream_args(self, args):
        # Validate bitstream args
        if args.bitinfo_note and len(args.bitinfo_note) > 256:
            msg.error('Length of bitinfo note must be less than 256 ASCII chars')

    # Validate hardware runtime args
    def check_hardware_runtime_args(self, args):
        # This check has to be delayed because arguments are parsed before the number of accelerators is calculated
        def prev_power_of_2(num):
            if num & (num - 1) != 0:
                num = int(math.log2(num))
                num = int(pow(2, num))
            return num

        if args.num_accs > 16:
            msg.error(f'Max number of accelerator types supported by POM is 16, but design has {args.num_accs}')

        if args.cmdin_queue_len is not None:
            args.cmdin_subqueue_len = prev_power_of_2(int(args.cmdin_queue_len / max(2, args.num_instances)))
            msg.info(f'Setting --cmdin_subqueue_len to {args.cmdin_subqueue_len}')
        elif args.cmdin_subqueue_len is None:
            args.cmdin_subqueue_len = max(64, prev_power_of_2(int(1024 / max(2, args.num_instances))))
        if args.cmdout_queue_len is not None:
            args.cmdout_subqueue_len = prev_power_of_2(int(args.cmdout_queue_len / max(2, args.num_instances)))
            msg.info(f'Setting --cmdout_subqueue_len to {args.cmdout_subqueue_len}')
        elif args.cmdout_subqueue_len is None:
            args.cmdout_subqueue_len = max(64, prev_power_of_2(int(1024 / max(2, args.num_instances))))

        # The subqueue length has to be checked here in the case the user provides the cmdin queue length
        if args.cmdin_subqueue_len < 34:
            msg.warning(f'Value of --cmdin_subqueue_len={args.cmdin_subqueue_len} is less than 34, which is the length of the longest command possible. This design might not work with tasks with enough arguments.')
        if not args.disable_spawn_queues and args.spawnout_queue_len < 79:
            msg.warning(f'Value of --spawnout_queue_len={args.spawnout_queue_len} is less than 79, which is the length of the longest task possible. This design might not work if an accelerator creates SMP tasks with enough copies, dependencies and/or arguments.')

    def check_picos_args(self, args):
        # Validate Picos args
        if (args.picos_dm_hash == 'P_PEARSON' and args.picos_hash_t_size != 64):
            msg.error('With P_PEARSON hash function, picos_hash_t_size must be 64')
        if (args.picos_hash_t_size > args.picos_dm_size):
            msg.error(f'Invalid picos_hash_t_size ({args.picos_hash_t_size}) value, maximum value is picos_dm_size ({args.picos_dm_size})')
        if (math.ceil(math.log2(args.picos_hash_t_size)) + math.ceil(math.log2(args.picos_num_dcts)) > 8):
            msg.error(f'Binary representation of (picos_hash_t_size + picos_num_dcts) must be smaller than 8 bits ({math.ceil(math.log2(args.picos_hash_t_size)) + math.ceil(math.log2(args.picos_num_dcts))} > 8)')

    def check_misc_args(self, args):
        if args.jobs and args.jobs > getNumJobs(args.mem_per_job):
            msg.warning(f'Using more Vivado jobs ({args.jobs}) than the recommended default ({getNumJobs(args.mem_per_job)}). Performance of the compilation process might be affected')

    def is_default(self, dest, backend):
        value = False
        if self.defaults[backend][dest] is not None:
            value = self.defaults[backend][dest]['used']
        return value
