#!/usr/bin/env python3
"""Copy file timestamps from source to destination directory.

Mostly to fix erroneous file copies that did not carry over file access
and modification times. This can be due to missing 'cp -p', or wrong
destination owner during 'rsync -a' copies (especially to NFS destinations).

The two pass is absolutely intended, because it'll be so easy to accidentally
copy timestamps FROM the destination TO the source instead. Time wasted better
than time lost :P

Examples:
    $ python transplant_file_timestamps.py ./src
    $ python transplant_file_timestamps.py ./dest

Changelog:
    2025-09-02 Justin: Init
"""

import argparse
import os
import pathlib
import pickle
import sys

parser = argparse.ArgumentParser(description="Transplants file access and modify timestamps, to failed copies.", usage=__doc__)
parser.add_argument("-f", default="transplant_file_timestamps.py.pickle", help="auxiliary file to store mtime.")
parser.add_argument("--ignore-same-dir", action="store_true")
parser.add_argument("DIR", help="directory to apply on.")
args = parser.parse_args()

pathdir = pathlib.Path(args.DIR)

# Update timestamps
if pathlib.Path(args.f).exists():
    with open(args.f, "rb") as f:
        source, data = pickle.load(f)

    if pathdir.absolute() == source:
        if not args.ignore_same_dir:
            raise ValueError("Transplanting to same directory: ignored.")
        print(f"Ignoring same directory: '{pathdir.absolute()}'", file=sys.stderr)

    count = 0
    for path in pathdir.rglob("*"):
        relpath = str(path.relative_to(args.DIR))
        if relpath in data:
            times = data[relpath]
            os.utime(path, ns=times)
            count += 1

    print(f"Transplanted {count} files!", file=sys.stderr)

# Store timestamps only
else:
    data = {}
    count = 0
    for path in pathdir.rglob("*"):
        relpath = str(path.relative_to(args.DIR))
        stat = path.stat()
        times = (stat.st_atime_ns, stat.st_mtime_ns)
        data[relpath] = times
        count += 1

    with open(args.f, "wb") as f:
        pickle.dump([pathdir.absolute(), data], f)

    print(f"Copied {count} timestamps!", file=sys.stderr)