#! /usr/bin/env python2.3

# retrieve full text from DL
import os.path
from datetime import datetime
import urllib2
import base64
import optparse
from oaipmh import metadata, error, client

# register metadata reader
prefix = 'dl'
namespaces = {prefix: 'http://infrae.com/ns/documentlibrary'}
fields = {
    'plaintext_url': ('textList', 'dl:dl/dl:file_urls/dl:plaintext_url/text()')
    }

reader = metadata.MetadataReader(fields, namespaces)
metadata.global_metadata_registry.registerReader(prefix, reader)

def _fetchAndStore(url, filename, username=None, passwd=None, redo=None):
    request = urllib2.Request(url)

    if username is not None and passwd is not None:
        info = '%s:%s' % (username, passwd)
        credentials = base64.encodestring(urllib2.unquote(info))
    else:
        credentials = None
            
    if credentials is not None:
        request.add_header('Authorization', 'Basic ' + credentials)

    dst = open(filename, 'w')
    # XXX support the hacked up fulltext storage format
    dst.write(url + '\n')
    dst.write('|FULLTEXT|\n\n')
    dst.write('|DOCUMENT|\n\n')
    src = urllib2.urlopen(request)
    blocksize = 2<<16
    while 1:
        buffer = src.read(blocksize)
        dst.write(buffer)
        if len(buffer) < blocksize:
            break
    dst.close()
        
def retrieval(
    url, storage_dir, from_=None, until=None, username=None, passwd=None, 
    redo=False):
 
    if username is not None:
        credentials=(username, passwd)
    else:
        credentials = None
    documentlibrary = client.Client(url, credentials=credentials) # at a later stage we provide credentials
    
    # XXX this should not be needed AFAIK
    if until is None:
        until = datetime.utcnow()
    if from_ is None:
        from_ = datetime(1970, 1, 1)

    try:
        records = documentlibrary.listRecords(
            from_=from_, until=until, metadataPrefix=prefix)
    except error.NoRecordsMatchError, e:
        return
    
    for header, metadata, about in records:
        if header.isDeleted():
           continue    
        plaintext_urls = metadata.getMap().get('plaintext_url')
        if not len(plaintext_urls):
            continue
        filename = os.path.join(storage_dir, header.identifier())
        if not os.path.exists(filename) or redo:
            _fetchAndStore(plaintext_urls[0], filename, username, passwd)
        
if __name__ == '__main__':
    # command line UI
    version = '0.1'
    usage = ('usage: %prog [options] OAI_URL FULLTEXT_DIRECTORY\n\n')
    parser = optparse.OptionParser(usage=usage, version=version)
    parser.add_option(
        '--date', dest='from_', default=None, help=(
        'Start date of period for which records are harvested.'
        '(ISO8601 UTC format. YYYY-MM-DD or YYYY-MM-DDThh:mmZ '
        'For example: 1957-03-20 or with time, 1957-03-20T20:30Z)'))
    parser.add_option(
        '--enddate', dest='until', default=None, help=(
        'End date of period for which records are harvested '
        '(ISO8601 UTC format. YYYY-MM-DD or YYYY-MM-DDThh:mmZ '
        'For example: 1957-03-20 or with time, 1957-03-20T20:30Z)'))
    parser.add_option(
        '--username', dest='username', default=None, help=(
        'Username used for connecting to the OAI repository and plain '
        'text URLs'))
    parser.add_option(
        '--passwd', dest='passwd', default=None, help=(
        'Password used for connecting to the OAI repository and plain '
        'text URLs'))
    parser.add_option(
        '--redo', dest='redo', default=False, action='store_true', help=(
        'Retrieve fulltext even the corresponding file already exists.'))
        
    options, args = parser.parse_args()
    url, dir = args
    
    # XXX make sure until and from_ are datetime objects (or None)
    
    retrieval(
        url, dir, options.from_, options.until, options.username, 
        options.passwd, options.redo)
