#!/usr/bin/env python

#

# Copyright (C) 2007 Brendan Cully <brendan@kublai.com>

#

# This software may be used and distributed according to the terms

# of the GNU General Public License, incorporated herein by reference.

import os, sys

class ConfigError(Exception): pass

def usage():

    print """relink <source> <destination>

    Recreate hard links between source and destination repositories"""

class Config:

    def __init__(self, args):

        if len(args) != 3:

            raise ConfigError("wrong number of arguments")

        self.src = os.path.abspath(args[1])

        self.dst = os.path.abspath(args[2])

        for d in (self.src, self.dst):

            if not os.path.exists(os.path.join(d, '.hg')):

                raise ConfigError("%s: not a mercurial repository" % d)

try:

    cfg = Config(sys.argv)

except ConfigError, inst:

    print str(inst)

    usage()

    sys.exit(1)

def collect(src):

    seplen = len(os.path.sep)

    candidates = []

    for dirpath, dirnames, filenames in os.walk(src):

        relpath = dirpath[len(src) + seplen:]

        for filename in filenames:

            if not (filename.endswith('.i') or filename.endswith('.d')):

                continue

            st = os.stat(os.path.join(dirpath, filename))

            candidates.append((os.path.join(relpath, filename), st))

    return candidates

def prune(candidates, dst):

    targets = []

    for fn, st in candidates:

        tgt = os.path.join(dst, fn)

        try:

            ts = os.stat(tgt)

        except OSError:

            # Destination doesn't have this file?

            continue

        if st.st_ino == ts.st_ino:

            continue

        if st.st_dev != ts.st_dev:

            raise Exception('Source and destination are on different devices')

        if st.st_size != ts.st_size:

            continue

        targets.append((fn, ts.st_size))

    return targets

def relink(src, dst, files):

    CHUNKLEN = 65536

    relinked = 0

    savedbytes = 0

    for f, sz in files:

        source = os.path.join(src, f)

        tgt = os.path.join(dst, f)

        sfp = file(source)

        dfp = file(tgt)

        sin = sfp.read(CHUNKLEN)

        while sin:

            din = dfp.read(CHUNKLEN)

            if sin != din:

                break

            sin = sfp.read(CHUNKLEN)

        if sin:

            continue

        try:

            os.rename(tgt, tgt + '.bak')

            try:

                os.link(source, tgt)

            except OSError:

                os.rename(tgt + '.bak', tgt)

                raise

            print 'Relinked %s' % f

            relinked += 1

            savedbytes += sz

            os.remove(tgt + '.bak')

        except OSError, inst:

            print '%s: %s' % (tgt, str(inst))

    print 'Relinked %d files (%d bytes reclaimed)' % (relinked, savedbytes)

src = os.path.join(cfg.src, '.hg')

dst = os.path.join(cfg.dst, '.hg')

candidates = collect(src)

targets = prune(candidates, dst)

relink(src, dst, targets)