"""
To Do:
 - Add system for rewriting docstrings
 - Convert to using click
"""

import ast, os
import openai
import sys, getopt
from query_constants import INDENT, FUNC_EXAMPLE, CLASS_EXAMPLE

from typing import List, Tuple
from dataclasses import dataclass

HELP_STRING = """docstr_maker.py

Adds Python docstrings to functions
Appends _ds to output filename by default

Usage:
    python docstr_maker.py -i <input_file>

Optional args:
 -o <output_file>
 -v Verbose
 -s Use input file as output file
    Incompatable when used with -o
 -h Displays help message

Example:
    python docstr_maker.py -i <input_file> -s
    python docstr_maker.py -i <input_file>
    python docstr_maker.py -i <input_file> -o <output_file> -v
"""


@dataclass
class CodeDefHolder:
    """
    A dataclass to hold the code definition of a function or class.

    Attributes:
        name (str):
            Name of the function or class
        unindented (str):
            The code definition of the function or class, unindented
        indented (str):
            The code definition of the function or class, indented
        node (ast.stmt):
            The node object of the function or class
        generated_docstring (str):
            The docstring generated by the code definition
        single_indent (str):
            The indentation used for the indented code definition
        full_indent (str):
            The indentation used for the indented code definition, including
            the single_indent
    """

    name: str = None
    unindented: str = None
    indented: str = None
    node: ast.stmt = None
    generated_docstring: str = None
    single_indent: str = ""
    full_indent: str = ""


def vprint(value: str, verbose: bool):
    """
    Prints a value only if verbose is true.

    Parameters:
        value (str): A string to be printed
        verbose (bool): A boolean value

    Returns:
        None
    """
    if verbose:
        print(value)


def process_input(argv: List) -> Tuple[str, str, bool]:
    """
    Processes the command line arguments and returns the input and output files.

    Parameters:
        argv (List): A list of command line arguments

    Returns:
        input_file (str): The input file
        output_file (str): The output file
        verbose (bool): Whether or not to print verbose output
    """
    input_file = None
    output_file = None
    verbose = False
    same_file = False

    try:
        opts, args = getopt.gnu_getopt(argv, "hsvi:o:", [])
    except getopt.GetoptError:
        print(HELP_STRING)
        sys.exit(2)

    for opt, arg in opts:
        if opt == "-h":
            print(HELP_STRING)
            sys.exit(2)

        elif opt in ("-i"):
            if os.path.isfile(arg):
                input_file = arg
            else:
                print("Invalid input file")
                sys.exit(2)

        elif opt in ("-o"):
            output_file = arg

        elif opt in ("-s"):
            same_file = True

        elif opt in ("-v"):
            verbose = True

    if not input_file:
        print("No file entered")
        sys.exit(2)

    if not output_file and same_file:
        output_file = input_file

    elif output_file and same_file:
        print("Error: -s (same file specified with output_file")
        sys.exit(2)

    elif not output_file and not same_file:
        period = input_file.rfind(".")
        output_file = input_file[:period] + "_ds" + input_file[period:]

    return (input_file, output_file, verbose)


def get_node_indent(node: ast.stmt, file_code: str) -> Tuple[str, str]:
    """
    Returns the indentation of a node in a file.

    Parameters:
        node (ast.stmt): A node in a file
        file_code (str): The code of the file

    Returns:
        single_indent (str): The indentation of the node
        full_indentation (str): The indentation of the node and all its children
    """
    child = node.body[0]

    start_col = node.col_offset
    end_col = child.col_offset
    row = child.lineno

    r = file_code.splitlines()[row - 1]

    single_indent = r[start_col:end_col]
    full_indentation = r[:end_col]

    return single_indent, full_indentation


def get_body_offset(node: CodeDefHolder) -> int:
    """
    Returns the number of lines between the start of the function and the start of the body

    Parameters:
        node (CodeDefHolder): A CodeDefHolder object

    Returns:
        row (int): The number of lines between the start of the function and the start of the body
    """
    start_row = node.node.lineno

    child = node.node.body[0]

    child_row = child.lineno

    row = child_row - start_row

    while node.unindented.splitlines()[row - 1].strip() == "":
        row -= 1

    return row


def isolate_nodes(file_code: str) -> List[CodeDefHolder]:
    """
    Returns a list of CodeDefHolder objects, each one representing a function or class
    definition in the file.

    Parameters:
        file_code (str): The code of the file to be analyzed

    Returns:
        ans (List[CodeDefHolder]): A list of CodeDefHolder objects
    """
    tree = ast.parse(file_code)
    ans = []

    for stmt in ast.walk(tree):
        if isinstance(stmt, (ast.FunctionDef, ast.ClassDef)):
            name = stmt.name
            unindented = ast.get_source_segment(file_code, stmt, padded=False)
            indented = ast.get_source_segment(file_code, stmt, padded=True)

            node = CodeDefHolder(name, unindented, indented, stmt)
            node.single_indent, node.full_indent = get_node_indent(stmt, file_code)

            for decorator in stmt.decorator_list:
                dec = ""
                if isinstance(decorator, ast.Name):
                    dec = "@" + decorator.id + "\n"

                elif isinstance(decorator, ast.Attribute):
                    d = decorator
                    li = []
                    while isinstance(d, ast.Attribute):
                        li.insert(0, d.attr)
                        d = d.value
                    li.insert(0, d.id)

                    dec = "@" + ".".join(li) + "\n"

                unindented = dec + unindented

                offset = node.full_indent.replace(node.single_indent, "", 1)

                indented = offset + dec + indented

            node.unindented, node.indented = unindented, indented

            ans.append(node)

    return ans


def get_codex_query(query: str) -> str:
    """
    Returns a string of code from the OpenAI API.

    Parameters:
        query (str): A string of code

    Returns:
        ds (str): A string of code
    """
    response = openai.Completion.create(
        engine="davinci-codex",
        prompt=query,
        temperature=0.1,
        max_tokens=1000,
        top_p=1.0,
        frequency_penalty=0.1,
        user="$007",
        stop=['"""', "def "],
    )

    ds = response.choices[0]["text"]

    return ds


def format_docstring(node: CodeDefHolder, docstring: str, indent: str):
    """
    Returns a formatted docstring for a function, class, or method.

    Parameters:
        node (CodeDefHolder): The node to which the docstring belongs
        docstring (str): The docstring to be formatted
        indent (str): The indentation to be used

    Returns:
        ds (str): The formatted docstring
    """
    ds = indent + '"""\n'
    ds += indent + docstring

    if ds[-1] != "\n":
        ds += "\n"

    ds += indent + '"""\n'

    ds = [line.replace(indent, "", 1) for line in ds.splitlines()]

    ds = [node.full_indent + line for line in ds]

    # Remove indentation from empty lines
    ds = [line.rstrip(node.single_indent) for line in ds]

    # rejoin docstring with newline characters between lines
    ds = "\n".join(ds)

    ds = ds.replace(indent, node.single_indent)

    return ds


def get_docstring(block: CodeDefHolder) -> str:
    """
    Returns the docstring of a code block using the codex query.

    Parameters:
        block (CodeDefHolder): A code block to be documented

    Returns:
        docstring (str): The docstring of the code block
    """
    query = None
    if isinstance(block.node, ast.FunctionDef):
        query = f'''{FUNC_EXAMPLE}
{block.unindented.replace(block.single_indent, INDENT)}
"""
{INDENT}'''

    elif isinstance(block.node, ast.ClassDef):
        query = f'''{CLASS_EXAMPLE}
### CLASS BEGIN ###
{block.unindented.replace(block.single_indent, INDENT)}

### DOCSTRING BEGIN ###
"""
{INDENT}'''

    else:
        print("ERROR!")

    docstring = get_codex_query(query)

    return format_docstring(block, docstring, INDENT)


def main():
    """
    Main function to run the program.
    """
    fpath, opath, verbose = process_input(sys.argv[1:])

    with open(fpath) as file:
        file_code = file.read()

    blocks = isolate_nodes(file_code)

    docstring_count = 0

    for i, block in enumerate(blocks):
        if ast.get_docstring(block.node):
            vprint(f"{block.name} already has a docstring", verbose)
            continue

        vprint(f"{i+1}/{len(blocks)} Generating docstring for {block.name}", verbose)

        block.generated_docstring = get_docstring(block)

        docstring_count += 1

    print(f"Generated {docstring_count} docstrings")

    for block in blocks:
        if block.generated_docstring:
            edited_snip = block.indented.splitlines()

            row = get_body_offset(block)

            decorator_count = len(block.node.decorator_list)

            edited_snip.insert(row + decorator_count, block.generated_docstring)

            edited_snip = "\n".join(edited_snip)

            file_code = file_code.replace(block.indented, edited_snip)

    with open(opath, "w+") as file:
        file.write(file_code)

    print("Successfully altered file and wrote to", opath)


openai.api_key = os.getenv("OPENAI_API_KEY")

if __name__ == "__main__":
    main()
