#!/usr/bin/env python3

import argparse
import sys
import re

MARKERSTART = "[TOCSTART]: #"
MARKEREND = "[TOCEND]: #"


def get_opts():
    parser = argparse.ArgumentParser(
        description="TOC for markdown.",
        formatter_class=argparse.RawTextHelpFormatter,
    )
    parser.add_argument("-n", help="Number TOC", action="store_true")
    parser.add_argument("-l", help="TOC as links", action="store_true")
    parser.add_argument("-f", help="Flat list, i.e. not bulleted list. Good if your headers contain section numbers.", action="store_true")
    parser.add_argument(
        "-a", help="Print whole Markdown, with TOC included.", action="store_true"
    )
    parser.add_argument(
        "markdown",
        help="Filename to read. - for stdin",
        action="store",
        default="-",
        nargs="?",
    )
    args = parser.parse_args()
    return args


def main():
    opts = get_opts()
    if opts.markdown == "-":
        fp = sys.stdin
    else:
        fp = open(opts.markdown, "rt")
    markdown = fp.read().splitlines()


    bullet = "" if opts.f else "- "
    postspace = "  " if opts.f else ""
    counters = [0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
    lastlev = 0
    re_anchor = re.compile("[^0-9a-z-]")

    entries = []
    marker_found = False
    marker_end_found = False
    to_remove = []
    in_code = False

    for i, row in enumerate(markdown):
        in_table = row.startswith("|")
        if row.startswith("```"):
            in_code = not in_code

        if any((in_code, in_table)):
            continue

        if row.startswith("#"):
            headers, title = row.strip().split(" ", 1)
            if opts.l:
                anchor = re_anchor.sub("", title.lower().replace(" ", "-"))
                title = "[{}](#{})".format(title, anchor)
            pad = headers.count("#") - 1
            if pad < lastlev:
                counters = [x if i <= pad else 0 for i, x in enumerate(counters)]
            counters[pad] += 1
            lastlev = pad
            if opts.n:
                bullet = "{:d}. ".format(counters[pad])
            if opts.f:
                pad = 0
            entries.append("{}{}{}{}".format(pad * "    ", bullet, title, postspace))
        if row == MARKERSTART:
            marker_found = True
        if marker_found and not marker_end_found:
            to_remove.append(i)
            if row == MARKEREND:
                marker_end_found = True

    # Just print the TOC
    if not opts.a:
        print("\n".join(entries))
        return
    # Print whole document
    if marker_found and not marker_end_found:
        raise ValueError(f"End marker '{MARKEREND}' not found in document")

    markdown = [row for i, row in enumerate(markdown) if i not in to_remove]
    entries = [MARKERSTART, *entries, "", MARKEREND]
    toc_position = 0
    if len(to_remove) > 0:
        toc_position = to_remove[0]

    for i, row in enumerate(markdown):
        if i == toc_position:
            print("\n".join(entries))
        print(row)


if __name__ == "__main__":
    main()