import os
import re
import fnmatch
import datetime

import simplejson as json
import mercurial
import mercurial.ui
import mercurial.localrepo
import mercurial.context

class Directory(object):
    def __init__(self, name='', parent=None, pattern=re.compile('.*'),
                 file_changesets=None):
        if isinstance(pattern, basestring):
            pattern = re.compile(fnmatch.translate(pattern))

        self.name = name
        if parent:
            self.path = os.path.join(parent.path, name)
        else:
            self.path = name

        if self.path:
            rawfiles = os.listdir(self.path)
        else:
            rawfiles = os.listdir('.')

        self.parent = parent
        self.files = []
        self.subdirs = []
        files = [filename for filename in rawfiles
                 if filename[0] != '.']
        for filename in files:
            file_path = os.path.join(self.path, filename)
            if (pattern.match(filename) and
                os.path.isfile(file_path)):
                self.files.append(File(filename, self, file_changesets))
            elif (os.path.isdir(file_path) and
                  'test' not in file_path):
                self.subdirs.append(Directory(filename,
                                              self,
                                              pattern,
                                              file_changesets))
        file_sizes = [file.size for file in self.files]
        subdir_sizes = [subdir.size for subdir in self.subdirs]
        self.size = sum(file_sizes) + sum(subdir_sizes)

        self.changes = set()
        for file in self.files:
            self.changes.update(file.changes)
        for subdir in self.subdirs:
            self.changes.update(subdir.changes)

    def __repr__(self):
        return "<Directory '%s'>" % self.path

class File(object):
    def __init__(self, filename, parent, file_changesets):
        self.parent = parent
        self.name = filename
        self.path = os.path.join(parent.path, filename)
        self.size = os.stat(self.path).st_size
        if self.path in file_changesets:
            self.changes = file_changesets[self.path]
        else:
            self.changes = set()

    def __repr__(self):
        return "<File '%s'>" % self.name

def sort_by_size(a, b):
    return cmp(a.size, b.size)

def get_json(directory, depth=-1):
    if depth == 0:
        return

    subdirs = [subdir for subdir in directory.subdirs
               if subdir.size > 0]
    subdirs.sort(sort_by_size, reverse=True)

    obj = {'name': directory.name,
           'size': directory.size,
           'changes': len(directory.changes)}

    json_subdirs = []
    for subdir in subdirs:
        json_subdirs.append(get_json(subdir, depth - 1))
    if json_subdirs:
        obj['subdirs'] = json_subdirs

    return obj

def get_file_changesets(path):
    a_year_ago = datetime.datetime.now() - datetime.timedelta(days=365)
    ui = mercurial.ui.ui()
    repo = mercurial.localrepo.localrepository(ui, path)
    ctx = mercurial.context.changectx(repo)
    to_visit = [ctx]
    files = {}

    while to_visit:
        ctx = to_visit.pop()
        rev = ctx.rev()
        date = datetime.datetime.fromtimestamp(ctx.date()[0])
        if date < a_year_ago:
            break
        for filename in ctx.files():
            if 'test' in filename:
                continue
            if filename not in files:
                files[filename] = set()
            files[filename].add(rev)
        for anc in ctx.parents():
            to_visit.append(anc)

    return files

if __name__ == '__main__':
    import sys

    if len(sys.argv) < 2:
        print ("Please pass the path to the root of the mozilla "
               "tree as the first argument to this program.")
        sys.exit(1)

    if len(sys.argv) > 2:
        pattern = sys.argv[2]
    else:
        pattern = "*.*"

    file_changesets = get_file_changesets(sys.argv[1])

    curdir = os.getcwd()
    os.chdir(sys.argv[1])
    directory = Directory(pattern=pattern,
                          file_changesets=file_changesets)
    os.chdir(curdir)

    filename = 'moztree.json'
    open(filename, 'w').write(json.dumps(get_json(directory)))
    print "Wrote '%s'." % filename