"""
A collection of functions for modifying source code that is embeded into the Sphinx documentation.
"""
from __future__ import print_function

import sys
import os
import re
import tokenize
import importlib
import inspect
import subprocess
import tempfile
import unittest
import traceback
from docutils import nodes

from six import StringIO
from six.moves import range, zip, cStringIO as cStringIO

from sphinx.errors import SphinxError
from sphinx.writers.html import HTMLTranslator
from redbaron import RedBaron

if sys.version_info[0] == 2:
    import cgi as cgiesc
else:
    import html as cgiesc

from openmdao.utils.general_utils import printoptions

sqlite_file = 'feature_docs_unit_test_db.sqlite'    # name of the sqlite database file
table_name = 'feature_unit_tests'   # name of the table to be queried

_sub_runner = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'run_sub.py')


class skipped_or_failed_node(nodes.Element):
    pass


def visit_skipped_or_failed_node(self, node):
    pass


def depart_skipped_or_failed_node(self, node):
    if not isinstance(self, HTMLTranslator):
        self.body.append("output only available for HTML\n")
        return

    html = '<div class="cell border-box-sizing code_cell rendered"><div class="output"><div class="inner_cell"><div class="{}"><pre>{}</pre></div></div></div></div>'.format(node["kind"], node["text"])
    self.body.append(html)


class in_or_out_node(nodes.Element):
    pass


def visit_in_or_out_node(self, node):
    pass


def depart_in_or_out_node(self, node):
    """
    This function creates the formatting that sets up the look of the blocks.
    The look of the formatting is controlled by _theme/static/style.css
    """
    if not isinstance(self, HTMLTranslator):
        self.body.append("output only available for HTML\n")
        return
    if node["kind"] == "In":
        html = '<div class="highlight-python"><div class="highlight"><pre>{}</pre></div></div>'.format(node["text"])
    elif node["kind"] == "Out":
        html = '<div class="cell border-box-sizing code_cell rendered"><div class="output_area"><pre>{}</pre></div></div>'.format(node["text"])

    self.body.append(html)


def node_setup(app):
    app.add_node(skipped_or_failed_node, html=(visit_skipped_or_failed_node, depart_skipped_or_failed_node))
    app.add_node(in_or_out_node, html=(visit_in_or_out_node, depart_in_or_out_node))


def remove_docstrings(source):
    """
    Return 'source' minus docstrings.

    Parameters
    ----------
    source : str
        Original source code.

    Returns
    -------
    str
        Source with docstrings removed.
    """
    io_obj = StringIO(source)
    out = ""
    prev_toktype = tokenize.INDENT
    last_lineno = -1
    last_col = 0
    for tok in tokenize.generate_tokens(io_obj.readline):
        token_type = tok[0]
        token_string = tok[1]
        start_line, start_col = tok[2]
        end_line, end_col = tok[3]
        # ltext = tok[4] # in original code but not used here
        # The following two conditionals preserve indentation.
        # This is necessary because we're not using tokenize.untokenize()
        # (because it spits out code with copious amounts of oddly-placed
        # whitespace).
        if start_line > last_lineno:
            last_col = 0
        if start_col > last_col:
            out += (" " * (start_col - last_col))
        # This series of conditionals removes docstrings:
        if token_type == tokenize.STRING:
            if prev_toktype != tokenize.INDENT:
                # This is likely a docstring; double-check we're not inside an operator:
                if prev_toktype != tokenize.NEWLINE:
                    # Note regarding NEWLINE vs NL: The tokenize module
                    # differentiates between newlines that start a new statement
                    # and newlines inside of operators such as parens, brackes,
                    # and curly braces.  Newlines inside of operators are
                    # NEWLINE and newlines that start new code are NL.
                    # Catch whole-module docstrings:
                    if start_col > 0:
                        # Unlabelled indentation means we're inside an operator
                        out += token_string
                    # Note regarding the INDENT token: The tokenize module does
                    # not label indentation inside of an operator (parens,
                    # brackets, and curly braces) as actual indentation.
                    # For example:
                    # def foo():
                    #     "The spaces before this docstring are tokenize.INDENT"
                    #     test = [
                    #         "The spaces before this string do not get a token"
                    #     ]
        else:
            out += token_string
        prev_toktype = token_type
        last_col = end_col
        last_lineno = end_line
    return out


def remove_redbaron_node(node, index):
    """
    Utility function for removing a node using RedBaron.

    RedBaron has some problems with modifying code lines that run across
    multiple lines. ( It is mentioned somewhere online but cannot seem to
    find it now. )

    RedBaron throws an Exception but when you check, it seems like it does
    what you asked it to do. So, for now, we ignore the Exception.
    """

    try:
        node.value.remove(node.value[index])
    except Exception as e:  # no choice but to catch the general Exception
        if str(e).startswith('It appears that you have indentation in your CommaList'):
            pass
        else:
            raise


def replace_asserts_with_prints(src):
    """
    Replace asserts with print statements.

    Using RedBaron, replace some assert calls with print statements that print the actual
    value given in the asserts. Depending on the calls, the actual value can be the first or second
    argument.

    Parameters
    ----------
    src : str
        String containing source lines.

    Returns
    -------
    str
        String containing source with asserts replaced by prints.
    """
    rb = RedBaron(src)  # convert to RedBaron internal structure

    # findAll is slow, so only check the ones that are present.
    base_assert = ['assertAlmostEqual', 'assertLess', 'assertGreater', 'assertEqual',
                   'assert_equal_arrays', 'assertTrue', 'assertFalse']
    used_assert = [item for item in base_assert if item in src]

    for assert_type in used_assert:
        assert_nodes = rb.findAll("NameNode", value=assert_type)
        for assert_node in assert_nodes:
            assert_node = assert_node.parent
            remove_redbaron_node(assert_node, 0)  # remove 'self' from the call
            assert_node.value[0].replace('print')
            if assert_type not in ['assertTrue', 'assertFalse']:
                # remove the expected value argument
                remove_redbaron_node(assert_node.value[1], 1)

    if 'assert_rel_error' in src:
        assert_nodes = rb.findAll("NameNode", value='assert_rel_error')
        for assert_node in assert_nodes:
            assert_node = assert_node.parent
            # If relative error tolerance is specified, there are 4 arguments
            if len(assert_node.value[1]) == 4:
                # remove the relative error tolerance
                remove_redbaron_node(assert_node.value[1], -1)
            remove_redbaron_node(assert_node.value[1], -1)  # remove the expected value
            # remove the first argument which is the TestCase
            remove_redbaron_node(assert_node.value[1], 0)
            #
            assert_node.value[0].replace("print")

    if 'assert_almost_equal' in src:
        assert_nodes = rb.findAll("NameNode", value='assert_almost_equal')
        for assert_node in assert_nodes:
            assert_node = assert_node.parent
            # If relative error tolerance is specified, there are 3 arguments
            if len(assert_node.value[1]) == 3:
                # remove the relative error tolerance
                remove_redbaron_node(assert_node.value[1], -1)
            remove_redbaron_node(assert_node.value[1], -1)  # remove the expected value
            assert_node.value[0].replace("print")

    return rb.dumps()


def remove_initial_empty_lines(source):
    """
    Some initial empty lines were added to keep RedBaron happy.
    Need to strip these out before we pass the source code to the
    directive for including source code into feature doc files.
    """

    idx = re.search(r'\S', source, re.MULTILINE).start()
    return source[idx:]


def get_source_code(path):
    """
    Return source code as a text string.

    Parameters
    ----------
    path : str
        Path to a file, module, function, class, or class method.

    Returns
    -------
    str
        The source code.
    int
        Indentation level.
    module or None
        The imported module.
    class or None
        The class specified by path.
    """

    indent = 0
    cls = None

    if path.endswith('.py'):
        if not os.path.isfile(path):
            raise SphinxError("Can't find file '%s' cwd='%s'" % (path, os.getcwd()))
        with open(path, 'r') as f:
            source = f.read()
        module = None
    else:
        # First, assume module path since we want to support loading a full module as well.
        try:
            module = importlib.import_module(path)
            source = inspect.getsource(module)

        except ImportError:

            # Second, assume class and see if it works
            try:
                parts = path.split('.')

                module_path = '.'.join(parts[:-1])
                module = importlib.import_module(module_path)
                class_name = parts[-1]
                cls = getattr(module, class_name)
                source = inspect.getsource(cls)
                indent = 1

            except ImportError:

                # else assume it is a path to a method
                module_path = '.'.join(parts[:-2])
                module = importlib.import_module(module_path)
                class_name = parts[-2]
                method_name = parts[-1]
                cls = getattr(module, class_name)
                meth = getattr(cls, method_name)
                source = inspect.getsource(meth)
                indent = 2

    return remove_leading_trailing_whitespace_lines(source), indent, module, cls


def remove_raise_skip_tests(src):
    """
    Remove from the code any raise unittest.SkipTest lines since we don't want those in
    what the user sees.
    """
    rb = RedBaron(src)
    raise_nodes = rb.findAll("RaiseNode")
    for rn in raise_nodes:
        # only the raise for SkipTest
        if rn.value[:2].dumps() == 'unittestSkipTest':
            rn.parent.value.remove(rn)
    return rb.dumps()


def remove_leading_trailing_whitespace_lines(src):
    """
    Remove any leading or trailing whitespace lines.

    Parameters
    ----------
    src : str
        Input code.

    Returns
    -------
    str
        Code with trailing whitespace lines removed.
    """
    lines = src.splitlines()

    non_whitespace_lines = []
    for i, l in enumerate(lines):
        if l and not l.isspace():
            non_whitespace_lines.append(i)
    imin = min(non_whitespace_lines)
    imax = max(non_whitespace_lines)

    return '\n'.join(lines[imin: imax+1])


def is_output_node(node):
    """
    Determine whether a RedBaron node may be expected to generate output.

    Parameters
    ----------
    node : <Node>
        a RedBaron Node.

    Returns
    -------
    bool
        True if node may be expected to generate output, otherwise False.
    """
    if node.type == 'print':
        return True

    # lines with the following signatures and function names may generate output
    output_signatures = [
        ('name', 'name', 'call'),
        ('name', 'name', 'name', 'call')
    ]
    output_functions = [
        'setup', 'run_model', 'run_driver',
        'check_partials', 'check_totals',
        'list_inputs', 'list_outputs', 'list_problem_vars'
    ]

    if node.type == 'atomtrailers' and len(node.value) in (3, 4):
        sig = []
        for val in node.value:
            sig.append(val.type)
        func_name = node.value[-2].value
        if tuple(sig) in output_signatures and func_name in output_functions:
            return True

    return False


def split_source_into_input_blocks(src):
    """
    Split source into blocks; the splits occur at inserted prints.

    Parameters
    ----------
    src : str
        Input code.

    Returns
    -------
    list
        List of input code sections.
    """
    input_blocks = []

    current_block = []

    for line in src.split('\n'):
        if 'print(">>>>>' in line:
            input_blocks.append('\n'.join(current_block))
            current_block = []
        else:
            current_block.append(line)

    if current_block and current_block[0]:
        input_blocks.append('\n'.join(current_block))

    return input_blocks


def insert_output_start_stop_indicators(src):
    """
    Insert identifier strings so that output can be segregated from input.

    Parameters
    ----------
    src : str
        String containing input and output lines.

    Returns
    -------
    str
        String with output demarked.
    """
    lines = src.split('\n')
    print_producing = ['.setup(',
                       'print(',
                       '.run_model(',
                       '.run_driver(',
                       '.check_partials(',
                       '.check_totals(',
                       '.list_inputs(',
                       '.list_outputs(',
                       '.list_problem_vars(',
                       ]

    newlines = []
    input_block_number = 0
    in_try = False
    in_continuation = False
    head_indent = ''
    for line in lines:
        newlines.append(line)

        # Check if we are concluding a continuation line.
        if in_continuation:
            line = line.rstrip()
            if not (line.endswith(',') or line.endswith('\\') or line.endswith('(')):
                newlines.append('%sprint(">>>>>%d")' % (head_indent, input_block_number))
                input_block_number += 1
                in_continuation = False

        # Don't print if we are in a try block.
        if in_try:
            if 'except' in line:
                in_try = False
            continue

        if 'try:' in line:
            in_try = True
            continue

        # Searching for 'print(' is a little ambiguous.
        if 'set_solver_print(' in line:
            continue

        for item in print_producing:
            if item in line:
                indent = ' ' * (len(line) - len(line.lstrip()))

                # Line continuations are a litle tricky.
                line = line.rstrip()
                if line.endswith(',') or line.endswith('\\') or line.endswith('('):
                    in_continuation = True
                    head_indent = indent
                    break

                newlines.append('%sprint(">>>>>%d")' % (indent, input_block_number))
                input_block_number += 1
                break

    return '\n'.join(newlines)


def clean_up_empty_output_blocks(input_blocks, output_blocks):
    """Some of the blocks do not generate output. We only want to have
    input blocks that have outputs.
    """

    new_input_blocks = []
    new_output_blocks = []
    current_in_block = ''

    for in_block, out_block in zip(input_blocks, output_blocks):
        if current_in_block and not current_in_block.endswith('\n'):
            current_in_block += '\n'
        current_in_block += in_block
        if out_block:
            current_in_block = remove_leading_trailing_whitespace_lines(current_in_block)
            out_block = remove_leading_trailing_whitespace_lines(out_block)
            new_input_blocks.append(current_in_block)
            new_output_blocks.append(out_block)
            current_in_block = ''

    # if there was no output, return the one input block and empty output block
    if current_in_block:
        new_input_blocks.append(current_in_block)
        new_output_blocks.append('')

    return new_input_blocks, new_output_blocks


def extract_output_blocks(run_output):
    """
    Identify and extract outputs from source.

    Parameters
    ----------
    run_output : str or list of str
        Source code with outputs.

    Returns
    -------
    list of str
        List containing output text blocks.
    """

    # Look for start and end lines that look like this:
    #  <<<<<4
    #  >>>>>4

    if isinstance(run_output, list):
        return sync_multi_output_blocks(run_output)

    output_blocks = []
    output_block = None

    for line in run_output.splitlines():
        if output_block is None:
            output_block = []
        if line[:5] == '>>>>>':
            output_blocks.append('\n'.join(output_block))
            output_block = None
        else:
            output_block.append(line)

    if output_block is not None:
        output_blocks.append('\n'.join(output_block))

    return output_blocks


def strip_header(src):
    """
    Directly manipulating function text to strip header.

    This function assumes that the docstring and header, if any, have already been removed.

    Parameters
    ----------
    src : str
        sourec code for method
    """
    lines = src.split('\n')
    first_len = None
    for i, line in enumerate(lines):
        n1 = len(line)
        newline = line.lstrip()
        tab = n1 - len(newline)
        if first_len is None:
            first_len = tab
        elif n1 == 0:
            continue
        if tab != first_len:
            return '\n'.join(lines[i:])

    return ''


def dedent(src):
    """
    Directly manipulating function text to remove leading whitespace.

    Parameters
    ----------
    src : str
        sourec code for method
    """

    lines = src.split('\n')
    if lines:
        for i, line in enumerate(lines):
            lstrip = line.lstrip()
            if lstrip: # keep going if first line(s) are blank.
                tab = len(line) - len(lstrip)
                return '\n'.join(l[tab:] for l in lines[i:])
    return ''


def sync_multi_output_blocks(blocks):
    """
    Combine output from different procs into the same output blocks.

    Parameters
    ----------
    blocks : list of str
        List of outputs from individual procs.

    Returns
    -------
    list of list of str
        List of synced output blocks from all procs.
    """
    if blocks:
        split_blocks = [extract_output_blocks(b) for b in blocks]
        n_out_blocks = len(split_blocks[0])
        synced_blocks = []
        for i in range(n_out_blocks):
            synced_blocks.append('\n'.join(["(rank %d) %s" % (j, m[i])
                for j, m in enumerate(split_blocks) if m[i]]))
        return synced_blocks
    else:
        return []


def run_code(code_to_run, path, module=None, cls=None, shows_plot=False):
    """
    Run the given code chunk and collect the output.
    """

    skipped = False
    failed = False

    if cls is None:
        use_mpi = False
    else:
        try:
            import mpi4py
        except ImportError:
            use_mpi = False
        else:
            N_PROCS = getattr(cls, 'N_PROCS', 1)
            use_mpi = N_PROCS > 1

    try:
        # use subprocess to run code to avoid any nasty interactions between codes

        # Move to the test directory in case there are files to read.
        save_dir = os.getcwd()

        if module is None:
            code_dir = os.path.dirname(os.path.abspath(path))
        else:
            code_dir = os.path.dirname(os.path.abspath(module.__file__))

        os.chdir(code_dir)

        if use_mpi:
            env = os.environ.copy()

            # output will be written to one file per process
            env['USE_PROC_FILES'] = '1'

            env['OPENMDAO_CURRENT_MODULE'] = module.__name__
            env['OPENMDAO_CODE_TO_RUN'] = code_to_run

            p = subprocess.Popen(['mpirun', '-n', str(N_PROCS), 'python', _sub_runner],
                                 env=env)
            p.wait()

            # extract output blocks from all output files & merge them
            output = []
            for i in range(N_PROCS):
                with open('%d.out' % i) as f:
                    output.append(f.read())
                os.remove('%d.out' % i)

        elif shows_plot:
            if module is None:
                # write code to a file so we can run it.
                fd, code_to_run_path = tempfile.mkstemp()
                with os.fdopen(fd, 'w') as tmp:
                    tmp.write(code_to_run)
                try:
                    p = subprocess.Popen(['python', code_to_run_path],
                                         stdout=subprocess.PIPE, stderr=subprocess.STDOUT, env=os.environ)
                    output, _ = p.communicate()
                    if p.returncode != 0:
                        failed = True

                finally:
                    os.remove(code_to_run_path)
            else:
                env = os.environ.copy()

                env['OPENMDAO_CURRENT_MODULE'] = module.__name__
                env['OPENMDAO_CODE_TO_RUN'] = code_to_run

                p = subprocess.Popen(['python', _sub_runner],
                                     stdout=subprocess.PIPE, stderr=subprocess.STDOUT, env=env)
                output, _ = p.communicate()
                if p.returncode != 0:
                    failed = True

            output = output.decode('utf-8', 'ignore')
        else:
            # just exec() the code for serial tests.

            # capture all output
            stdout = sys.stdout
            stderr = sys.stderr
            strout = cStringIO()
            sys.stdout = strout
            sys.stderr = strout

            # We need more precision from numpy
            with printoptions(precision=8):

                if module is None:
                    globals_dict = {
                        '__file__': path,
                        '__name__': '__main__',
                        '__package__': None,
                        '__cached__': None,
                    }
                else:
                    globals_dict = module.__dict__

                try:
                    exec(code_to_run, globals_dict)
                except Exception:
                    # print code (with line numbers) to facilitate debugging
                    for n, line in enumerate(code_to_run.split('\n')):
                        print('%4d: %s' % (n, line), file=stderr)
                    raise
                finally:
                    sys.stdout = stdout
                    sys.stderr = stderr

            output = strout.getvalue()

    except subprocess.CalledProcessError as e:
        output = e.output.decode('utf-8', 'ignore')
        # Get a traceback.
        if 'raise unittest.SkipTest' in output:
            reason_for_skip = output.splitlines()[-1][len('unittest.case.SkipTest: '):]
            output = reason_for_skip
            skipped = True
        else:
            output = "Running of embedded code {} in docs failed due to: \n\n{}".format(path, output)
            failed = True
    except unittest.SkipTest as skip:
        output = str(skip)
        skipped = True
    except Exception as exc:
        output = "Running of embedded code {} in docs failed due to: \n\n{}".format(path, traceback.format_exc())
        failed = True
    finally:
        os.chdir(save_dir)

    return skipped, failed, output


def get_skip_output_node(output):
    output = "Test skipped because " + output
    return skipped_or_failed_node(text=output, number=1, kind="skipped")


def get_interleaved_io_nodes(input_blocks, output_blocks):
    nodelist = []
    n = 1
    output_blocks = [cgiesc.escape(ob) for ob in output_blocks]
    for input_block, output_block in zip(input_blocks, output_blocks):
        input_node = nodes.literal_block(input_block, input_block)
        input_node['language'] = 'python'
        nodelist.append(input_node)
        if len(output_block) > 0:
            output_node = in_or_out_node(kind="Out", number=n, text=output_block)
            nodelist.append(output_node)
        n += 1
    return nodelist


def get_output_block_node(output_blocks):
    output_block = '\n'.join([cgiesc.escape(ob) for ob in output_blocks])
    return in_or_out_node(kind="Out", number=1, text=output_block)
