#!/usr/bin/python3

import sys, os, bsddb3, struct, getopt, pwd, time, hashlib
bd = bsddb3.db
pj = os.path.join
deadlock = bd.DBLockDeadlockError
notfound = bd.DBNotFoundError

class txn(object):
    def __init__(self, env, flags=bd.DB_TXN_WRITE_NOSYNC):
        self.tx = env.txn_begin(None, flags)
        self.env = env
        self.done = False

    def commit(self):
        self.done = True
        self.tx.commit(0)

    def abort(self):
        self.done = True
        self.tx.abort()

    def __enter__(self):
        return self

    def __exit__(self, etype, exc, tb):
        if not self.done:
            self.abort()
        return False

class dbcursor(object):
    def __init__(self, db, tx):
        self.bk = db.cursor(txn=tx.tx)

    def __enter__(self):
        return self.bk

    def __exit__(self, *args):
        self.bk.close()
        return False

def txnfun(envfun):
    def fxf(fun):
        def wrapper(self, *args, tx=None, **kwargs):
            if tx is None:
                while True:
                    try:
                        with txn(envfun(self)) as ltx:
                            ret = fun(self, *args, tx=ltx, **kwargs)
                            ltx.commit()
                            return ret
                    except deadlock:
                        continue
            else:
                return fun(self, *args, tx=tx, **kwargs)
        return wrapper
    return fxf

class prefix(object):
    use = None

    def __init__(self, root, envdir):
        self.root = root
        self.envdir = envdir
        self._env = None
        self.dbs = {}

    @property
    def env(self):
        if self._env is None:
            if not os.path.isdir(self.envdir):
                sys.stderr.write("tpkg: creatings %s...\n" % (self.envdir))
                os.makedirs(self.envdir)
            env = bd.DBEnv()
            env.set_lk_detect(bd.DB_LOCK_RANDOM)
            fl = bd.DB_THREAD | bd.DB_INIT_MPOOL | bd.DB_INIT_LOCK | bd.DB_INIT_LOG | bd.DB_INIT_TXN | bd.DB_CREATE
            mode = 0o666
            env.open(self.envdir, fl, mode)
            self._env = env
        return self._env

    def maint(self):
        env = self._env
        if env is not None:
            env.txn_checkpoint(1024)
            env.log_archive(bd.DB_ARCH_REMOVE)

    @txnfun(lambda self: self.env)
    def db(self, name, *, dup=False, tx):
        if name not in self.dbs:
            db = bd.DB(self.env)
            if dup:
                db.set_flags(bd.DB_DUPSORT)
            fl = bd.DB_THREAD | bd.DB_CREATE
            mode = 0o666
            db.open(name, None, bd.DB_BTREE, fl, mode, txn=tx.tx)
            self.dbs[name] = db
        return self.dbs[name]

    def close(self):
        if self._env is not None:
            for db in self.dbs.values():
                db.close()
            self._env.close()
            self._env = None

    def __del__(self):
        self.close()

    @txnfun(lambda self: self.env)
    def unregfile(self, path, *, tx):
        epath = path.encode("utf-8")
        db = self.db("filedata")
        if db.has_key(epath, txn=tx.tx):
            db.delete(epath, txn=tx.tx)
        db = self.db("file-pkg")
        epkg = db.get(epath, None, txn=tx.tx)
        if epkg is not None:
            db.delete(epath, txn=tx.tx)
            with dbcursor(self.db("pkg-file", dup=True), tx) as cur:
                try:
                    cur.get_both(epkg, epath)
                except notfound:
                    pass
                else:
                    cur.delete()

    @txnfun(lambda self: self.env)
    def regfile(self, path, pkg, digest, *, tx):
        epath, epkg = path.encode("utf-8"), pkg.encode("utf-8")
        self.unregfile(path, tx=tx)
        filedata = b"digest\0" + digest.encode("utf-8") + b"\0"
        self.db("filedata").put(epath, filedata, txn=tx.tx)
        self.db("file-pkg").put(epath, epkg, txn=tx.tx)
        self.db("pkg-file", dup=True).put(epkg, epath, flags=bd.DB_NODUPDATA, txn=tx.tx)

    @txnfun(lambda self: self.env)
    def filedata(self, path, default=KeyError, *, tx):
        epath = path.encode("utf-8")
        data = self.db("filedata").get(epath, None, txn=tx.tx)
        if data is None:
            if default is KeyError:
                raise KeyError(path)
            else:
                return default
        data = data.split(b'\0')
        if data[-1] != b"" or len(data) % 2 != 1:
            raise Exception("invalid filedata")
        ret = {}
        for i in range(0, len(data) - 1, 2):
            ret[data[i].decode("utf-8")] = data[i + 1].decode("utf-8")
        return ret

    @txnfun(lambda self: self.env)
    def filepkg(self, path, default=KeyError, *, tx):
        epath = path.encode("utf-8")
        epkg = self.db("file-pkg").get(epath, None, txn=tx.tx)
        if epkg is None:
            if default is KeyError:
                raise KeyError(path)
            else:
                return default
        return epkg.decode("utf-8")

    @txnfun(lambda self: self.env)
    def pkgfiles(self, pkg, default=KeyError, *, tx):
        epkg = pkg.encode("utf-8")
        with dbcursor(self.db("pkg-file", dup=True), tx) as cur:
            try:
                edat = cur.set(epkg)
                if edat is None:
                    raise notfound()
                fpkg, epath = edat
                assert fpkg == epkg
            except notfound:
                if default is KeyError:
                    raise KeyError(pkg)
                else:
                    return default
            ret = []
            while fpkg == epkg:
                ret.append(epath.decode("utf-8"))
                edat = cur.next()
                if edat is None:
                    break
                fpkg, epath = edat
        return ret

    @classmethod
    def home(cls):
        home = pwd.getpwuid(os.getuid()).pw_dir
        return cls(pj(home, "sys"), pj(home, ".tpkg/db"))

    @classmethod
    def test(cls):
        home = pwd.getpwuid(os.getuid()).pw_dir
        return cls(pj(home, "tpkgtest"), pj(home, ".tpkg/testdb"))

    @classmethod
    def local(cls):
        return cls("/usr/local", "/usr/local/etc/tpkg/db")

class vfsfile(object):
    def __init__(self, path, fullpath):
        self.path = path
        self.fullpath = fullpath

    def open(self):
        return open(self.fullpath, "rb")

    def stat(self):
        return os.stat(self.fullpath)

class vfspkg(object):
    def __init__(self, root):
        self.root = root

    def __iter__(self):
        def scan(lp, fp):
            dpre = "" if (lp is "") else lp + "/"
            for dent in os.scandir(fp):
                dpath = dpre + dent.name
                if dent.is_dir():
                    yield from scan(dpath, dent.path)
                else:
                    yield vfsfile(dpath, dent.path)
        return scan("", self.root)

def copy(dst, src):
    dig = hashlib.sha256()
    while True:
        buf = src.read(65536)
        if buf == b"":
            return dig.hexdigest().lower()
        dst.write(buf)
        dig.update(buf)

def digest(fp):
    dig = hashlib.sha256()
    while True:
        buf = fp.read(65536)
        if buf == b"":
            return dig.hexdigest().lower()
        dig.update(buf)

def install(pfx, pkg, pkgname):
    for fl in pkg:
        if os.path.exists(pj(pfx.root, fl.path)):
            sys.stderr.write("tpkg: %s: already exists\n" % (fl.path))
            sys.exit(1)
    for fl in pkg:
        tp = pj(pfx.root, fl.path)
        tpdir = os.path.dirname(tp)
        if not os.path.isdir(tpdir):
            os.makedirs(tpdir)
        tmpp = tp + ".tpkg-new"
        sb = fl.stat()
        with open(tmpp, "wb") as ofp:
            os.fchmod(ofp.fileno(), sb.st_mode & 0o7777)
            with fl.open() as ifp:
                dig = copy(ofp, ifp)
        pfx.regfile(fl.path, pkgname, dig)
        os.rename(tmpp, tp)
        os.utime(tp, ns=(time.time(), sb.st_mtime))

def uninstall(pfx, pkg):
    for fn in pfx.pkgfiles(pkg):
        fpath = pj(pfx.root, fn)
        if not os.path.exists(fpath):
            sys.stderr.write("tpkg: warning: %s does not exist\n" % (fn))
        else:
            fdat = pfx.filedata(fn)
            with open(fpath, "rb") as fp:
                if digest(fp) != fdat.get("digest", ""):
                    sys.stderr.write("tpkg: %s does not match registered hash\n" % (fn))
                    sys.exit(1)
    for fn in pfx.pkgfiles(pkg):
        fpath = pj(pfx.root, fn)
        try:
            os.unlink(fpath)
        except FileNotFoundError:
            pass
        pfx.unregfile(fn)

cmds = {}

def cmd_install(argv):
    def usage(out):
        out.write("usage: tpkg install [-n NAME] SOURCEDIR\n")
    opts, args = getopt.getopt(argv, "n:")
    pkgname = None
    for o, a in opts:
        if o == "-n":
            pkgname = a
    if len(args) < 1:
        usage(sys.stderr)
        sys.exit(1)
    srcpath = args[0]
    if not os.path.isdir(srcpath):
        sys.stderr.write("tpkg: %s: not a directory\n" % (srcpath))
        sys.exit(1)
    if pkgname is None:
        pkgname = os.path.basename(os.path.realpath(srcpath))
    if not pkgname:
        sys.stderr.write("tpkg: could not determine package name\n")
        sys.exit(1)
    install(prefix.use, vfspkg(srcpath), pkgname)
cmds["install"] = cmd_install

def cmd_uninstall(argv):
    def usage(out):
        out.write("usage: tpkg uninstall NAME\n")
    opts, args = getopt.getopt(argv, "")
    if len(args) < 1:
        usage(sys.stderr)
        sys.exit(1)
    pkgname = args[0]
    uninstall(prefix.use, pkgname)
cmds["uninstall"] = cmd_uninstall

def cmd_list(argv):
    def usage(out):
        out.write("usage: tpkg list NAME\n")
    opts, args = getopt.getopt(argv, "")
    if len(args) < 1:
        usage(sys.stderr)
        sys.exit(1)
    pkgname = args[0]
    try:
        files = prefix.use.pkgfiles(pkgname)
    except KeyError:
        sys.stderr.write("tpkg: %s: no such package\n" % (pkgname))
        sys.exit(1)
    for fn in files:
        sys.stdout.write("%s\n" % pj(prefix.use.root, fn))
cmds["list"] = cmd_list

def usage(file):
    file.write("usage:\ttpkg help\n")
cmds["help"] = lambda argv: usage(sys.stdout)

def main(argv):
    pfx = None
    opts, args = getopt.getopt(argv, "hp:")
    for o, a in opts:
        if o == "-h":
            usage(sys.stdout)
            sys.exit(0)
        elif o == "-p":
            if a == "sys":
                pfx = prefix.home()
            elif a == "local":
                pfx = prefix.local()
            elif a == "test":
                pfx = prefix.test()
            else:
                sys.stderr.write("tpkg: %s: undefined prefix\n" % (a))
                sys.exit(1)
    if pfx is None:
        sys.stderr.write("tpkg: no prefix specified\n")
        sys.exit(1)
    prefix.use = pfx
    try:
        if len(args) > 0 and args[0] in cmds:
            cmds[args[0]](args[1:])
            pfx.maint()
            sys.exit(0)
        else:
            usage(sys.stderr)
            sys.exit(1)
    finally:
        pfx.close()

if __name__ == "__main__":
    main(sys.argv[1:])