#!/usr/bin/env python
# Script to ramdiskify things

import argparse
import os
import shutil
import subprocess
import sys
import time
from pathlib import Path

def eprint(*args, **kwargs):
    print(*args, file=sys.stderr, **kwargs)

class RamDiskItem():
    def __init__(self, target, persist):
        self.target = Path(target)
        self.persist = Path(persist)
        self.ramdisk = Path("/dev/shm/ramdisk")
        self.full_rd_path = self.ramdisk.joinpath(self.target.name)
        self.full_persist_path = self.persist.joinpath(self.target.name)
        self.rsync_destination = self.full_persist_path if self.target.is_file() else self.persist


    def init(self):
        self.persist.mkdir(parents=True, exist_ok=True)
        self.ramdisk.mkdir(parents=True, exist_ok=True)

        if self.target.is_file():
            shutil.copy2(self.target, self.ramdisk)
        else:
            shutil.copytree(self.target, self.full_rd_path, dirs_exist_ok=True)
        shutil.move(self.target, self.persist)

        os.symlink(self.full_rd_path, self.target)


    def sync(self):
        # TODO: actually look at output from this...
        subprocess.run(["rsync", "-aP", "--delete", self.full_rd_path, self.rsync_destination])

    def loop_sync(self, delay):
        while True:
            try:
                self.sync()
                time.sleep(delay)
            except KeyboardInterrupt:
                eprint("Interrupt received, exiting...")
                return

    def cleanup(self):
        if not self.target.is_symlink():
            eprint(f"Target: {self.target} isn't a symlink!")
            return
        self.sync()
        self.target.unlink()
        shutil.move(self.full_persist_path, self.target.parent)
        if self.full_rd_path.is_file():
            self.full_rd_path.unlink()
        else:
            shutil.rmtree(self.full_rd_path)


def main():
    parser = argparse.ArgumentParser()
    # Arguments to specify the action to take
    action = parser.add_mutually_exclusive_group(required=True)
    action.add_argument("-i", "--init", action="store_true",
        help="Initialize the target file/directory into ramdisk")
    action.add_argument("-s", "--sync", action="store_true",
        help="Perform a single sync of the target from ramdisk to persistent storage")
    action.add_argument("-l", "--loop_sync", type=int,
        help="Starts a repeating sync that loops every LOOP_SYNC seconds")
    action.add_argument("-c", "--cleanup", action="store_true",
        help="clean up an existing ramdiskified target")
    # Arguments to specify what to act on
    parser.add_argument("target", help="The target file/directory to ramdiskify")
    parser.add_argument("persist", nargs="?", default=".persist",
        help="The location to store the persistent copy while ramdiskified, defaults to '.persist'")

    args = parser.parse_args()

    if not Path(args.target).exists():
        eprint(f"Target: {args.target} doesn't exist!")
        return

    rd_item = RamDiskItem(args.target, args.persist)

    if args.init:
        rd_item.init()
    elif args.sync:
        rd_item.sync()
    elif args.loop_sync:
        rd_item.loop_sync(args.loop_sync)
    elif args.cleanup:
        rd_item.cleanup()
    else:
        pass

if __name__ == "__main__":
    main()