#! /usr/bin/env python3

"""Convert TG-LRX2 command script to binary encoded device specific downlink messages."""

import sys
import msgpack
import yaml
import re
import base64
import copy
import logging
import argparse

log = logging.getLogger(__name__)


class CommandGroup:
    """Group of commands sharing same prefix (tokens)"""

    def __init__(self, prefix_list, value=None):
        self.prefix = prefix_list
        self.suffix_list = []
        self.encoded = None
        if value is not None:
            self.suffix_list.append(value)

    def insert(self, command):
        if len(self.suffix_list) == 0 or isinstance(self.suffix_list[-1], str):
            self.suffix_list.append(CommandGroup(command, "exe"))
        else:
            last = self.suffix_list[-1]
            match, remainder = CommandGroup._first_matching(last.prefix, command)

            if len(match) > 0:
                if len(match) >= len(last.prefix):
                    if len(remainder):
                        self.suffix_list[-1].insert(remainder)
                    else:
                        self.suffix_list[-1].suffix_list.append("exe")
                else:
                    last.prefix = last.prefix[len(match):]
                    self.suffix_list[-1] = CommandGroup(match, last)
                    self.suffix_list[-1].insert(remainder)
            else:
                self.suffix_list.append(CommandGroup(command, "exe"))

    def make_token_struct(self):
        """Create device specific nested token array structure with embedded commands

        Token sequence at the end of array is executed as a command. Nested arrays can be used to
        create commands with same prefix.

        Example structure [1,2,3,[10,11],[12,13],[14,15]] will execute following commands:
        1,2,3,10,11
        1,2,3,12,13
        1,2,3,14,15
        """

        tokens = []
        tokens.extend(self.prefix)

        for suffix in self.suffix_list:
            if isinstance(suffix, CommandGroup):
                tokens.append(suffix.make_token_struct())

        # Array with single nested array can be simplified to reduce size.
        while len(tokens) == 1 and isinstance(tokens[0], list):
            tokens = tokens[0]

        # Last nested array can be placed inline to reduce total size.
        while len(tokens) > 1 and isinstance(tokens[-1], list):
            make_inline = tokens[-1]
            tokens.pop()
            tokens.extend(make_inline)

        return tokens

    @staticmethod
    def _first_matching(a, b):
        matching_prefix = []
        for x, y in zip(a, b):
            if x != y:
                break
            matching_prefix.append(x)
        return matching_prefix, b[len(matching_prefix):]

    def get_bytes(self):
        token_struct = self.make_token_struct()

        if len(token_struct) == 0:
            return b""

        self.encoded = msgpack.packb(token_struct)
        return self.encoded

    def text_repr(self, nest_level=0):
        txt = "   " * nest_level + f"{self.prefix}, bytes: {len(self.get_bytes())}"

        for suffix in self.suffix_list:
            if isinstance(suffix, CommandGroup):
                txt += "\n"
                txt += suffix.text_repr(nest_level + 1)
            else:
                txt += " "

        return txt

    def __str__(self):
        return self.text_repr()


def error(msg, expert=False):
    if not expert:
        raise ValueError(msg)
    log.error(msg)


def prefix_compression(cmd_list, max_msg_size, is_expert):
    """Group sequential messages with same prefix and encode to binary

    Since user can specify arbitrary maximal DL message size, there is a possibility that
    we can't encode command in requested size.
    """

    root = CommandGroup([])
    msg_list = []
    invalid_size = 0

    for cmd in cmd_list:
        new = copy.deepcopy(root)
        new.insert(cmd)

        if len(new.get_bytes()) > max_msg_size:
            msg_list.append(root.get_bytes())

            if len(msg_list[-1]) > max_msg_size:
                error(f"size violation for #{len(msg_list) - 1}: '{root.make_token_struct()}' "
                      f"requires {len(msg_list[-1])}, but limit is {max_msg_size} bytes.",
                      is_expert)

                invalid_size += 1

            root = CommandGroup([])
            root.insert(cmd)
        else:
            root = new

    msg_list.append(root.get_bytes())

    return msg_list, invalid_size


def tokenize_cmd(cmd):
    token_list = re.split(r',', cmd)
    token_list = re.split(r' ', token_list[0], maxsplit=2) + token_list[1:]
    return [x.strip() for x in token_list]


def alias_subs(cmd_tokens, dev_desc, iface_name):
    """Replace group and command text tokens with integer codes to reduce size"""
    spec = dev_desc[iface_name]

    if cmd_tokens[0] in spec:
        gr_name = cmd_tokens[0]
        gr_id = spec[gr_name]['id']
        cmd_tokens[0] = str(gr_id)

        if cmd_tokens[1] in spec[gr_name]['commands']:
            cmd_name = spec[gr_name]['commands'][cmd_tokens[1]]
            cmd_id = cmd_name['id']
            cmd_tokens[1] = str(cmd_id)


def sanitize(src, cmd_tokens, dev_desc, iface_name, is_expert=False):
    """Check if valid command and convert arguments for minimal encoding size

    Group and command names are converted to integers. All numeric arguments are converted
    to int or text representation of float. Known frame type names and variable names are
    converted to their respective numeric IDs. All arguments are validated using device
    description spec. In expert mode unknown or invalid commands are reported, but still encoded
    as text.
    """
    src_line = src[0]
    src_str = src[1]

    for n in range(len(cmd_tokens)):
        if len(cmd_tokens[n]) == 0:
            cmd_tokens[n] = None
        else:
            try:
                cmd_tokens[n] = float(cmd_tokens[n])
                if cmd_tokens[n].is_integer():
                    cmd_tokens[n] = int(cmd_tokens[n])
                continue
            except ValueError:
                pass

            try:
                cmd_tokens[n] = int(cmd_tokens[n], 16)
            except ValueError:
                pass

    groups_by_id = {v['id']: {'name': k, 'commands': v['commands']} for k, v in
                    dev_desc[iface_name].items()}
    variables_by_id = {v['id']: v for k, v in dev_desc['variables'].items()}
    frames_by_id = {v['id']: v for k, v in dev_desc['frames'].items()}

    if cmd_tokens[0] not in groups_by_id:
        error(f"unknown command: '{src_str}', line {src_line}", is_expert)
    else:
        group = groups_by_id[cmd_tokens[0]]
        cmds_by_id = {v['id']: {'name': k, 'param-list': v['param-list']} for k, v in
                      group['commands'].items()}

        if cmd_tokens[1] not in cmds_by_id:
            error(f"unknown command: '{src_str}', line {src_line}", is_expert)
        else:
            cmd = cmds_by_id[cmd_tokens[1]]
            cmd_name = f"/{group['name']}/{cmd['name']}"
            full_cmd_str = f"{cmd_name}/{cmd_tokens[2:]}"

            cmd_args = {num: val for num, val in enumerate(cmd_tokens[2:])}

            log.debug(f"command: {full_cmd_str}")

            arg_count = len(cmd_tokens[2:])
            param_count = len(cmd['param-list'])

            if arg_count > param_count and cmd['param-list'][-1]['flexible_count']:
                for n in range(arg_count - param_count):
                    cmd['param-list'].append(cmd['param-list'][-1])
                    cmd['param-list'][-1]['optional'] = True
                    param_count += 1

            # sanitize command arguments
            for n, param_spec in enumerate(cmd['param-list']):
                if n not in cmd_args or cmd_args[n] is None:
                    if not param_spec['optional']:
                        if is_expert and param_spec['val_is_missing'] is not None:
                            cmd_args[n] = int(param_spec['val_is_missing'])
                            log.warning(f"using implicit value: '{cmd_args[n]}' for mandatory "
                                        f"argument '{param_spec['name']}' in '{src_str}'")
                        else:
                            error(
                                f"missing mandatory argument: {param_spec['name']}, line "
                                f"{src_line}", is_expert)
                            continue
                    else:
                        continue

                value_in_range = False

                for val_set in param_spec['value_set']:
                    if val_set['min'] is not None and (
                            isinstance(cmd_args[n], int) or isinstance(cmd_args[n], float)):
                        if val_set['max'] is not None and val_set['min'] <= cmd_args[n] <= \
                                val_set['max']:
                            value_in_range = True
                            break
                        elif cmd_args[n] == val_set['min']:
                            value_in_range = True
                            break

                    if val_set['built_in_type'] == 'var-id':
                        if cmd_args[n] in dev_desc['variables']:
                            cmd_tokens[2 + n] = dev_desc['variables'][cmd_args[n]]['id']
                            value_in_range = True

                        if cmd_args[n] in variables_by_id:
                            value_in_range = True

                    if val_set['built_in_type'] == 'fr-id':
                        if cmd_args[n] in dev_desc['frames']:
                            cmd_tokens[2 + n] = dev_desc['frames'][cmd_args[n]]['id']
                            value_in_range = True

                        if cmd_args[n] in frames_by_id:
                            value_in_range = True

                    if val_set['built_in_type'] == 'texp' and bool(
                            re.fullmatch(r'\^?[\S ]+ *[<>=!%]+ *[+-]? *\d+(.\d+)?', cmd_args[n])):
                        for var_name, var_data in dev_desc['variables'].items():
                            if var_name in cmd_args[n]:
                                cmd_tokens[2 + n] = cmd_tokens[2 + n].replace(var_name,
                                                                              str(var_data['id']))
                        cmd_tokens[2 + n] = cmd_tokens[2 + n].replace(" ", "")
                        value_in_range = True

                    if val_set['built_in_type'] == 'iexp' and bool(
                            re.fullmatch(r'([&|]?!?\d+<?)+', cmd_args[n])):
                        cmd_tokens[2 + n] = cmd_tokens[2 + n].replace(" ", "")
                        value_in_range = True

                    if val_set['built_in_type'] == 'hstr' and bool(
                            re.fullmatch(r'[a-fA-F0-9]+', cmd_args[n])):
                        value_in_range = True

                    if val_set['built_in_type'] == 'fp':
                        cmd_args[n] = str(cmd_args[n])
                        res = re.fullmatch(r'[+-]? *\d+(.\d+)?', cmd_args[n])
                        if bool(res):
                            value_in_range = True
                            cmd_tokens[2 + n] = cmd_args[n]

                if not value_in_range:
                    error(f"argument {n}: '{cmd_args[n]}' in '{full_cmd_str}', line {src_line} "
                          f"out of spec: '{param_spec['value_set_str']}'", is_expert)

                if isinstance(cmd_args[n], float):
                    cmd_tokens[2 + n] = str(cmd_args[n])

            if arg_count > param_count:
                error(f"line {src_line}: '{cmd_name}' accepts {len(cmd['param-list'])} arguments, "
                      f"but command is '{full_cmd_str}'!", is_expert)


def order_key(x):
    if x is None:
        return 0, 0
    if isinstance(x, (int, float)):
        return 1, x
    if isinstance(x, str):
        return 2, x
    if isinstance(x, list):
        return 3, tuple(order_key(i) for i in x)
    return 4, x


def encode_commands(text_list, dev_desc, iface_name, sort_cmd=False, max_msg_size=64, expert=False):
    log.info(f"encoding {len(text_list)} lines")

    cmd_list = []
    for line_num, line in text_list:
        log.debug(f"source line: #{line_num}: '{line}'")
        cmd_token_list = tokenize_cmd(line)
        if len(cmd_token_list) > 0:
            cmd_list.append(cmd_token_list)
            log.debug(f"tokens: {cmd_token_list}")

            if dev_desc:
                alias_subs(cmd_token_list, dev_desc, iface_name)
                sanitize((line_num, line), cmd_token_list, dev_desc, iface_name, expert)
                log.debug(f"final tokens: {cmd_token_list}")
            elif not expert:
                raise Exception(f"Missing device description!")

    log.info(f"tokenized list of {len(cmd_list)} commands")

    if sort_cmd:
        log.info(f"sorting commands")
        cmd_list.sort(key=order_key)

    log.debug(f"binary encoding and compression")
    msg_list, invalid_size = prefix_compression(cmd_list, max_msg_size, expert)

    log.info(f"commands encoded to {len(msg_list)} DL messages with total size: "
             f"{sum(map(len, msg_list))} bytes")

    return msg_list, invalid_size


def process_script(text):
    cmd_list = []
    for n, line in enumerate(text.splitlines()):
        # Exclude comment lines and meta commands
        cmd = re.match(r' *(?!#)(?P<meta>\$\S+)? *(?P<cmd>[\S ]+)?', line)
        if cmd and cmd.groupdict()['cmd']:
            cmd_list.append((n + 1, cmd.groupdict()['cmd']))

    return cmd_list


def read_cmd_script(script_file):
    with open(script_file, "r") as f:
        text = f.read()
        log.debug(f"read {len(text)} bytes from {script_file}")
        return process_script(text)


def read_dev_description(dev_desc_file):
    with open(dev_desc_file, "r") as desc_file:
        return yaml.load(desc_file, Loader=yaml.FullLoader)


def get_param_spec(dev_desc, iface_name):
    """Extract command parameter spec from device description"""
    if iface_name not in dev_desc:
        return False
    iface_data = dev_desc[iface_name]

    for group_name, cmd_group in iface_data.items():
        for cmd_name, cmd in cmd_group['commands'].items():
            cmd['param-list'] = []
            params_list = cmd['params'].split(",")
            for param_str in params_list:
                if len(param_str) > 0:
                    match = re.match(
                        r" *(?:(?P<rep>\d+)\*)? *"
                        r"(?P<opt>\[)? *"
                        r"(?P<fill>\+)? *"
                        r"(?P<name>[^\s(\]{=]+) *"
                        r"(?:\((?P<input>[^)]+)\))? *"
                        r"(?:{(?P<hint>[\S ]+)})? *"
                        r"(?:=(?P<defv>[^\s\]]+))? *"
                        r"]?.*",
                        param_str)

                    if match:
                        input_set = match['input']
                        param_value_set = []
                        if input_set:
                            input_set = input_set.replace("u8max", "255")
                            input_set = input_set.replace("u16max", "65535")
                            input_set = input_set.replace("u32max", "4294967295")
                            input_set = input_set.replace("u8min", "0")
                            input_set = input_set.replace("u16min", "0")
                            input_set = input_set.replace("u32min", "0")

                            input_set = input_set.replace("i8max", "127")
                            input_set = input_set.replace("i16max", "32767")
                            input_set = input_set.replace("i32max", "2147483647")
                            input_set = input_set.replace("i8min", "-128")
                            input_set = input_set.replace("i16min", "-32768")
                            input_set = input_set.replace("i32min", "-2147483648")

                            input_set = input_set.replace("u8", "0..255")
                            input_set = input_set.replace("i8", "-128..127")
                            input_set = input_set.replace("u16", "0..65535")
                            input_set = input_set.replace("i16", "-32768..32767")
                            input_set = input_set.replace("u32", "0..4294967295")
                            input_set = input_set.replace("i32", "-2147483648..2147483647")

                            input_set = input_set.replace("i32", "-2147483648..2147483647")

                            input_choices = input_set.split('/')

                            for n, val_range in enumerate(input_choices):
                                set_match = re.match(
                                    r'(?:(?P<label>[^:]+):)?'
                                    r'(((?P<min>[+-]?\d+)'
                                    r'(?:\.\.(?P<max>[+-]?\d+))?)|'
                                    r'(?P<built_in_type>(texp|iexp|hstr|fr-id|var-id|fp)))',
                                    val_range)

                                if set_match:
                                    value_set = set_match.groupdict()
                                    if value_set['min']:
                                        value_set['min'] = int(value_set['min'])
                                    if value_set['max']:
                                        value_set['max'] = int(value_set['max'])

                                    param_value_set.append(value_set)

                        param_desc = {
                            'optional': match['opt'] is not None,
                            'flexible_count': match['fill'] is not None,
                            'name': match['name'],
                            'value_set': param_value_set,
                            'hint': match['hint'],
                            'val_is_missing': match['defv'],
                            'value_set_str': input_set
                        }
                        if match['rep']:
                            repeat = int(match['rep'])
                            name = param_desc['name']
                            for n in range(repeat):
                                param_desc['name'] = name + f"-{n}"
                                cmd['param-list'].append(copy.deepcopy(param_desc))
                        else:
                            cmd['param-list'].append(copy.deepcopy(param_desc))

    return True


def main():
    parser = argparse.ArgumentParser(description="TG-LRX2 control utility")
    parser.add_argument("-v", type=str, default="WARN", metavar="LEVEL",
                        help="Log level (ERROR, WARN, INFO, DEBUG).")
    parser.add_argument("--dev-desc", type=str, default=None, metavar="DESC.yml",
                        help="Device description file")
    parser.add_argument("--max-len", type=int, default=50, metavar="LEN",
                        help="Maximal uplink message length")
    parser.add_argument("--sort", action=argparse.BooleanOptionalAction,
                        help="Sort commands for best compression")
    parser.add_argument("--port", type=int, default=None, metavar="PORT",
                        help="Override uplink port")
    parser.add_argument("--out", type=str, default=None, metavar="OUT.yml",
                        help="Write output to file")
    parser.add_argument("--expert", action=argparse.BooleanOptionalAction,
                        help="Allow encoding unknown/invalid commands and arguments!")

    group = parser.add_mutually_exclusive_group()
    group.add_argument('--script', type=str, metavar="script.txt", default=None,
                       help="Command script file")
    group.add_argument('--cmd', type=str, metavar="command", default=None,
                       help="Command string")

    args = parser.parse_args()

    log_fmt = '%(asctime)s.%(msecs)03d: %(module)s <%(levelname)s>: %(message)s'
    logging.basicConfig(level="ERROR",
                        format=log_fmt,
                        datefmt='%m/%d %H:%M:%S')
    log.setLevel(args.v)

    dev_desc = None
    cmd_list = None

    if args.script:
        cmd_list = read_cmd_script(args.script)
    elif args.cmd:
        cmd_list = process_script(args.cmd)
    else:
        print("Error: no input script file specified!")
        return 1

    if args.dev_desc:
        dev_desc = read_dev_description(args.dev_desc)
        get_param_spec(dev_desc, 'rpc')
    elif not args.expert:
        print("Error: no device description file specified!")
        return 1

    if len(cmd_list):
        if args.port is None and dev_desc is not None:
            cmd_port = dev_desc['ports']['cmd']
            cfg_port = dev_desc['ports']['cfg']
        else:
            cmd_port = args.port
            cfg_port = args.port

        ul_msg_list, invalid_size = encode_commands(cmd_list, dev_desc, 'rpc', args.sort,
                                                    args.max_len,
                                                    args.expert)

        uplinks = [
            {'port': cmd_port, 'hex': msg.hex(), 'b64': base64.b64encode(msg).decode('ascii'),
             'len': len(msg)}
            for msg in ul_msg_list]

        # Last configuration uplink is sent to "cfg" port to trigger configuration checksum reply.
        uplinks[-1]['port'] = cfg_port

        if args.out is not None:
            with open(args.out, 'w') as outfile:
                yaml.dump(uplinks, outfile)

        print(yaml.dump(uplinks))

        return 0 if invalid_size == 0 else 1

    return 0


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