import re

# Option default value
STYLE_DOCSTRING = 'reStructuredText'

def getOptions():
    """ Get options from the options.txt file

    Returns
    -------
    None
    """
    global STYLE_DOCSTRING
    try:
        with open('options.ini', 'r') as file:
            file_content = file.read()
            style_docstring_match = re.search(r'(?<=style_docstring = )\w+', file_content)
    except FileNotFoundError:
        style_docstring_match = None

    if style_docstring_match:
        STYLE_DOCSTRING = style_docstring_match.group()

def getFunctions(filename: str) -> list:
    """ Get the name, arguments and return type of all functions in a file

    Parameters
    ----------
    filename : str
        Name of the file to get the functions from

    Returns
    -------
    list
        List of functions name, arguments and return type
    """
    with open(filename, 'r') as file:
        content = file.read()
        pattern_function = re.compile(r'(?:async\s+)?def\s+(\w+)\s*\(([^)]*)\)\s*(?:->\s*([^:\n]+))?\s*:', re.MULTILINE)
        functions = re.findall(pattern_function, content)
    return functions

def parseArgument(arg: str) -> str:
    """ Parse the argument name and type

    Parameters
    ----------
    arg : str
        Argument to parse

    Returns
    -------
    str
        Argument name and type
    """
    arg = arg.split('=')[0]
    try:
        arg_name, arg_type = arg.split(':')
        arg_type = arg_type.strip()
    except ValueError:
        arg_name = arg
        arg_type = ''
    return arg_name, arg_type

def writeFunctionCommentLines_google(content: list, index: int, function: list, indent: str) -> list:
    """ Write the function comment lines in Google style

    Parameters
    ----------
    content : list
        Content of the file
    index : int
        Line index where the function is located
    function : list
        Function name, arguments and return type
    indent : str
        Indentation of the comment block

    Returns
    -------
    list
        Content of the file with the comment block added
    """
    docstring_lines = [f'{indent}""" Description\n\n']
    args = [arg.strip() for arg in function[1].split(',') if arg.strip()]
    if args:
        docstring_lines.append(f'{indent}Args:\n')
        for arg in args:
            arg_name, arg_type = parseArgument(arg)
            docstring_lines.append(f'{indent}    {arg_name} {"(" + arg_type + ")" if arg_type else ""}:\n')
    docstring_lines.append(f'{indent}\n')
    docstring_lines.append(f'{indent}Returns:\n')
    docstring_lines.append(f'{indent}    {function[2] + ":" if function[2] else ""}\n')
    docstring_lines.append(f'{indent}"""\n')
    content[index:index] = docstring_lines
    return content

def writeFunctionCommentLines_reStructuredText(content: list, index: int, function: list, indent: str) -> list:
    """ Write the function comment lines in reStructuredText style

    Parameters
    ----------
    content : list
        Content of the file
    index : int
        Line index where the function is located
    function : list
        Function name, arguments and return type
    indent : str
        Indentation of the comment block

    Returns
    -------
    list
        Content of the file with the comment block added
    """
    docstring_lines = [f'{indent}""" Description\n\n']
    args = [arg.strip() for arg in function[1].split(',') if arg.strip()]
    for arg in args:
        arg_name, arg_type = parseArgument(arg)
        if arg_name:
            docstring_lines.append(f'{indent}:param {arg_name}:\n')
            docstring_lines.append(f'{indent}:type {arg_name}: {arg_type}\n')
    docstring_lines.append(f'{indent}:return:\n')
    docstring_lines.append(f'{indent}:rtype: {function[2] if function[2] else ""}\n')
    docstring_lines.append(f'{indent}"""\n')
    content[index:index] = docstring_lines
    return content

def writeFunctionCommentLines_numpy(content: list, index: int, function: list, indent: str) -> list:
    """ Write the function comment lines in Numpy style

    Parameters
    ----------
    content : list
        Content of the file
    index : int
        Line index where the function is located
    function : list
        Function name, arguments and return type
    indent : str
        Indentation of the comment block

    Returns
    -------
    list
        Content of the file with the comment block added
    """
    docstring_lines = [f'{indent}""" Description\n\n']
    args = [arg.strip() for arg in function[1].split(',') if arg.strip()]
    if args:
        docstring_lines.append(f'{indent}Parameters\n')
        docstring_lines.append(f'{indent}----------\n')
        for arg in args:
            arg_name, arg_type = parseArgument(arg)
            docstring_lines.append(f'{indent}{arg_name} : {arg_type}\n\n')
    docstring_lines.append(f'\n')
    docstring_lines.append(f'{indent}Returns\n')
    docstring_lines.append(f'{indent}-------\n')
    docstring_lines.append(f'{indent}{function[2] if function[2] else ""}\n\n')
    docstring_lines.append(f'{indent}"""\n')
    content[index:index] = docstring_lines
    return content

def writeFunctionCommentLines_epytext(content: list, index: int, function: list, indent: str) -> list:
    """ Write the function comment lines in Epytext style

    Parameters
    ----------
    content : list
        Content of the file
    index : int
        Line index where the function is located
    function : list
        Function name, arguments and return type
    indent : str
        Indentation of the comment block

    Returns
    -------
    list
        Content of the file with the comment block added
    """
    docstring_lines = [f'{indent}""" Description\n\n']
    args = [arg.strip() for arg in function[1].split(',') if arg.strip()]
    for arg in args:
        arg_name, arg_type = parseArgument(arg)
        if arg_name:
            docstring_lines.append(f'{indent}@param {arg_name}: {arg_type}\n')
            docstring_lines.append(f'{indent}@type {arg_name}: {arg_type}\n')
    docstring_lines.append(f'{indent}@return:\n')
    docstring_lines.append(f'{indent}@rtype: {function[2] if function[2] else ""}\n')
    docstring_lines.append(f'{indent}"""\n')
    content[index:index] = docstring_lines
    return content

def writeFunctionCommentLines(content: list, index: int, function: list, indentation: int) -> list:
    """ Write the function comment lines in the selected style

    Parameters
    ----------
    content : list
        Content of the file
    index :  int
        Line index where the function is located
    function : list
        Function name, arguments and return type
    indentation : int
        Indentation of the comment block

    Returns
    -------
    list
        Content of the file with the comment block added, or the same content if the function already has a comment block
    """
    if content[index].strip().startswith('"""'):
        return content

    indent = '  ' * indentation if indentation != 0 else '    '

    if STYLE_DOCSTRING == 'Google':
        content = writeFunctionCommentLines_google(content, index, function, indent)
        return content

    if STYLE_DOCSTRING == 'reStructuredText':
        content = writeFunctionCommentLines_reStructuredText(content, index, function, indent)
        return content

    if STYLE_DOCSTRING == 'Numpy':
        content = writeFunctionCommentLines_numpy(content, index, function, indent)
        return content

    if STYLE_DOCSTRING == 'Epytext':
        content = writeFunctionCommentLines_epytext(content, index, function, indent)
        return content

def writeFunctionComment(filename: str, function: list):
    """ Write the comment block for a function

    Parameters
    ----------
    filename : str
        Name of the file to write the comment block
    function : str
        Function name, arguments and return type
    
    Returns
    -------
    None
    """
    with open(filename, 'r') as file:
        content = file.readlines()

    for i, line in enumerate(content):
        if line.strip().startswith((f'def {function[0]}', f'async def {function[0]}')):
            indentation = len(line) - len(line.lstrip())
            content = writeFunctionCommentLines(content, i + 1, function, indentation)
            break

    with open(filename, 'w') as file:
        file.writelines(content)

def generateAllComments(filename: str):
    """ Write the comment block for all functions in a file

    Parameters
    ----------
    filename : str
        Name of the file to write the comment blocks

    Returns
    -------
    None 
    """
    getOptions()
    functions = getFunctions(filename)
    for function in functions:
        writeFunctionComment(filename, function)