#!/usr/bin/env python3
"""Render rsync's per-file output as a single updating status line.

Fed the size and path rsync reports with --out-format='%l %n', one per line.
The size is what makes an estimate possible: rsync's own rate is not exposed
per file, but bytes completed over time elapsed is the same arithmetic and
needs nothing rsync does not already print. Prints one
line that rewrites itself, showing how far through the transfer is and which
album is currently going across, rather than either scrolling fifty thousand
filenames past or -- as rsync does while it builds its file list -- saying
nothing at all for several minutes.

Falls back to periodic plain lines when stderr is not a terminal, so a log does
not fill up with carriage returns.
"""

import argparse
import collections
import os
import shutil
import sys
import time

# The estimate is taken over a trailing window rather than the whole run, so it
# follows a device that slows down instead of averaging the slowdown away.
RATE_WINDOW_SECONDS = 30.0

# Below this the window is too narrow to divide by: the first few files arrive
# in microseconds and produce a rate in the gigabytes per second, and an ETA of
# nothing at all. Better to show neither until the figure means something.
RATE_MINIMUM_SPAN_SECONDS = 2.0


def parse(line):
    """Return (bytes, path) for one line of rsync output.

    Tolerates a bare path, in case someone runs this against --out-format='%n'.
    """
    line = line.rstrip("\n")
    size, separator, path = line.partition(" ")
    if separator and size.isdigit():
        return int(size), path
    return 0, line


def human_bytes(count):
    size = float(count)
    for unit in ("B", "KiB", "MiB", "GiB", "TiB"):
        if size < 1024 or unit == "TiB":
            return f"{size:.1f} {unit}"
        size /= 1024


def human_duration(seconds):
    """Return a duration nobody has to do arithmetic on."""
    seconds = int(seconds)
    if seconds < 60:
        return f"{seconds}s"
    if seconds < 3600:
        return f"{seconds // 60}m{seconds % 60:02d}s"
    return f"{seconds // 3600}h{(seconds % 3600) // 60:02d}m"


class Rate:
    """Bytes per second over a trailing window."""

    def __init__(self, window=RATE_WINDOW_SECONDS):
        self.window = window
        self.samples = collections.deque()

    def add(self, when, total_bytes):
        self.samples.append((when, total_bytes))
        while len(self.samples) > 2 and when - self.samples[0][0] > self.window:
            self.samples.popleft()

    def per_second(self):
        if len(self.samples) < 2:
            return 0.0
        (first_time, first_bytes), (last_time, last_bytes) = (
            self.samples[0],
            self.samples[-1],
        )
        elapsed = last_time - first_time
        if elapsed < RATE_MINIMUM_SPAN_SECONDS:
            return 0.0
        return (last_bytes - first_bytes) / elapsed


def album_of(path):
    """Return "Artist / Album" for a mirror-relative path."""
    parts = [part for part in path.strip("/").split("/") if part]
    if len(parts) >= 3:
        return f"{parts[0]} / {parts[1]}"
    if len(parts) == 2:
        return parts[0]
    return ""


def fit(text, width):
    """Trim to the terminal, from the left: the album matters more than the artist."""
    if width <= 1 or len(text) <= width:
        return text
    return "…" + text[-(width - 1) :]


def render(done, total, copied, expected, rate, label, width):
    """Build the status line, giving whatever room is left to the album."""
    if total > 0:
        share = min(100, done * 100 // total)
        head = f"[{share:>3}%] {done:,}/{total:,}"
    else:
        head = f"[{done:,} files]"

    if expected > 0:
        head += f"  {human_bytes(copied)}/{human_bytes(expected)}"
    if rate > 0:
        head += f"  {human_bytes(rate)}/s"
        remaining = expected - copied
        if remaining > 0:
            head += f"  ETA {human_duration(remaining / rate)}"
    head += "  "
    return head + fit(label, max(0, width - len(head)))


def main(argv=None, stream=None, out=None):
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--total", type=int, default=0, help="files expected")
    parser.add_argument("--bytes", type=int, default=0, help="bytes expected")
    parser.add_argument("--interval", type=float, default=0.1, help="seconds between redraws")
    args = parser.parse_args(argv)

    stream = stream or sys.stdin
    out = out or sys.stderr
    interactive = out.isatty()
    width = shutil.get_terminal_size((100, 24)).columns - 1

    done = 0
    copied = 0
    rate = Rate()
    started = time.monotonic()
    rate.add(started, 0)
    last_drawn = 0.0
    label = ""

    for line in stream:
        size, path = parse(line)
        # rsync reports directories too, with a trailing slash and an inode
        # size. Counting them puts the percentage past a hundred and the byte
        # total well over what will actually be transferred.
        if not path or path.endswith("/"):
            continue
        done += 1
        copied += size
        label = album_of(path) or os.path.basename(path)

        now = time.monotonic()
        rate.add(now, copied)
        if interactive:
            if now - last_drawn >= args.interval:
                out.write(
                    "\r\033[2K"
                    + render(done, args.total, copied, args.bytes, rate.per_second(),
                             label, width)
                )
                out.flush()
                last_drawn = now
        elif now - last_drawn >= 30:
            out.write(
                render(done, args.total, copied, args.bytes, rate.per_second(), label, width)
                + "\n"
            )
            out.flush()
            last_drawn = now

    elapsed = max(1e-9, time.monotonic() - started)
    if interactive:
        out.write("\r\033[2K")
    summary = render(done, args.total, copied, args.bytes, 0, label, width).rstrip()
    out.write(f"{summary}\n")
    if copied:
        out.write(
            f"copied {human_bytes(copied)} in {human_duration(elapsed)}"
            f" at {human_bytes(copied / elapsed)}/s\n"
        )
    out.flush()
    return 0


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