"""A convenience tool that helps to convert Jupyter notebook into plain python
scripts.

The major function of this module is called `to_py`. The function takes path
to the folder with notebooks (or a single notebook name), and converts them
into a bunch of python files. Only the cells marked with `#export` are
considered.

With this tool, you can easily separate your testing and mocking code from
the production-ready implementations and export into .py files only relevant
cells without need to clean-up the notebook at first.

The script can optionally strip notebooks names to exclude the numerical
prefixes at the beginning of file names. For example, you can structure your
code in the following way:

/notebooks
    01_prepare.ipynb
    02_train.ipynb
    03_visualize.ipynb

Then the script will go though the notebooks in accordance with their order,
and saves everything into output folder.

output:
    prepare.py
    train.py
    visualize.py

"""
import json
import re
from operator import itemgetter
from pathlib import Path
from pprint import pprint
import shutil


__all__ = ['to_py']


HEADER = """
# -----------------------------------------
# THIS FILE WAS AUTOGENERATED! DO NOT EDIT!
# -----------------------------------------
# file to edit: {edit_file}

"""


def export(cell: dict, min_size: int = 7, types: tuple = ('code',)) -> bool:
    """Determines if a cell should be exported into source file.

    The function inspects a content of a cell and returns True only if the
    following conditions are met:
        (1) a cell is not empty;
        (2) a cell contains '#export' comment as its first line;
        (3) a cell is 'big enough' (more then a few chars).

    Args:
        cell: A Jupyter cell represented as a JSON object.
        min_size: Minimal length of cell's source code to consider it as being
            ready for export.
        types: Types of cells to export. By default, only 'code' cells are
            exported.

    """
    if cell['cell_type'] not in types:
        return False
    if not cell['source']:
        return False
    src = cell['source'][0]
    if not src or len(src) < min_size:
        return False
    match = re.match(r'^\s*#\s*export\s*$', src, re.IGNORECASE)
    return match is not None


def find_notebooks(folder, ordered=True, pattern='*.ipynb', exclude=None):
    """Goes through the folder and converts notebooks into Python scripts

    If ordered is True, then each notebook's name should start with a number.
    Otherwise, a notebook is ignored.
    """
    p = Path(folder)
    notebooks = list(p.glob(pattern))
    if ordered:
        items = []
        for filename in notebooks:
            m = re.match(r'^(\d+)([a-z]+)?_', filename.stem, re.IGNORECASE)
            if m is None:
                continue
            items.append((m.groups(), filename))
        items.sort(key=itemgetter(0))
        notebooks = [fn for _, fn in items]
    if exclude:
        notebooks = [fn for fn in notebooks if fn not in exclude]
    return notebooks


def export_notebooks(filenames: list, output: str = None, gap: int = 2,
                     name_as_is: bool = False):
    """Convert list of notebooks into .py files.

    Args:
        filenames: Notebooks files to export.
        output: A folder to export the notebooks.
        gap: Number of empty lines between exported cells.
        name_as_is: Don't strip numerical prefix from notebook names.

    """
    count = 0
    empty_lines = "\n" * gap
    for fn in filenames:
        fn = Path(fn)
        content = json.load(fn.open('r'))
        module = HEADER.format(edit_file=fn.name)
        for cell in (c for c in content['cells'] if export(c)):
            module += f"{''.join(cell['source'][1:])}{empty_lines}"
        module = re.sub(r'[ \t]+$', '', module, flags=re.MULTILINE)
        if name_as_is:
            pyfile = output / f'{fn.stem}.py'
        else:
            pyfile = output / f"{''.join(fn.stem.partition('_')[2:])}.py"
        with pyfile.open('w') as file:
            file.write(module.strip() + "\n")
        print(f'Exported: {fn} -> {pyfile}')
        count += 1
    return count


def copy_dependencies(deps, output):
    """Goes through the collection of files and copies them to the folder with
    exported notebooks.
    """
    output = Path(output)
    for path, alias in deps.items():
        p = Path(path)
        if not p.exists():
            print(f'Warning! File not found: {p}')
            continue
        shutil.copy(p, output / alias)


def to_py(
    source_dir: str = None,
    deps: dict = None,
    keep_names: bool = True,
    nb_name: str = None,
    exclude_notebooks: list = None,
    output_dir: str = None,
    gap: int = 3
) -> tuple:
    """Converts Jupyter notebook(s) into plain .py files.

    Args:
        source_dir: A directory to search notebooks. (Default: workdir)
        deps: A dictionary with additional files that should be copied into
            the output folder. Keys are original paths, values are aliases
            referred from the source code. For example, consider the following
            structure:
            {
                '../path/to/utility.py': 'extra_utils.py'
            }
            The file 'utility.py' will be copied to the output folder and
            renamed into 'jupyter_utils.py'
        keep_names: If True, then the names of exported files match to the names
            of the notebooks.
        nb_name: A name of notebook file. If provided, then only this specific
            notebook is exported.
        exclude_notebooks: List of notebooks names to exclude from export.
        output_dir: A directory to save the exported notebooks.
        gap: A number of new lines between exported cells.

    Returns:
         tuple (notebooks, py_files): The export results.

    """
    if nb_name is None:
        source_dir = Path(source_dir or Path.cwd())
        ordered = not keep_names
        notebooks = find_notebooks(
            source_dir, ordered=ordered,
            exclude=exclude_notebooks)
        print(f'Notebooks directory: {source_dir}')
        print('Automatically discovered notebooks:')
        pprint(notebooks)
    else:
        notebooks = [nb_name]
    output_dir = Path(output_dir or Path.cwd())
    if not output_dir.exists():
        output_dir.mkdir(parents=True)

    n = export_notebooks(notebooks, output_dir, gap, name_as_is=keep_names)
    if deps is not None:
        copy_dependencies(deps, output_dir)
    py_files = list(output_dir.iterdir())
    print(f'{n} notebook(s) exported into folder: {output_dir}')
    return notebooks, py_files
