#!/usr/bin/env python3
import argparse
from collections import defaultdict
import io
import re
import os
import subprocess
import sys

def trim_remote(remote):
    remote = remote.removesuffix(" (untrusted)")
    # user@host:/path -> host (but not for removable media)
    if m := re.match(r"^\w+@(\w+):/home/.+/Attic/.+", remote):
        remote = m.group(1)
    # user@host:/run/media/foo -> foo
    if m := re.match(r"^\w+@\w+:/run/media/\w+/(.+?)/", remote):
        remote = m.group(1)
    # foo (bar) -> foo
    if m := re.match(r"^(\S+) \(.+\)", remote):
        remote = m.group(1)
    return remote

def fmt_remote(remote):
    if remote == "here":                return fg(remote, -3)
    elif remote == "ember":             return fg(remote, 208)
    elif remote == "frost":             return fg(remote, 109)
    elif remote == "midnight":          return fg(remote, 146)
    elif remote == "rain":              return fg(remote, 82)
    elif re.match(r"^archive", remote): return fg(remote, 251)
    elif re.match(r"^vol\d", remote):   return fg(remote, 198 + int(remote[3:])*2)
    else:                               return fg(remote, 15)

def fmt_absent(remote):
    return fg(remote, 240)

def parse_annex_list(lines):
    is_header = True
    all_remotes = []
    for line in lines:
        line = line.rstrip()
        if line == "(recording state in git...)":
            continue
        elif re.match(r"^\(merging .*\.\.\.\)$", line):
            continue
        elif is_header:
            if m := re.match(r"^[|]+$", line):
                is_header = False
            elif m := re.match(r"^[|]*([^|]+)$", line):
                remote = m.group(1)
                remote = trim_remote(remote)
                all_remotes.append(remote)
            else:
                raise ValueError(f"bad header line {line!r}")
        else:
            if m := re.match(r"([Xx_]+) (.+)$", line):
                # 'X' indicates a normal remote, 'x' untrusted
                bits, path = m.groups()
                if len(bits) != len(all_remotes):
                    raise ValueError(f"bitmap length {len(bits)} of {bits!r} != "
                                     f"remote count {len(all_remotes)} of {all_remotes!r}")
                locations = {all_remotes[i] for i in range(len(bits)) if bits[i] in "Xx"}
                yield path, frozenset(locations)
            else:
                raise ValueError(f"bad list line {line!r}")

def whereis(find_args):
    # Pass --allrepos to see remotes that aren't in .git/config
    with subprocess.Popen(["git-annex", "list", "--allrepos", *find_args],
                          stdout=subprocess.PIPE) as proc:
        yield from parse_annex_list(io.TextIOWrapper(proc.stdout))

def fmt(string, ansifmt):
    return f"\033[{ansifmt}m{string}\033[m" if ansifmt else f"{string}"

def fg(string, color):
    if color < 0:
        string = fmt(string, "1")
        color = abs(color)
    return fmt(string, f"38;5;{color}")

def loc_fullstr(present, total, only_present):
    loc = sorted(present if only_present else total)
    loc = [fmt_remote(r) if r in present else fmt_absent(r) for r in loc]
    return fg("{", 8) + " ".join(loc) + fg("}", 8)

def loc_diffstr(old, new):
    loc = sorted(old | new)
    loc = [x for x in loc if (x in old) != (x in new)]
    loc = ["+"+fmt_remote(x) if x in new else "-"+fmt_absent(x) for x in loc]
    return fg("{", 8) + " ".join(loc) + fg("}", 8)

only_present = True
show_diffs_for_files = False
show_unabridged = False
show_non_diverging = False
always_hide = {"web"}

parser = argparse.ArgumentParser()
parser.add_argument("-a", "--all-remotes", action="store_true",
                    help="show all remotes, not only 'present' ones")
parser.add_argument("-f", "--all-files", action="store_true",
                    help="show all files, not only diverging ones")
parser.add_argument("-d", "--diff-locations", action="store_true",
                    help="show ±diff for diverging file locations")
args, find_args = parser.parse_known_args()

if args.all_remotes:
    only_present = False
if args.all_files:
    show_unabridged = True
if args.diff_locations:
    show_diffs_for_files = True

# Note: all_remotes only contains those remotes that have been used for at least one file.
all_remotes = set()
locs_by_dir = defaultdict(lambda: defaultdict(int))
locs_by_file = defaultdict(dict)

for path, locations in whereis(find_args):
    locations -= always_hide
    all_remotes |= locations
    dir_name = os.path.dirname(path)
    base_name = os.path.basename(path)
    locs_by_dir[dir_name][locations] += 1
    locs_by_file[dir_name][base_name] = locations

# Make it feel snappier by outputting whole chunks at once
sys.stdout.reconfigure(line_buffering=False)

for dir_name, loc_usage in sorted(locs_by_dir.items()):
    # All distinct location-sets for this directory, sorted by num. of occurences
    loc_sets = sorted(loc_usage.keys(), key=lambda i: loc_usage[i])
    loc_sets.reverse()

    # Print the most common set, approximately representing the whole directoy
    main_loc = loc_sets[0]
    main_loc_str = loc_fullstr(main_loc, all_remotes, only_present)
    print(f"{main_loc_str} {dir_name}/")

    if show_unabridged:
        # Show all sets in full
        full = True
    else:
        # Are top and bottom sets roughly equally common? If so, print all in full.
        top_usage = loc_usage[loc_sets[0]]
        btm_usage = loc_usage[loc_sets[-1]]
        full = (top_usage - btm_usage <= 2)

    # If any files diverge, or if -a -a was given, print those bitmaps
    if show_non_diverging or len(loc_sets) > 1:
        files = locs_by_file[dir_name]
        for file, file_loc in sorted(files.items()):
            if not full and file_loc == main_loc:
                continue
            if show_diffs_for_files:
                file_loc_str = loc_diffstr(main_loc, file_loc)
                if file_loc != main_loc:
                    file = fmt(file, "1")
            else:
                file_loc_str = loc_fullstr(file_loc, all_remotes, only_present)
            # The Perl version had a check to skip printing if the set was empty,
            # but I'm not sure if this can happen (due to the file==main check above).
            assert file_loc_str
            print(f"  {file_loc_str} {file}")
