#!/usr/bin/env python3
# Copyright (c) 2022 Robin Jarry

"""
Colorize plaintext email. Write to stdout.
"""

import re
import sys


TERM_RESET = "\x1b[0m"
TERM_BOLD = "\x1b[1m"
TERM_BG_GRAY = "\x1b[48;5;244m"
TERM_FG_GRAY = "\x1b[38;5;244m"
TERM_FG_RED = "\x1b[38;5;1m"
TERM_FG_GREEN = "\x1b[38;5;2m"
TERM_FG_CYAN = "\x1b[38;5;6m"
TERM_FG_YELLOW = "\x1b[38;5;229m"
TERM_FG_WHITE = "\x1b[38;5;15m"
TERM_FG_BLUE = "\x1b[38;5;75m"
TERM_FG_PURPLE = "\x1b[38;5;141m"
TERM_FG_ORANGE = "\x1b[38;5;208m"
TERM_FG_PINK = "\x1b[38;5;171m"
QUOTE_COLORS = {
    1: TERM_FG_BLUE,
    2: TERM_FG_ORANGE,
    3: TERM_FG_PURPLE,
    4: TERM_FG_PINK,
}
URL_RE = re.compile(
    r"""
    (
        https?://[\w,;:!/#%^=@~\&\*\+\?\.\-]+
        |
        [\w\-\+\.~/]*\w@\w[\w\-\.]+\w
    )
    """,
    re.VERBOSE,
)
HEADER_RE = re.compile(r"^[A-Z][\w-]+:")
DIFF_META = ("diff --git ", "index ", "--- ", "+++ ")


def replace_match(color, context_color):
    def replace_func(match):
        return color + match.group(0) + context_color

    return replace_func


def main():
    in_patch = in_signature = in_headers = in_body = False

    for line in sys.stdin.buffer:
        line = line.decode("utf-8", errors="replace").rstrip("\r\n")
        if in_patch:
            if line in ("--", "-- "):
                in_signature = True
                in_body = in_patch = in_headers = False
                line = TERM_FG_PURPLE + line + TERM_RESET
            elif line.startswith("@@ "):
                line = TERM_FG_CYAN + line + TERM_RESET
            elif any(line.startswith(m) for m in DIFF_META):
                line = TERM_BOLD + TERM_FG_WHITE + line + TERM_RESET
            elif line.startswith("+"):
                line = TERM_FG_GREEN + line + TERM_RESET
            elif line.startswith("-"):
                line = TERM_FG_RED + line + TERM_RESET

        elif in_signature:
            line = (
                TERM_FG_PURPLE
                + URL_RE.sub(replace_match(TERM_FG_YELLOW, TERM_FG_PURPLE), line)
                + TERM_RESET
            )

        elif in_headers:
            if line == "":
                in_body = True
                in_headers = False
            else:
                line = HEADER_RE.sub(replace_match(TERM_FG_PURPLE, TERM_RESET), line)
                line = URL_RE.sub(replace_match(TERM_FG_YELLOW, TERM_RESET), line)

        elif in_body:
            if line.startswith(">"):
                level = 0
                quotes = ""
                while line.startswith(">"):
                    quotes += ">"
                    line = line[1:]
                    level += 1
                    while line.startswith(" "):
                        quotes += " "
                        line = line[1:]
                quote_color = QUOTE_COLORS.get(quotes.count(">"), TERM_FG_GRAY)

                if line.startswith("+"):
                    line = quote_color + quotes + TERM_FG_GREEN + line + TERM_RESET
                elif line.startswith("-"):
                    line = quote_color + quotes + TERM_FG_RED + line + TERM_RESET
                else:
                    line = (
                        quote_color
                        + quotes
                        + URL_RE.sub(
                            replace_match(TERM_FG_YELLOW, quote_color), line
                        )
                        + TERM_RESET
                    )
            elif line.startswith("diff --git "):
                in_patch = True
                in_body = in_headers = False
                line = TERM_BOLD + TERM_FG_WHITE + line + TERM_RESET
            elif line in ("--", "-- "):
                in_signature = True
                in_body = in_patch = in_headers = False
                line = TERM_FG_PURPLE + line + TERM_RESET
            else:
                line = URL_RE.sub(replace_match(TERM_FG_YELLOW, TERM_RESET), line)

        elif line.startswith("diff --git "):
            in_patch = True
            in_body = in_headers = False
            line = TERM_BOLD + TERM_FG_WHITE + line + TERM_RESET

        elif line in ("--", "-- "):
            in_signature = True
            in_body = in_patch = in_headers = False
            line = TERM_FG_PURPLE + line + TERM_RESET

        elif HEADER_RE.search(line):
            in_headers = True
            line = HEADER_RE.sub(replace_match(TERM_FG_PURPLE, TERM_RESET), line)
            line = URL_RE.sub(replace_match(TERM_FG_YELLOW, TERM_RESET), line)

        else:
            in_body = True
            if line.startswith(">"):
                level = 0
                quotes = ""
                while line.startswith(">"):
                    quotes += ">"
                    line = line[1:]
                    level += 1
                    while line.startswith(" "):
                        quotes += " "
                        line = line[1:]
                quote_color = QUOTE_COLORS.get(quotes.count(">"), TERM_FG_GRAY)

                if line.startswith("+"):
                    line = quote_color + quotes + TERM_FG_GREEN + line + TERM_RESET
                elif line.startswith("-"):
                    line = quote_color + quotes + TERM_FG_RED + line + TERM_RESET
                else:
                    line = (
                        quote_color
                        + quotes
                        + URL_RE.sub(
                            replace_match(TERM_FG_YELLOW, quote_color), line
                        )
                        + TERM_RESET
                    )

        sys.stdout.write(line + "\r\n")
        sys.stdout.flush()


if __name__ == "__main__":
    sys.exit(main())