#!/usr/bin/env python

"""
$ svn co URL wc1
$ svn co URL wc2 --depth=empty
$ rsync -a wc1/.svn/pristine/ wc2/.svn/pristine/
$ populate-pristine.py wc2
"""

# TODO: increment refcount upon collision
# TODO: add <given file>, not just argv[1]/.svn/pristine/??/*

import hashlib
import os
import re
import sqlite3
import sys

# ### Currently, require this specific format.  This could require any other
# ### format that has the same PRISTINE schema and semantics.
FORMAT = 22
BUFFER_SIZE = 4 * 1024

class UnknownFormat(Exception):
  def __init__(self, formatno):
    self.formatno = formatno

def open_db(wc_path):
  wc_db = os.path.join(wc_path, '.svn', 'wc.db')
  conn = sqlite3.connect(wc_db)
  curs = conn.cursor()
  curs.execute('pragma user_version;')
  formatno = int(curs.fetchone()[0])
  if formatno > FORMAT:
    raise UnknownFormat(formatno)
  return conn

_sha1_re = re.compile(r'^[0-9a-f]{40}$')

def md5_of(path):
  fd = os.open(path, os.O_RDONLY)
  ctx = hashlib.md5()
  while True:
    s = os.read(fd, BUFFER_SIZE)
    if len(s):
      ctx.update(s)
    else:
      os.close(fd)
      return ctx.hexdigest()

def populate(wc_path):
  conn = open_db(wc_path)
  for dirname, dirs, files in os.walk(os.path.join(wc_path, '.svn', 'pristine')):
    # skip everything but .svn/pristine/xx/
    if os.path.basename(os.path.dirname(dirname)) == 'pristine':
      sys.stdout.write("Updating '%s'..." % os.path.basename(dirname))
      for f in filter(lambda x: _sha1_re.match(x), files):
      	fullpath = os.path.join(dirname, f)
      	conn.execute("""INSERT OR REPLACE INTO pristine VALUES (:checksum,:compression,:size,:refcount,:md5_checksum)""",
      	             ('$sha1$'+f, None, os.stat(fullpath).st_size, 1, '$md5 $'+md5_of(fullpath)))
      # periodic transaction commits, for efficiency
      conn.commit()
      sys.stdout.write("\n")




if __name__ == '__main__':
  paths = sys.argv[1:]
  if not paths:
    paths = ['.']
  for wc_path in paths:
    try:
      populate(wc_path)
    except UnknownFormat, e:
      print "Don't know how to handle '%s' (format %d)'" % (wc_path, e.formatno)