#!/usr/bin/env python3

import sys, os, re
from argparse import ArgumentParser, RawDescriptionHelpFormatter

sys.path.append(os.path.dirname(os.path.realpath(__file__)))
import ansicodes as ansi

__author__ = "Ville Rantanen <ville.q.rantanen@gmail.com>"
__version__ = "0.3"

""" Rules modified from mistune project """


def setup_options():
    bc = ansi.code()
    """ Create command line options """
    usage = (
        """
Markdown syntax color in ansi codes.
Special syntaxes:
- Colors: insert string e.g. ${C}.
- Any ANSI control code:  ${3A}, ${1;34;42m}, see the table..

"""
        + ansi.demo()
    )

    parser = ArgumentParser(
        formatter_class=RawDescriptionHelpFormatter,
        description=usage,
        epilog=__author__,
    )

    parser.add_argument("-v", "--version", action="version", version=__version__)
    parser.add_argument(
        "-D", action="store_true", dest="debug", default=False, help="Debug mode"
    )
    parser.add_argument(
        "--dc",
        action="store_true",
        dest="dark_colors",
        default=False,
        help="Use dark colorscheme, better for white background terminals.",
    )
    parser.add_argument(
        "--no-color",
        "-n",
        action="store_false",
        dest="color",
        default=True,
        help="Disable color.",
    )
    parser.add_argument(
        "--print-template",
        action="store_true",
        dest="print_template",
        default=False,
        help="Print customizable color template.",
    )
    parser.add_argument(
        "--template",
        action="store",
        type=str,
        dest="template",
        default=None,
        help="Use a custom template file for colorization.",
    )
    parser.add_argument(
        "-z",
        action="store_true",
        dest="zero",
        default=False,
        help="Reset coloring at the end of each line.",
    )
    parser.add_argument(
        "-Z",
        action="store_false",
        dest="zero_final",
        default=True,
        help="Disable reset of colors at the end of file.",
    )
    parser.add_argument(
        "filename", type=str, nargs="?", help="File to show, - for stdin"
    )
    opts = parser.parse_args()
    return opts


def parse(data):
    data = [[None, row] for row in data]
    block = "text"
    new_block = "text"
    multiline_block = False
    # Parse styles
    for i, line in enumerate(data):
        row = line[1]
        if line[0] is not None:
            # Previous lines have set the style already
            continue
        for match in blocks:
            if block_match[match]["re"].match(row):
                new_block = match
                if match.startswith("multiline"):
                    if multiline_block:
                        multiline_block = False
                    else:
                        multiline_block = match
                break
        if multiline_block:
            new_block = multiline_block
        # Lists must end with empty line
        if new_block not in ("empty", "list_bullet") and block.startswith("list_"):
            new_block = "list_loose"

        if "mod" in block_match[match]:
            # Style sets block in previous or next lines
            data[i + block_match[match]["mod"]["pos"]][0] = block_match[match]["mod"][
                "name"
            ]
        data[i][0] = new_block
        if new_block != block:
            block = new_block
    return data


def colorize(data, remove_colors=False, dark_colors=False, debug=False):
    # Start inserting colors, and printing
    bc = ansi.code()
    colorized = []
    cs = "dc" if dark_colors else "bc"
    csc = cs + "c"
    for i, line in enumerate(data):
        row = line[1]
        block = line[0]
        multiline_block = block.startswith("multiline")
        if multiline_block:
            row = block_match[block][csc] + row
        if block_match[block][cs]:
            row = block_match[block]["re"].sub(block_match[block][cs], row)
        # No coloring inside block_code, nor multiline
        if not (multiline_block or block == "block_code"):
            for match in inlines:
                if inline_match[match]["re"].search(row):
                    row = inline_match[match]["re"].sub(
                        inline_match[match][cs] + block_match[block][csc], row
                    )
        if remove_colors:
            colored = bc.nocolor_string(row)
        else:
            colored = bc.color_string(row)
        if debug:
            multistr = "*" if multiline_block else " "
            colored = "{:<18}{:}:".format(data[i][0], multistr) + colored
        colorized.append(colored)
    return colorized


def md_re_compile(d):
    """ Returns a re.compiled dict """
    n = {}
    for t in d:
        n[t] = {}
        for i in d[t]:
            n[t][i] = d[t][i]
        try:
            if n[t]["re"]:
                n[t]["re"] = re.compile(n[t]["re"])
        except err:
            print("Error compiling: %s" % (n[t]["re"]))
            sys.exit(1)
    return n


def read_data2(fp):
    data = []
    # Read data
    for row in f:
        if not row:
            continue
        row = row.decode("utf-8").rstrip("\n\r ")
        data.append(row)
    return data


def read_data3(fp):
    data = []
    # Read data
    for row in f:
        if not row:
            continue
        row = row.rstrip("\n\r ")
        data.append(row)
    return data


def write_colored2(colored, opts):
    for c in colored:
        sys.stdout.write(c.encode("utf-8"))
        if opts.zero:
            sys.stdout.write(bc.Z)
        sys.stdout.write("\n")
    if opts.zero_final:
        sys.stdout.write(bc.Z.encode("utf-8"))


def write_colored3(colored, opts):
    for c in colored:
        sys.stdout.write(c)
        if opts.zero:
            sys.stdout.write(bc.Z)
        sys.stdout.write("\n")
    if opts.zero_final:
        sys.stdout.write(bc.Z)


# re: regular expression, bc: bright colors, bcc: continue with this color after inline
# dc: dark colors, dcc: continued color after inline
block_match_str = {
    "block_code": {
        "re": "^( {4}[^*])(.*)$",
        "bc": "${Z}${c}\\1\\2",
        "bcc": "${Z}${c}",
        "dc": "${Z}${m}\\1\\2",
        "dcc": "${Z}${m}",
    },  #    code
    "multiline_code": {
        "re": "^ *(`{3,}|~{3,}) *(\S*)",
        "bc": "${Z}${c}\\1\\2",
        "bcc": "${Z}${c}",
        "dc": "${Z}${m}\\1\\2",
        "dcc": "${Z}${m}",
    },  # ```lang
    "block_quote": {
        "re": "^(>[ >]* )",
        "bc": "${K}\\1${Z}",
        "bcc": "${Z}",
        "dc": "${Y}\\1${Z}",
        "dcc": "${Z}",
    },  # > > quote
    "hrule": {
        "re": "^ {0,3}[-*_]([-*_]){2,}$",
        "bc": "False",
        "bcc": "${Z}",
        "dc": "False",
        "dcc": "${Z}",
    },  # ----
    "heading": {
        "re": "^ *(#{1,6}) *([^\n]+?) *#* *(?:\n+|$)",
        "bc": "${W}\\1 ${U}\\2${Z}",
        "bcc": "${W}${U}",
        "dc": "${B}\\1 ${U}\\2${Z}",
        "dcc": "${B}${U}",
    },  # # heading
    "lheading": {
        "re": "^(=+|-+)$",
        "bc": "${W}\\1",
        "bcc": "${W}",
        "dc": "${B}\\1",
        "dcc": "${B}",
        "mod": {"pos": -1, "name": "lheading.mod"},
    },  # ======= under header
    "lheading.mod": {
        "re": "^([^\n]+)",
        "bc": "${W}\\1",
        "bcc": "${W}",
        "dc": "${B}\\1",
        "dcc": "${B}",
    },  # over the ======= under header
    "list_bullet": {
        "re": "^( *)([*+-]|[\d\.]+)( +)",
        "bc": "\\1${y}\\2${Z}\\3",
        "bcc": "${Z}",
        "dc": "\\1${r}\\2${Z}\\3",
        "dcc": "${Z}",
    },  # * or 1.
    "list_loose": {
        "re": "None",
        "bc": "False",
        "bcc": "${Z}",
        "dc": "False",
        "dcc": "${Z}",
    },
    "text": {
        "re": "^([^\n]+)",
        "bc": "${Z}\\1",
        "bcc": "${Z}",
        "dc": "${Z}\\1",
        "dcc": "${Z}",
    },
    "empty": {
        "re": "(^$)",
        "bc": "${Z}\\1",
        "bcc": "${Z}",
        "dc": "${Z}\\1",
        "dcc": "${Z}",
    },
    "preformatted": {
        "re": "a^",  # Never matches anything
        "bc": "",
        "bcc": "",
        "dc": "",
        "dcc": "",
    },
}
block_match = md_re_compile(block_match_str)
blocks = [
    "block_quote",
    "multiline_code",
    "hrule",
    "heading",
    "lheading",
    "list_bullet",
    "block_code",
    "text",
    "empty",
]

inline_match_str = {
    "bold1": {
        "re": r"(^| |})(_[^_]+_)",
        "bc": "\\1${W}\\2",
        "dc": "\\1${W}\\2",
    },  # _bold_
    "bold2": {
        "re": r"(^| |})(\*{1,2}[^\*]+\*{1,2})",
        "bc": "\\1${W}\\2",
        "dc": "\\1${W}\\2",
    },  # **bold**
    "code": {"re": r"([`]+[^`]+[`]+)", "bc": "${c}\\1", "dc": "${m}\\1"},  # `code`
    "code_special": {
        "re": r"([`]+[^`]+[`]+)([!>])",
        "bc": "${c}\\1${g}\\2",
        "dc": "${m}\\1${r}\\2",
    },  # `code`! or `code`> for markslider
    "link": {
        "re": r"(\[)([^\]]+)(\])\(([^\)]+)\)",
        "bc": "${B}\\1${Z}\\2${B}\\3(${U}\\4${u})",
        "dc": "${b}\\1${Z}\\2${b}\\3(${U}\\4${u})",
    },  # [text](link)
    "image": {
        "re": r"(!\[[^\]]+\]\([^\)]+\))",
        "bc": "${r}\\1",
        "dc": "${g}\\1",
    },  # ![text](image)
    "underline": {
        "re": r"(^|\W)(__)([^_]+)(__)",
        "bc": "\\1\\2${U}\\3${Z}\\4",
        "dc": "\\1\\2${U}\\3${Z}\\4",
    },  # __underline__
    "strikethrough": {
        "re": r"(~~)([^~]+)(~~)",
        "bc": "\\1${st}\\2${so}\\3",
        "dc": "\\1${st}\\2${so}\\3",
    },  # ~~strike~~
    "checkbox": {"re": r"(\[[x ]\])", "bc": "${y}\\1", "dc": "${r}\\1"},  # [x] [ ]
}
inline_match = md_re_compile(inline_match_str)
inlines = [
    "bold1",
    "bold2",
    "code_special",
    "code",
    "image",
    "link",
    "underline",
    "strikethrough",
    "checkbox",
]

if __name__ == "__main__":
    opts = setup_options()
    if opts.print_template:
        import json

        print(
            json.dumps(
                {
                    "about": 're: regular expression, bc: bright colors, bcc: continue with this color after inline, dc: dark colors, dcc: continued color after inline. "blocks" and "inlines" list keys of matchers. ',
                    "blocks": blocks,
                    "block_match": block_match_str,
                    "inline_match": inline_match_str,
                    "inlines": inlines,
                },
                indent=2,
                sort_keys=True,
            )
        )
        sys.exit(0)
    if opts.template:
        import json

        template = json.load(open(opts.template, "r"))
        if "inlines" in template:
            inlines = template["inlines"]
        if "blocks" in template:
            blocks = template["blocks"]
        if "block_match" in template:
            block_match_str = template["block_match"]
            block_match = md_re_compile(block_match_str)
        if "inline_match" in template:
            inline_match_str = template["inline_match"]
            inline_match = md_re_compile(inline_match_str)

    bc = ansi.code()
    if opts.filename == "-" or opts.filename == None:
        f = sys.stdin
    else:
        f = open(opts.filename, "r")
    if sys.version_info > (3, 0):
        data = read_data3(f)
    else:
        data = read_data2(f)

    data = parse(data)
    colored = colorize(data, not opts.color, opts.dark_colors, opts.debug)
    if sys.version_info > (3, 0):
        write_colored3(colored, opts)
    else:
        write_colored2(colored, opts)