"""Read a Python file into string in IPython notebook form."""

from __future__ import absolute_import, division, print_function

import sys
import tokenize
from io import StringIO

if sys.version_info[0] == 2:
    # Python 2.x
    from collections import namedtuple
    TokenInfo = namedtuple('TokenInfo', 'type string start end line')

    def _generate_tokens(readline):
        return map(lambda x: TokenInfo(*x), tokenize.generate_tokens(readline))

else:
    # Python 3.x
    TokenInfo = tokenize.TokenInfo
    _generate_tokens = tokenize.tokenize



def read_python(fin):
    """
    Read a regular Python file with special formatting, and perform
    preprocessing on it.  The result is a string that conforms to the
    IPython notebook version 3 Python script format.
    """

    token_gen = _generate_tokens(fin.readline)
    cvt_docstr_gen = convert_toplevel_docstring(token_gen)
    nl_gen = fix_newlines(cvt_docstr_gen)
    out = list(nl_gen)
    formatted = tokenize.untokenize(out).decode('utf-8')

    return fix_empty_lines(formatted)


# ======================================================================
#                                   Helpers
# ======================================================================

def convert_toplevel_docstring(tokens):
    for token in tokens:
        # For each string
        if token.type == tokenize.STRING:
            text = token.string
            # Must be a docstring
            if text.startswith('"""') or text.startswith("'''"):
                startline, startcol = token.start
                # Starting column MUST be 0
                if startcol == 0:
                    endline, endcol = token.end
                    lines = ['# ' + line
                             for line in text.strip('"\' \n').split('\n')]
                    text = '\n'.join(lines)
                    fmt = '# <markdowncell>\n{0}\n# <codecell>'.format(text)
                    yield TokenInfo(type=tokenize.COMMENT,
                                    start=(startline, startcol),
                                    end=(endline, endcol),
                                    string=fmt,
                                    line='#')
                    # To next token
                    continue
        # Return untouched
        yield token


def fix_newlines(tokens):
    first = True
    curline = 1
    for token in tokens:
        if first:
            first = False
            curline = token.end[0] + 1
        else:
            # Fill NEWLINE token in between
            while curline < token.start[0]:
                yield TokenInfo(type=tokenize.NEWLINE,
                                string='\n',
                                start=(curline, 0),
                                end=(curline, 0),
                                line='\n', )
                curline += 1

            curline = token.end[0] + 1
        yield token


def fix_empty_lines(text):
    def gen():
        for line in text.splitlines():
            if not line.strip():
                # Empty line
                yield ''
            else:
                yield line

    return '\n'.join(gen())
