#!/usr/bin/env python

import sys, os, re

# Languages mapping as used by markdown/pandoc
shortname2language = dict(
    c='C',
    cpp='Cpp',
    f='Fortran',
    html='HTML',
    js='JavaScript',
    r='R',
    rb='Ruby',
    pl='Perl',
    py='Python',
    sh='Bash',
    tex='Tex',
    )

def read(text, argv=sys.argv[3:]):
    lines = text.splitlines()
    # First read all include statements
    for i in range(len(lines)):
        if lines[i].startswith('#include "'):
            filename = lines[i].split('"')[1]
            with open(filename, 'r') as f:
                include_text = f.read()
            lines[i] = include_text
    text = '\n'.join(lines)

    # Run Mako
    mako_kwargs = {}
    for arg in argv:
        key, value = arg.split('=')
        mako_kwargs[key] = value

    try:
        import mako
        has_mako = True
    except ImportError:
        print('Cannot import mako - mako is not run')
        has_mako = False

    if has_mako:
        from mako.template import Template
        from mako.lookup import TemplateLookup
        lookup = TemplateLookup(directories=[os.curdir])
#        text = text.encode('utf-8')
        temp = Template(text=text, lookup=lookup, strict_undefined=True)
        text = temp.render(**mako_kwargs)

    # Parse the cells
    lines = text.splitlines()
    cells = []
    inside = None    # indicates which type of cell we are inside
    fullname = None  # full language name in code cells
    for line in lines:
        if line.startswith('-----'):
            # New cell, what type?
            m = re.search(r'-----([a-z0-9-]+)?', line)
            if m:
                shortname = m.group(1)
                if shortname:
                    # Check if code is to be typeset as static
                    # Markdown code (e.g., shortname=py-t)
                    astext = shortname[-2:] == '-t'
                    if astext:
                        # Markdown
                        shortname = shortname[:-2]
                        inside = 'markdown'
                        cells.append(['markdown', 'code', ['\n']])
                        cells[-1][2].append('```%s\n' % fullname)
                    else:
                        # Code cell
                        if shortname in shortname2language:
                            fullname = shortname2language[shortname]
                        inside = 'codecell'
                        cells.append(['codecell', fullname, []])
                else:
                    # Markdown cell
                    inside = 'markdown'
                    cells.append(['markdown', 'text', ['\n']])
            else:
                raise SyntaxError('Wrong syntax of cell delimiter:\n{}'
                    .format(repr(line)))
        else:
            # Ordinary line in a cell
            if inside in ('markdown', 'codecell'):
                cells[-1][2].append(line)
            else:
                raise SyntaxError('line\n {}\nhas no beginning cell delimiter'
                    .format(line))
    # Merge the lines in each cell to a string
    for i in range(len(cells)):
        if cells[i][0] == 'markdown' and cells[i][1] == 'code':
            # Add an ending ``` of code
            cells[i][2].append('```\n')
        cells[i][2] = '\n'.join(cells[i][2])
    return cells

def write(cells):
    """Turn cells list into valid IPython notebook code."""
    # Use Jupyter nbformat functionality for writing the notebook

    from nbformat.v4 import (
        new_code_cell, new_markdown_cell, new_notebook, writes)
    nb_cells = []

    for cell_tp, language, block in cells:
        if cell_tp == 'markdown':
            nb_cells.append(
                new_markdown_cell(source=block))
        elif cell_tp == 'codecell':
            nb_cells.append(new_code_cell(source=block))

    nb = new_notebook(cells=nb_cells)
    filestr = writes(nb)
    return filestr

def driver():
    """Compile a document and its variables."""
    try:
        inputfile = sys.argv[1]
        with open(inputfile, 'r') as f:
            text = f.read()
        # Assuming file extension .gj (generate Jupyter); TODO: less strict
        outputfile = inputfile[:-3]+'.ipynb' if (len(sys.argv)<=2 or sys.argv[2]=='-') \
            else sys.argv[2]
    except (IndexError, IOError) as e:
        print('Usage: %s inputfile [outputfile|- [Mako args]]' % (sys.argv[0]))
        print(e)
        sys.exit(1)
    cells = read(text, argv=sys.argv[3:])
    filestr = write(cells)
    with open(outputfile, 'w') as f:
        f.write(filestr)

if __name__ == '__main__':
    driver()