from pygments.formatters.html import HtmlFormatter


class CustomHtmlFormatter(HtmlFormatter):
    """Custom HTML formatter. Adds an option to wrap lines into HTML <p> tags."""
    def __init__(self, **options):
        super().__init__(**options)
        self.lineparagraphs = options.get('lineparagraphs', '')

    def _wrap_lineparagraphs(self, inner):
        """Wrap lines into <p> tags

        :param inner: iterator of tuples of format (code, text)
        :return: iterator of tuples containing updated wrapped lines
        """
        s = self.lineparagraphs
        i = self.linenostart - 1
        for t, line in inner:
            if t:
                i += 1
                yield 1, '<p id="%s-%d">%s</p>' % (s, i, line)
            else:
                yield 0, line

    def format_unencoded(self, tokensource, outfile):
        """
        The formatting process uses several nested generators; which of
        them are used is determined by the user's options.

        Each generator should take at least one argument, ``inner``,
        and wrap the pieces of text generated by this.

        Always yield 2-tuples: (code, text). If "code" is 1, the text
        is part of the original tokensource being highlighted, if it's
        0, the text is some piece of wrapping. This makes it possible to
        use several different wrappers that process the original source
        linewise, e.g. line number generators.
        """
        source = self._format_lines(tokensource)

        # As a special case, we wrap line numbers before line highlighting
        # so the line numbers get wrapped in the highlighting tag.
        if not self.nowrap and self.linenos == 2:
            source = self._wrap_inlinelinenos(source)

        if self.hl_lines:
            source = self._highlight_lines(source)

        if not self.nowrap:
            if self.lineanchors:
                source = self._wrap_lineanchors(source)
            if self.linespans:
                source = self._wrap_linespans(source)
            # vvv customization of bepasty start:
            if self.lineparagraphs:
                source = self._wrap_lineparagraphs(source)
            # ^^^ customization of bepasty end.
            source = self.wrap(source)
            if self.linenos == 1:
                source = self._wrap_tablelinenos(source)
            source = self._wrap_div(source)
            if self.full:
                source = self._wrap_full(source, outfile)

        for t, piece in source:
            outfile.write(piece)
