# -*- coding: utf-8 -*-
"""
Custom docutils writer for IPython Notebooks.
"""

import os
import re
import textwrap
from itertools import groupby

from nbformat import v4 as ipynb

from six.moves import zip_longest

from docutils import nodes, writers
from docutils.utils import column_width

from sphinx import addnodes
from sphinx.locale import admonitionlabels, _

from pprint import pprint

MAXWIDTH = 80
STDINDENT = 2

class IpynbWriter(writers.Writer):
    supported = ('ipynb',)
    settings_spec = ('No options here.', '', ())
    settings_defaults = {}

    output = None

    def __init__(self, builder, kernel, indent, cont):
        writers.Writer.__init__(self)
        self.builder = builder
        self.kernel = kernel
        self.code_indent = indent
        self.code_cont = cont
        self.translator_class = self.builder.translator_class or IpynbTranslator

    def translate(self):
        visitor = self.translator_class(self.document, self.builder, self.kernel, self.code_indent, self.code_cont)
        self.document.walkabout(visitor)
        self.output = visitor.body


class IpynbTranslator(nodes.NodeVisitor):
    sectionchars = '*=-~"+`'

    def __init__(self, document, builder, kernel, indent, cont):
        nodes.NodeVisitor.__init__(self, document)
        self.builder = builder
        self.kernel = kernel
        self.code_indent = indent
        self.code_cont = cont

        self.nl = '\n\n' # newline in markdown ?
        self.sectionchars = builder.config.text_sectionchars
        self.cells = []
        self.list_counter = []
        self.sectionlevel = 1
        self.lineblocklevel = 0
        self.stateindent = [0]
        self.table = None
        self.current_state = None

    def add_text(self, text):
        self.cells[-1]["source"] += text

    def new_state(self, cell_type):
        if cell_type == "code":
            cell = ipynb.new_code_cell()
        elif cell_type == "text":
            cell = ipynb.new_markdown_cell()
        else:
            raise ValueError("Unknown cell type '%s'" % cell_type)
        self.current_state = cell_type
        self.cells.append(cell)

    def end_state(self):
        # post processing the last cell
        # break up the code cells and remove the python indentation
        # the idea is to make each cells executable such that the user discovers the result
        if self.current_state == 'code':
            code = self.cells.pop()['source']
            codeblocks = [[]]
            output_line = False
            for line in code.splitlines():
                for start in [self.code_indent, self.code_cont]:
                    if line.startswith(start):
                        if output_line == True:
                            codeblocks.append([])
                            output_line = False
                        line = line[len(start):]
                        codeblocks[-1].append(line)
                        break
                else:
                    # TODO that comment syntax is R and Python specific. Make this a configuration.
                    if len(line) > 0:
                        line = '# ' + line
                        codeblocks[-1].append(line)
                    output_line = True
            for block in codeblocks:
                if len(block) > 0:
                    self.new_state('code')
                    self.add_text('\n'.join(block))

    def visit_document(self, node):
        # self.new_state()
        pass

    def depart_document(self, node):
        #self.end_state()

        nb = ipynb.new_notebook()
        nb["cells"] = self.cells

        nb["metadata"]["kernelspec"] = {"name" : self.kernel}

        self.body = ipynb.writes_json(nb)

    def visit_highlightlang(self, node):
        raise nodes.SkipNode

    def visit_section(self, node):
        self._title_char = "#" * self.sectionlevel
        self.sectionlevel += 1

    def depart_section(self, node):
        self.sectionlevel -= 1

    def visit_topic(self, node):
        self.new_state("text")

    def depart_topic(self, node):
        self.end_state()

    visit_sidebar = visit_topic
    depart_sidebar = depart_topic

    def visit_rubric(self, node):
        self.new_state("text")
        self.add_text('-[ ')

    def depart_rubric(self, node):
        self.add_text(' ]-')
        self.end_state()

    def visit_compound(self, node):
        pass

    def depart_compound(self, node):
        pass

    def visit_glossary(self, node):
        pass

    def depart_glossary(self, node):
        pass

    def visit_title(self, node):
        if isinstance(node.parent, nodes.Admonition):
            self.add_text(node.astext()+': ')
            raise nodes.SkipNode
        self.new_state("text")
        self.add_text(self._title_char + " ")

    def depart_title(self, node):
        pass
        #if isinstance(node.parent, nodes.section):
        #    char = self._title_char
        #else:
        #    char = '^'
        #text = ''.join(x[1] for x in self.states.pop() if x[0] == -1)
        #self.cells[-1].append(
        #    (0, ['', text, '%s' % (char * column_width(text)), '']))

    def visit_subtitle(self, node):
        pass

    def depart_subtitle(self, node):
        pass

    def visit_attribution(self, node):
        self.add_text('-- ')

    def depart_attribution(self, node):
        pass

    def visit_desc(self, node):
        pass

    def depart_desc(self, node):
        pass

    def visit_desc_signature(self, node):
        self.new_state("text")

    def depart_desc_signature(self, node):
        # XXX: wrap signatures in a way that makes sense
        self.end_state()

    def visit_desc_name(self, node):
        pass

    def depart_desc_name(self, node):
        pass

    def visit_desc_addname(self, node):
        pass

    def depart_desc_addname(self, node):
        pass

    def visit_desc_type(self, node):
        pass

    def depart_desc_type(self, node):
        pass

    def visit_desc_returns(self, node):
        self.add_text(' -> ')

    def depart_desc_returns(self, node):
        pass

    def visit_desc_parameterlist(self, node):
        self.add_text('(')
        self.first_param = 1

    def depart_desc_parameterlist(self, node):
        self.add_text(')')

    def visit_desc_parameter(self, node):
        if not self.first_param:
            self.add_text(', ')
        else:
            self.first_param = 0
        self.add_text(node.astext())
        raise nodes.SkipNode

    def visit_desc_optional(self, node):
        self.add_text('[')

    def depart_desc_optional(self, node):
        self.add_text(']')

    def visit_desc_annotation(self, node):
        pass

    def depart_desc_annotation(self, node):
        pass

    def visit_desc_content(self, node):
        self.new_state("text")
        self.add_text(self.nl)

    def depart_desc_content(self, node):
        self.end_state()

    def visit_figure(self, node):
        self.new_state("text")

    def depart_figure(self, node):
        self.end_state()

    def visit_caption(self, node):
        pass

    def depart_caption(self, node):
        pass

    def visit_productionlist(self, node):
        self.new_state("text")
        names = []
        for production in node:
            names.append(production['tokenname'])
        maxlen = max(len(name) for name in names)
        lastname = None
        for production in node:
            if production['tokenname']:
                self.add_text(production['tokenname'].ljust(maxlen) + ' ::=')
                lastname = production['tokenname']
            elif lastname is not None:
                self.add_text('%s    ' % (' '*len(lastname)))
            self.add_text(production.astext() + self.nl)
        self.end_state()
        raise nodes.SkipNode

    def visit_footnote(self, node):
        self._footnote = node.children[0].astext().strip()
        self.new_state("text")

    def depart_footnote(self, node):
        self.add_text('[%s] ' % self._footnote)

    def visit_citation(self, node):
        if len(node) and isinstance(node[0], nodes.label):
            self._citlabel = node[0].astext()
        else:
            self._citlabel = ''
        self.new_state("text")

    def depart_citation(self, node):
        self.add_text('[%s] ' % self._citlabel)

    def visit_label(self, node):
        raise nodes.SkipNode

    def visit_legend(self, node):
        pass

    def depart_legend(self, node):
        pass

    # XXX: option list could use some better styling

    def visit_option_list(self, node):
        pass

    def depart_option_list(self, node):
        pass

    def visit_option_list_item(self, node):
        self.new_state("text")

    def depart_option_list_item(self, node):
        self.end_state()

    def visit_option_group(self, node):
        self._firstoption = True

    def depart_option_group(self, node):
        self.add_text('     ')

    def visit_option(self, node):
        if self._firstoption:
            self._firstoption = False
        else:
            self.add_text(', ')

    def depart_option(self, node):
        pass

    def visit_option_string(self, node):
        pass

    def depart_option_string(self, node):
        pass

    def visit_option_argument(self, node):
        self.add_text(node['delimiter'])

    def depart_option_argument(self, node):
        pass

    def visit_description(self, node):
        pass

    def depart_description(self, node):
        pass

    def visit_tabular_col_spec(self, node):
        raise nodes.SkipNode

    def visit_colspec(self, node):
        self.table[0].append(node['colwidth'])
        raise nodes.SkipNode

    def visit_tgroup(self, node):
        pass

    def depart_tgroup(self, node):
        pass

    def visit_thead(self, node):
        pass

    def depart_thead(self, node):
        pass

    def visit_tbody(self, node):
        self.table.append("sep")

    def depart_tbody(self, node):
        pass

    def visit_row(self, node):
        self.table.append([])

    def depart_row(self, node):
        pass

    def visit_entry(self, node):
        #if 'morerows' in node or 'morecols' in node:
        #    raise NotImplementedError('Column or row spanning cells are '
        #                              'not implemented.')
        pass

    def depart_entry(self, node):
        pass

    def visit_table(self, node):
        self.new_state("text")
        self.table = [[]]

    def depart_table(self, node):
        lines = self.table[1:]
        rows = []
        colwidths = self.table[0]

        for line in lines:
            if line == 'sep':
                pass
            else:
                cells = []
                for i, cell in enumerate(line):
                    cells.append(cell)
                rows.append(cells)

        def writerow(row):
            lines = zip_longest(*row)
            for line in lines:
                out += "<tr>"
                for i, cell in enumerate(line):
                    out += '<td>' + cell + "</td>"
                out += '</tr>'

        out = "<table>"
        for row in rows:
            writerow(row)
        out += "</table>"
        self.table = None
        self.end_state()

    def visit_acks(self, node):
        self.new_state("text")
        self.add_text(', '.join(n.astext() for n in node.children[0].children) +
                      '.')
        self.end_state()
        raise nodes.SkipNode

    def visit_image(self, node):
        if 'alt' in node.attributes:
            self.add_text(_('[image: %s]') % node['alt'])
        self.add_text(_('[image]'))
        raise nodes.SkipNode

    def visit_transition(self, node):
        indent = sum(self.stateindent)
        self.new_state("text")
        self.add_text('=' * (MAXWIDTH - indent))
        self.end_state()
        raise nodes.SkipNode

    def visit_bullet_list(self, node):
        self.list_counter.append(-1)

    def depart_bullet_list(self, node):
        self.list_counter.pop()

    def visit_enumerated_list(self, node):
        self.list_counter.append(node.get('start', 1) - 1)

    def depart_enumerated_list(self, node):
        self.list_counter.pop()

    def visit_definition_list(self, node):
        self.list_counter.append(-2)

    def depart_definition_list(self, node):
        self.list_counter.pop()

    def visit_list_item(self, node):
        if self.list_counter[-1] == -1:
            # bullet list
            self.new_state("text")
        elif self.list_counter[-1] == -2:
            # definition list
            pass
        else:
            # enumerated list
            self.list_counter[-1] += 1
            self.new_state("text")

    def depart_list_item(self, node):
        pass
        #if self.list_counter[-1] == -1:
        #    self.end_state(first='* ')
        #elif self.list_counter[-1] == -2:
        #    pass
        #else:
        #    self.end_state(first='%s. ' % self.list_counter[-1])

    def visit_definition_list_item(self, node):
        self._classifier_count_in_li = len(node.traverse(nodes.classifier))

    def depart_definition_list_item(self, node):
        pass

    def visit_term(self, node):
        self.new_state("text")

    def depart_term(self, node):
        if not self._classifier_count_in_li:
            self.end_state()

    def visit_termsep(self, node):
        self.add_text(', ')
        raise nodes.SkipNode

    def visit_classifier(self, node):
        self.add_text(' : ')

    def depart_classifier(self, node):
        self._classifier_count_in_li -= 1
        if not self._classifier_count_in_li:
            self.end_state()

    def visit_definition(self, node):
        self.new_state("text")

    def depart_definition(self, node):
        self.end_state()

    def visit_field_list(self, node):
        pass

    def depart_field_list(self, node):
        pass

    def visit_field(self, node):
        pass

    def depart_field(self, node):
        pass

    def visit_field_name(self, node):
        self.new_state("text")

    def depart_field_name(self, node):
        self.add_text(':')
        self.end_state()

    def visit_field_body(self, node):
        self.new_state("text")

    def depart_field_body(self, node):
        self.end_state()

    def visit_centered(self, node):
        pass

    def depart_centered(self, node):
        pass

    def visit_hlist(self, node):
        pass

    def depart_hlist(self, node):
        pass

    def visit_hlistcol(self, node):
        pass

    def depart_hlistcol(self, node):
        pass

    def visit_admonition(self, node):
        self.new_state("text")

    def depart_admonition(self, node):
        self.end_state()

    def _visit_admonition(self, node):
        self.new_state("text")

        if isinstance(node.children[0], nodes.Sequential):
            self.add_text(self.nl)

    def _make_depart_admonition(name):
        def depart_admonition(self, node):
            self.add_text(admonitionlabels[name] + ': ')
        return depart_admonition

    visit_attention = _visit_admonition
    depart_attention = _make_depart_admonition('attention')
    visit_caution = _visit_admonition
    depart_caution = _make_depart_admonition('caution')
    visit_danger = _visit_admonition
    depart_danger = _make_depart_admonition('danger')
    visit_error = _visit_admonition
    depart_error = _make_depart_admonition('error')
    visit_hint = _visit_admonition
    depart_hint = _make_depart_admonition('hint')
    visit_important = _visit_admonition
    depart_important = _make_depart_admonition('important')
    visit_note = _visit_admonition
    depart_note = _make_depart_admonition('note')
    visit_tip = _visit_admonition
    depart_tip = _make_depart_admonition('tip')
    visit_warning = _visit_admonition
    depart_warning = _make_depart_admonition('warning')
    visit_seealso = _visit_admonition
    depart_seealso = _make_depart_admonition('seealso')

    def visit_versionmodified(self, node):
        self.new_state("text")

    def depart_versionmodified(self, node):
        self.end_state()

    def visit_literal_block(self, node):
        self.new_state("code")

    def depart_literal_block(self, node):
        self.end_state()

    def visit_doctest_block(self, node):
        self.new_state("code")

    def depart_doctest_block(self, node):
        self.end_state()

    def visit_line_block(self, node):
        self.new_state("code")
        self.lineblocklevel += 1

    def depart_line_block(self, node):
        self.lineblocklevel -= 1
        self.end_state()
        if not self.lineblocklevel:
            self.add_text('\n')

    def visit_line(self, node):
        pass

    def depart_line(self, node):
        self.add_text('\n')

    def visit_block_quote(self, node):
        self.new_state("text")

    def depart_block_quote(self, node):
        self.end_state()

    def visit_compact_paragraph(self, node):
        pass

    def depart_compact_paragraph(self, node):
        pass

    def visit_paragraph(self, node):
        if not isinstance(node.parent, nodes.Admonition) or \
           isinstance(node.parent, addnodes.seealso):
            self.new_state("text")

    def depart_paragraph(self, node):
        if not isinstance(node.parent, nodes.Admonition) or \
           isinstance(node.parent, addnodes.seealso):
            self.end_state()

    def visit_target(self, node):
        raise nodes.SkipNode

    def visit_index(self, node):
        raise nodes.SkipNode

    def visit_toctree(self, node):
        raise nodes.SkipNode

    def visit_substitution_definition(self, node):
        raise nodes.SkipNode

    def visit_pending_xref(self, node):
        pass

    def depart_pending_xref(self, node):
        pass

    def visit_reference(self, node):
        pass

    def depart_reference(self, node):
        pass

    def visit_number_reference(self, node):
        text = nodes.Text(node.get('title', '#'))
        self.visit_Text(text)
        raise nodes.SkipNode

    def visit_download_reference(self, node):
        pass

    def depart_download_reference(self, node):
        pass

    def visit_emphasis(self, node):
        self.add_text('*')

    def depart_emphasis(self, node):
        self.add_text('*')

    def visit_literal_emphasis(self, node):
        self.add_text('*')

    def depart_literal_emphasis(self, node):
        self.add_text('*')

    def visit_strong(self, node):
        self.add_text('**')

    def depart_strong(self, node):
        self.add_text('**')

    def visit_literal_strong(self, node):
        self.add_text('**')

    def depart_literal_strong(self, node):
        self.add_text('**')

    def visit_abbreviation(self, node):
        self.add_text('')

    def depart_abbreviation(self, node):
        if node.hasattr('explanation'):
            self.add_text(' (%s)' % node['explanation'])

    def visit_title_reference(self, node):
        self.add_text('*')

    def depart_title_reference(self, node):
        self.add_text('*')

    def visit_literal(self, node):
        self.add_text('"')

    def depart_literal(self, node):
        self.add_text('"')

    def visit_subscript(self, node):
        self.add_text('_')

    def depart_subscript(self, node):
        pass

    def visit_superscript(self, node):
        self.add_text('^')

    def depart_superscript(self, node):
        pass

    def visit_footnote_reference(self, node):
        self.add_text('[%s]' % node.astext())
        raise nodes.SkipNode

    def visit_citation_reference(self, node):
        self.add_text('[%s]' % node.astext())
        raise nodes.SkipNode

    def visit_Text(self, node):
        self.add_text(node.astext())

    def depart_Text(self, node):
        pass

    def visit_generated(self, node):
        pass

    def depart_generated(self, node):
        pass

    def visit_inline(self, node):
        if 'xref' in node['classes'] or 'term' in node['classes']:
            self.add_text('*')

    def depart_inline(self, node):
        if 'xref' in node['classes'] or 'term' in node['classes']:
            self.add_text('*')

    def visit_container(self, node):
        pass

    def depart_container(self, node):
        pass

    def visit_problematic(self, node):
        self.add_text('>>')

    def depart_problematic(self, node):
        self.add_text('<<')

    def visit_system_message(self, node):
        self.new_state("text")
        self.add_text('<SYSTEM MESSAGE: %s>' % node.astext())
        self.end_state()
        raise nodes.SkipNode

    def visit_comment(self, node):
        raise nodes.SkipNode

    def visit_meta(self, node):
        # only valid for HTML
        raise nodes.SkipNode

    def visit_raw(self, node):
        if 'text' in node.get('format', '').split():
            self.body.append(node.astext())
        raise nodes.SkipNode

    def visit_math(self, node):
        self.builder.warn('using "math" markup without a Sphinx math extension '
                          'active, please use one of the math extensions '
                          'described at http://sphinx-doc.org/ext/math.html',
                          (self.builder.current_docname, node.line))
        raise nodes.SkipNode

    visit_math_block = visit_math
    visit_displaymath = visit_math

    def visit_autosummary_table(self, node):
        pprint("unprocessed node autosummary table:")
        pprint(node)
        pprint(type(node))
        #pprint(dir(node))

    def depart_autosummary_table(self, node):
        pass

    def visit_autosummary_toc(self, node):
        pass

    def depart_autosummary_toc(self, node):
        pass


    def unknown_visit(self, node):
        raise NotImplementedError('Unknown node: ' + node.__class__.__name__)

