#!/usr/bin/env python3
import sys
from argparse import ArgumentParser
from json import loads
from pathlib import Path
from subprocess import PIPE, Popen
from sys import stderr
from typing import TypedDict

parser = ArgumentParser(prog="restic-search")
parser.add_argument("search", type=str, nargs="+")
parser.add_argument(
    "-n", "--last", type=int, help="Only search through the n latest snapshots"
)


class Snapshot(TypedDict):
    id: str


def restic_snapshots() -> list[Snapshot]:
    ps = Popen(["restic", "snapshots", "--json"], stdout=PIPE)
    stdout, _ = ps.communicate()
    if ps.returncode != 0:
        raise RuntimeError(
            f"Error while running restic snapshots: returncode {ps.returncode}"
        )
    j = loads(stdout.decode("utf-8"))
    return j


def restic_ls(snapshot_id: str) -> list[str]:
    ps = Popen(["restic", "ls", snapshot_id], stdout=PIPE)
    stdout, _ = ps.communicate()
    if ps.returncode != 0:
        raise RuntimeError(f"Error while running restic ls: returncode {ps.returncode}")
    return stdout.decode("utf-8").split("\n")


if __name__ == "__main__":
    args = parser.parse_args()
    stderr.write("Listing snapshots...\n")
    snapshots = restic_snapshots()
    if args.last:
        snapshots = snapshots[-args.last :]
    results = []
    for snapshot in snapshots:
        stderr.write(f"Listing files in snapshot {snapshot['id']}...\n")
        files = restic_ls(snapshot["id"])
        for file in files:
            path = Path(file)
            for search in args.search:
                if path.match(search):
                    results.append((search, file, snapshot["id"]))

    for result in results:
        print(f"{result[0]}\t{result[1]}\t{result[2]}")