##
## Copyright (C) by Argonne National Laboratory
##     See COPYRIGHT in top-level directory
##

from local_python import MPI_API_Global as G
from local_python import RE
import re

def get_kind_map(lang, is_large=False):
    if lang.upper() == 'C':
        if is_large:
            return G.MAPS['BIG_C_KIND_MAP']
        else:
            return G.MAPS['SMALL_C_KIND_MAP']
    elif lang.upper() == 'F08':
        if is_large:
            return G.MAPS['BIG_F08_KIND_MAP']
        else:
            return G.MAPS['SMALL_F08_KIND_MAP']
    elif lang.upper() == 'F90':
        return G.MAPS['SMALL_F90_KIND_MAP']
    elif lang.upper() == "LIS":
        return G.MAPS['LIS_KIND_MAP']
    else:
        raise Exception("Kind Mapping for [%s] not available" % lang)

def function_has_POLY_parameters(func):
    for p in func['parameters']:
        if p['kind'].startswith('POLY'):
            return True
    return False

def get_function_name(func, is_large=False):
    if is_large:
        name = func['name'] + "_c"
        return name
    else:
        name = func['name']
        if 'mpix' in func:
            G.mpix_symbols[name] = "functions"
            name = re.sub(r'MPI_', 'MPIX_', name)
        return name

def split_line_with_break(s, tail, N=100):
    """Breaks a long line with proper indentations.
    This simplistic routine splits on ", ", thus only works with function declarations
    and simple function calls such as those generated by this script. """
    out_list = []

    tlist = []
    n = 0

    # by default, segments indent by additional 4 spaces
    if RE.match(r'(\s*)', s):
        n_lead = len(RE.m.group(1)) + 4

    # -------------
    def break_s_next(s_lead, s_next):
        nonlocal tlist, n
        for a in s_next.split(', '):
            if n == 0:
                # first line
                tlist = [s_lead, a]
                n = n_lead + len(a)
            elif n + 2 + len(a) < N:
                # just append to tlist
                tlist.append(', ')
                tlist.append(a)
                n += 2 + len(a)
            else:
                # break the line
                tlist.append(',')
                out_list.append(''.join(tlist))
                # start new line with leading spaces
                # if lead is too much, it won't look good
                if n_lead > N - 40:
                    tlist = [' ' * 20, a]
                    n = 20 + len(a)
                else:
                    tlist = [' ' * n_lead, a]
                    n = n_lead + len(a)
        # leave last segment with tail

    # -------------
    if len(s) < N:
        tlist.append(s)
        n = len(s)
    elif RE.match(r'(.*?\()(.*)', s):
        # line with function pattern, match indent at opening parenthesis
        s_lead, s_next = RE.m.group(1,2)
        n_lead = len(s_lead)

        break_s_next(s_lead, s_next)
    else:
        break_s_next('', s)

    # tail is mostly for "__attribute__ ((weak, alias(...))));",
    # which contains , that we do not desire to break
    if tail:
        if n + 1 + len(tail) < 100:
            out_list.append(''.join(tlist) + ' ' + tail)
        else:
            out_list.append(''.join(tlist))
            out_list.append(' ' * n_lead + tail)
    else:
        out_list.append(''.join(tlist))

    return out_list

def get_C_param(param, func, mapping):
    kind = param['kind']
    if kind == "VARARGS":
        return "..."

    want_star, want_bracket = '', ''
    param_type = mapping[kind]

    if param['func_type']:
        param_type = param['func_type']
        if mapping['_name'].startswith("BIG_"):
            param_type += "_c"

    if param_type in G.mpix_symbols:
        param_type = re.sub(r'MPI_', 'MPIX_', param_type)

    if not param_type:
        raise Exception("Type mapping [%s] %s not found!" % (mapping, kind))

    # We special treat a few cases to simplify the general rules
    if kind == "ARGUMENT_COUNT":
        if re.match(r'mpi_info_create_env', func['name'], re.IGNORECASE):
            # -> int argc
            pass
        else:
            # MPI_Init, MPI_Init_thread -> int *argc
            want_star = 1
    elif kind == "ARGUMENT_LIST":
        if re.match(r'mpi_info_create_env', func['name'], re.IGNORECASE):
            # -> char *argv[]
            want_star = 1
            want_bracket = 1
        else:
            # MPI_Init, MPI_Init_thread -> char ***argv
            want_star = 3
    elif kind == "STATUS" and param['length'] is not None:
        # MPI_{Wait,Test}{all,some}
        # gcc-11 warns when we pass MPI_STATUSES_IGNORE to MPI_Status statues[]
        want_star = 1
    elif not want_star:
        if is_pointer_type(param):
            if kind == "STRING_ARRAY":
                want_star = 1
                want_bracket = 1
            elif kind == "STRING_2DARRAY":
                want_star = 2
                want_bracket = 1
            elif param['pointer'] is not None and not param['pointer']:
                want_bracket = 1
            elif param['length'] is not None and kind != "STRING":
                want_bracket = 1
            else:
                want_star = 1

    s = ''
    if param['constant']:
        s += "const "
    s += param_type

    if want_star:
        s += " " + "*" * want_star
    else:
        s += " "
    s += param['name']

    if want_bracket:
        s += "[]"
    if isinstance(param['length'], list):
        s += "[%s]" % param['length'][-1]

    return s

def is_pointer_type(param):
    if RE.match(r'(STRING\w*)$', param['kind']):
        return 1
    elif RE.match(r'(STATUS|F90_STATUS|F08_STATUS)$', param['kind']):
        return 1
    elif RE.match(r'(ATTRIBUTE_VAL\w*|(C_)?BUFFER\d?|EXTRA_STATE\d*|TOOL_MPI_OBJ|(POLY)?FUNCTION\w*)$', param['kind']):
        return 1
    elif param['param_direction'] != 'in':
        return 1
    elif param['length']:
        return 1
    elif param['pointer']:
        return 1
    else:
        return 0

def param_is_in(p):
    if p['param_direction'] == 'in' or p['param_direction'] == 'inout':
        return True
    else:
        return False

def param_is_out(p):
    if p['param_direction'] == 'out' or p['param_direction'] == 'inout':
        return True
    else:
        return False

def get_userbuffer_group(func_name, parameters, i):
    """internal function used by process_func_parameters"""
    p = parameters[i]
    p2 = parameters[i + 1]
    p3 = parameters[i + 2]
    if RE.match(r'mpi_i?(alltoall|allgather|gather|scatter)', func_name, re.IGNORECASE):
        type = "inplace"
        if RE.search(r'send', p['name'], re.IGNORECASE) and RE.search(r'scatter', func_name, re.IGNORECASE):
            type = "noinplace"
        elif RE.search(r'recv', p['name'], re.IGNORECASE) and not RE.search(r'scatter', func_name, re.IGNORECASE):
            type = "noinplace"

        if RE.search(r'alltoallw', func_name, re.IGNORECASE):
            group_kind = "USERBUFFER-%s-w" % (type)
            group_count = 4
        elif p3['kind'] == "DATATYPE":
            group_kind = "USERBUFFER-%s" % (type)
            group_count = 3
        else:
            group_kind = "USERBUFFER-%s-v" % (type)
            group_count = 4
    elif RE.match(r'mpi_i?neighbor', func_name, re.IGNORECASE):
        if RE.search(r'alltoallw', func_name, re.IGNORECASE):
            group_kind = "USERBUFFER-neighbor-w"
            group_count = 4
        elif p3['kind'] == "DATATYPE":
            group_kind = "USERBUFFER-neighbor"
            group_count = 3
        else:
            group_kind = "USERBUFFER-neighbor-v"
            group_count = 4
    elif RE.match(r'mpi_i?(allreduce|reduce|scan|exscan)', func_name, re.IGNORECASE):
        group_kind = "USERBUFFER-reduce"
        group_count = 5
    elif RE.match(r'mpi_p(send|recv)_init', func_name, re.IGNORECASE):
        group_kind = "USERBUFFER-partition"
        group_count = 4
    elif RE.search(r'XFER_NUM_ELEM', p2['kind']) and RE.search(r'DATATYPE', p3['kind']):
        group_kind = "USERBUFFER-simple"
        group_count = 3
    else:
        group_kind, group_count = None, 0
    return (group_kind, group_count)