from gi.repository import OSTree
from gi.repository import Gio
from gi.repository import GLib

import os


class OSTreeRepo(object):

    """Class to wrap the OSTree API."""

    OSTREE_GIO_FAST_QUERYINFO = 'standard::name,standard::type,standard::' \
                                'size,standard::is-symlink,standard::syml' \
                                'ink-target,unix::device,unix::inode,unix' \
                                '::mode,unix::uid,unix::gid,unix::rdev'
    G_FILE_QUERY_INFO_NOFOLLOW_SYMLINKS = Gio.FileQueryInfoFlags(1)
    cancellable = Gio.Cancellable.new()
    
    def __init__(self, path, disable_fsync=True):
        self.path = path
        self.repo = self._open_repo(path, disable_fsync)

    def _open_repo(self, path, disable_fsync=True):
        """Create and open and OSTree.Repo, and return it."""
        repo_dir = Gio.file_new_for_path(path)
        repo = OSTree.Repo.new(repo_dir)
        try:
            repo.open(self.cancellable)
        except GLib.GError:
            if not os.path.exists(path):
                os.makedirs(path)
            repo.create(OSTree.RepoMode.ARCHIVE_Z2, self.cancellable)
        repo.set_disable_fsync(disable_fsync)
        return repo
    
    def refsdir(self):
        """Return the abspath to the refs/heads directory in the repo."""
        return os.path.join(os.path.abspath(self.path), 'refs/heads')

    def touch_ref(self, ref):
        """Update the mtime of a ref file in repo/refs/heads."""
        os.utime(os.path.join(self.refsdir(), ref), None)

    def resolve_rev(self, branch, allow_noent=True):
        """Return the SHA256 corresponding to 'branch'."""
        return self.repo.resolve_rev(branch, allow_noent)[1]
    
    def read_commit(self, branch):
        """Return an OSTree.RepoFile representing a committed tree."""
        return self.repo.read_commit(branch, self.cancellable)[1]
    
    def query_info(self, file_object):
        """Quickly return a Gio.FileInfo for file_object."""
        return file_object.query_info(self.OSTREE_GIO_FAST_QUERYINFO,
                                      self.G_FILE_QUERY_INFO_NOFOLLOW_SYMLINKS,
                                      self.cancellable)
    
    def checkout(self, branch, destdir):
        """Checkout branch into destdir."""
        checkout_path = destdir
        if not os.path.exists(checkout_path):
            os.makedirs(checkout_path)
        checkout = Gio.file_new_for_path(checkout_path)
    
        commit = self.read_commit(branch)
        commit_info = self.query_info(commit)
        self.repo.checkout_tree(0, 1, checkout, commit,
                                commit_info, self.cancellable)

    def commit(self, subject, srcdir, branch):
        """Commit the contents of 'srcdir' to 'branch'."""
        self.repo.prepare_transaction(self.cancellable)
        parent = self.resolve_rev(branch)
        if parent:
            parent_root = self.read_commit(parent)

        mtree = OSTree.MutableTree()
        src = Gio.file_new_for_path(srcdir)
        self.repo.write_directory_to_mtree(src, mtree, None, self.cancellable)
        root = self.repo.write_mtree(mtree, self.cancellable)[1]
        if parent and root.equal(parent_root):
            return
        checksum = self.repo.write_commit(parent, subject, '', None,
                                          root, self.cancellable)[1]
        self.repo.transaction_set_ref(None, branch, checksum)
        stats = self.repo.commit_transaction(self.cancellable)

    def cat_file(self, ref, path):
        """Return the file descriptor of path at ref."""
        commit = self.read_commit(ref)
        relative = commit.resolve_relative_path(path)
        ret, content, etag = relative.load_contents()
        return content

    def list_refs(self, resolved=False):
        """Return a list of all refs in the repo."""
        refs = self.repo.list_refs()[1]
        if not resolved:
            return refs.keys()
        return refs

    def delete_ref(self, ref):
        """Remove refspec from the repo."""
        if not self.list_refs(ref):
            raise Exception("Failed to delete ref, it doesn't exist")
        self.repo.set_ref_immediate(None, ref, None, self.cancellable)

    def prune(self):
        """Remove unreachable objects from the repo."""
        return self.repo.prune(OSTree.RepoPruneFlags.REFS_ONLY,
                               -1, self.cancellable)

    def add_remote(self, name, url):
        """Add a remote with a given name and url."""
        options_type = GLib.VariantType.new('a{sv}')
        options_builder = GLib.VariantBuilder.new(options_type)
        options = options_builder.end()
        self.repo.remote_add(name, url, options, self.cancellable)

    def remove_remote(self, name):
        """Remove a remote with a given name."""
        self.repo.remote_delete(name, self.cancellable)

    def get_remote_url(self, name):
        """Return the URL for a remote."""
        return self.repo.remote_get_url(name)[1]

    def list_remotes(self):
        """Return a list of all remotes for this repo."""
        return self.repo.remote_list()

    def has_remote(self, name):
        """Return True if name is a remote for the repo."""
        return name in self.list_remotes()

    def pull(self, refs, remote):
        """Pull ref from remote into the local repo."""
        flags = OSTree.RepoPullFlags.NONE
        self.repo.pull(remote, refs, flags, None, self.cancellable)