"""Colorize read_plus.zsh output."""
import sys
import re
import termcolor

SIZE = sys.argv[1]
FILE = open(sys.argv[2], 'a') if len(sys.argv) >= 3 else None
COLORS = {"data": "red", "hole": "green", "d": "cyan", "h": "yellow"}


def write_log(text: str, color: str = None,
              attrs: list = [], end: str = "") -> None:
    """Print a single formatted string to the screen and possibly a file."""
    if FILE:
        FILE.write(text + end)
    termcolor.cprint(text, color=color, attrs=attrs, end=end)


def print_allocation(prefix: str, name: str,
                     filetype: str, extra: str) -> None:
    """Format and print file allocation information."""
    write_log(prefix, color="white", attrs=["dark"])
    write_log(name, color=COLORS[filetype], attrs=["bold"])
    write_log(extra, color="white", attrs=["dark"])


def print_allocation_done(done: str, time: str, end: str) -> None:
    """Format and print the time to complete file allocation."""
    write_log(done, color="white", attrs=["dark"])
    write_log(time, color="blue", attrs=["bold"])
    write_log(end, color="white", attrs=["dark"])


def print_reading(prefix: str, name: str,
                  precache: str, cached: str, postcache: str) -> None:
    """Format and print which file we are currently reading."""
    match cached:
        case "uncached": ccolor = "yellow"
        case "cached": ccolor = "green"
        case "unknown": ccolor = "grey"
        case _: ccolor = "white"

    write_log(prefix, color="white", attrs=["dark"])
    write_log(name, color=COLORS[filetype], attrs=["bold"])
    write_log(precache, color="white", attrs=["dark"])
    write_log(cached, color=ccolor, attrs=["bold"])
    write_log(postcache, color="white", attrs=["dark"])
    write_log(" ...", color="white", attrs=["dark"])


def print_stats_1(walltime: str, speed: str) -> None:
    """Format and print the transfer time and rate information."""
    write_log(" ")
    write_log(walltime, color="blue", attrs=["bold"])
    write_log(", ", color="white", attrs=["dark"])
    write_log(speed, color="green", attrs=["bold"])
    write_log(", ", color="white", attrs=["dark"])


def print_stats_2(kerntime: str, kernunits: str, cpu: str) -> None:
    """Format and print time spent in the kernel and cpu usage."""
    write_log(kerntime, color="cyan", attrs=["bold"], end=" ")
    write_log(kernunits, color="cyan", attrs=["bold"])
    write_log(", ", color="white", attrs=["dark"])
    write_log(cpu, color="red", attrs=["bold"])


def print_iteration(prefix: str, cur: str, sep: str, total: str) -> None:
    """Format and print the current iteration header."""
    write_log(prefix, color="magenta", attrs=["bold", "underline"])
    write_log(cur, color="green", attrs=["bold", "underline"])
    write_log(sep, color="magenta", attrs=["bold", "underline"])
    write_log(total, color="blue", attrs=["bold", "underline"], end="\n")


text = ""
while c := sys.stdin.read(1):
    text += c
    line = text
    text = ""
    if match := re.match(r"(Allocating file: )"
                         rf"({SIZE}M-(data|hole|mixed-\d+(d|h)))( ...)", line):
        filetype = match.group(3) if match.group(4) is None else match.group(4)
        print_allocation(match.group(1), match.group(2),
                         filetype, match.group(5))
    elif match := re.match(r"( \[Done: )(\d+\.\d+s)(\]\n)", line):
        print_allocation_done(match.group(1), match.group(2), match.group(3))
    elif match := re.match(rf"(Reading: )({SIZE}M-(data|hole|mixed-\d+(d|h)))"
                           r"( \()((un)?(cached|known))( on [\w-]*\))\n",
                           line):
        filetype = match.group(3) if match.group(4) is None else match.group(4)
        print_reading(match.group(1), match.group(2), match.group(5),
                      match.group(6), match.group(9))
    elif match := re.match(r"\d+ bytes \([\w,\. ]*\) copied, (\d+\.\d+ s), "
                           r"(\d+\.?\d+ [MG]B/s)\n", line):
        print_stats_1(match.group(1), match.group(2))
    elif match := re.match(r"(\d+\.\d+)(s kern), (\d+% cpu)", line):
        print_stats_2(match.group(1), match.group(2), match.group(3))
    elif match := re.match(r"\d+\+\d+ records (in|out)\n", line):
        pass
    elif match := re.match(r"(Iteration )(\d+)( / )(\d+)\n", line):
        print_iteration(match.group(1), match.group(2),
                        match.group(3), match.group(4))
    elif c == "\n":
        write_log(line)
    else:
        text = line