# Copyright 2003, 2004 Infrae
# Released under the BSD license (see LICENSE.txt)
from __future__ import nested_scopes
import urllib2
from urllib import urlencode
from StringIO import StringIO
from types import SliceType
import libxml2
import time

WAIT_DEFAULT = 120 # two minutes
WAIT_MAX = 5

class Error(Exception):
    pass

# XXX are URIs bytes or unicode? right now we make sure they're bytes

def xpathEvalBytes(context, expr):
    return str(context.xpathEval(expr))

def xpathEvalBytesList(context, expr):
    result = []
    for node in context.xpathEval(expr):
        result.append(str(node.serialize()))
    return result

def xpathEvalText(context, expr):
    return unicode(context.xpathEval(expr), 'utf8')

def xpathEvalTextList(context, expr):
    result = []
    for node in context.xpathEval(expr):
        result.append(unicode(node.serialize('UTF-8'), 'utf8'))
    return result

class MetadataSchemaRegistry:
    def __init__(self):
        self._metadata_schemas = {}
        
    def addMetadataSchema(self, metadata_schema):
        self._metadata_schemas[
            metadata_schema.getMetadataPrefix()] = metadata_schema
        
    def createMetadata(self, metadata_prefix, xpath_context):
        metadata_schema = self._metadata_schemas.get(metadata_prefix)
        if metadata_schema is None:
            return None
        return metadata_schema.createMetadata(xpath_context)

globalMetadataSchemaRegistry = MetadataSchemaRegistry()

addMetadataSchema = globalMetadataSchemaRegistry.addMetadataSchema

class MetadataSchema:
    def __init__(self, metadata_prefix, namespaces):
        self._metadata_prefix = metadata_prefix
        self._namespaces = namespaces
        self._descriptions = {}

    def addFieldDescription(self, field_name, field_type, xpath):
        self._descriptions[field_name] = field_type, xpath

    def createMetadata(self, xpath_context):
        map = {}
        # setup any extra namespaces needed by this schema
        for prefix, uri in self._namespaces.items():
            xpath_context.xpathRegisterNs(prefix, uri)
        # now extra field info according to xpath expr
        for field_name, (field_type, expr) in self._descriptions.items():
            if field_type == 'bytes':
                value = xpathEvalBytes(xpath_context, expr)
            elif field_type == 'bytesList':
                value = xpathEvalBytesList(xpath_context, expr)
            elif field_type == 'text':
                value = xpathEvalText(xpath_context, expr)
            elif field_type == 'textList':
                value = xpathEvalTextList(xpath_context, expr)
            else:
                raise Error, "Unknown field type: %s" % field_type
            map[field_name] = value
        return Metadata(map)

    def getMetadataPrefix(self):
        return self._metadata_prefix

class Header:
    def __init__(self, identifier, datestamp, setspec, deleted):
        self._identifier = identifier
        self._datestamp = datestamp
        self._setspec = setspec
        self._deleted = deleted
        
    def identifier(self):
        return self._identifier
    
    def datestamp(self):
        return self._datestamp

    def setSpec(self):
        return self._setspec

    def isDeleted(self):
        return self._deleted
    
class Metadata:
    def __init__(self, map):
        self._map = map

    def getMap(self):
        return self._map
    
    def getField(self, name):
        return self._map[name]

    __getitem__ = getField

class ServerIdentify:
    def __init__(self, repositoryName, baseURL, protocolVersion, adminEmails,
                 earliestDatestamp, deletedRecord, granularity, compression):
        self._repositoryName = repositoryName
        self._baseURL = baseURL
        self._protocolVersion = protocolVersion
        self._adminEmails = adminEmails
        self._earliestDatestamp = earliestDatestamp
        self._deletedRecord = deletedRecord
        self._granularity = granularity
        self._compression = compression
        # XXX description
        
    def repositoryName(self):
        return self._repositoryName

    def baseURL(self):
        return self._baseURL

    def protocolVersion(self):
        return self._protocolVersion

    def adminEmails(self):
        return self._adminEmails

    def earliestDatestamp(self):
        return self._earliestDatestamp

    def deletedRecord(self):
        return self._deletedRecord

    def granularity(self):
        return self._granularity

    def compression(self):
        return self._compression
    
def buildHeader(xpath_context):
    identifier = xpathEvalBytes(
        xpath_context,
        'string(oai:identifier/text())')
    datestamp = xpathEvalBytes(
        xpath_context,
        'string(oai:datestamp/text())')
    setspec = xpathEvalBytesList(
        xpath_context,
        'oai:setSpec/text()')
    deleted = xpath_context.xpathEval(
        "@status = 'deleted'") 
    return Header(identifier, datestamp, setspec, deleted)

def buildRecords(server, metadata_prefix, namespaces, schema_registry, xml):
    dom = server.parse(xml)
    context = dom.xpathNewContext()
    try:
        for prefix, uri in namespaces.items():
            context.xpathRegisterNs(prefix, uri)
        # first find resumption token if available
        token = context.xpathEval(
            'string(/oai:OAI-PMH/*/oai:resumptionToken/text())')
        if token.strip() == '':
            token = None      
        record_nodes = context.xpathEval('/oai:OAI-PMH/*/oai:record')
        result = []
        for record_node in record_nodes:
            # go into header context
            context.setContextNode(record_node)
            header_node = context.xpathEval(
                'oai:header')[0]
            context.setContextNode(header_node)
            # create header
            header = buildHeader(context)
            # go into metadata context
            context.setContextNode(record_node)
            metadata_list = context.xpathEval(
                'oai:metadata')
            if metadata_list:
                metadata_node = metadata_list[0]
                context.setContextNode(metadata_node)
                # create metadata
                metadata = schema_registry.createMetadata(
                    metadata_prefix, context)
            else:
                metadata = None
            # XXX TODO: about, should be third element of tuple
            result.append((header, metadata, None))
    finally:
        dom.freeDoc()
        context.xpathFreeContext()
    return result, token

def buildIdentifiers(server, namespaces, xml):
    dom = server.parse(xml)
    context = dom.xpathNewContext()
    try:
        for prefix, uri in namespaces.items():
            context.xpathRegisterNs(prefix, uri)
        # first find resumption token is available
        token = context.xpathEval(
            'string(/oai:OAI-PMH/oai:ListIdentifiers/oai:resumptionToken/text())')
        if token.strip() == '':
            token = None    
        header_nodes = context.xpathEval(
                '/oai:OAI-PMH/oai:ListIdentifiers/oai:header')            
        result = []
        for header_node in header_nodes:
            context.setContextNode(header_node)
            header = buildHeader(context)
            result.append(header)
    finally:
        dom.freeDoc()
        context.xpathFreeContext()
    return result, token

def buildSets(server, namespaces, xml):
    dom = server.parse(xml)
    context = dom.xpathNewContext()
    try:
        for prefix, uri in namespaces.items():
            context.xpathRegisterNs(prefix, uri)
        # first find resumption token if available
        token = context.xpathEval(
            'string(/oai:OAI-PMH/oai:ListSets/oai:resumptionToken/text())')
        if token.strip() == '':
            token = None  
        set_nodes = context.xpathEval(
            '/oai:OAI-PMH/oai:ListSets/oai:set')
        sets = []
        for set_node in set_nodes:
            context.setContextNode(set_node)
            setSpec = xpathEvalBytes(
                context,
                'string(oai:setSpec/text())')
            setName = xpathEvalText(
                context,
                'string(oai:setName/text())')
            # XXX setDescription nodes
            sets.append((setSpec, setName, None))
    finally:
        dom.freeDoc()
        context.xpathFreeContext()
    return sets, token

def GetRecord(server, args, xml):
    records, token = buildRecords(
        server,
        args['metadataPrefix'],
        server.getNamespaces(),
        server.getMetadataSchemaRegistry(),
        xml
        )
    assert token is None
    return records[0]
    
def Identify(server, args, xml):
    dom = server.parse(xml)
    context = dom.xpathNewContext()
    try:
        for prefix, uri in server.getNamespaces().items():
            context.xpathRegisterNs(prefix, uri)
        identify_node = context.xpathEval(
            '/oai:OAI-PMH/oai:Identify')[0]
        context.setContextNode(identify_node)
        repositoryName = xpathEvalText(
            context, 
            'string(oai:repositoryName/text())')
        baseURL = xpathEvalBytes(
            context,
            'string(oai:baseURL/text())')
        protocolVersion = xpathEvalBytes(
            context,
            'string(oai:protocolVersion/text())')
        adminEmails = xpathEvalBytesList(
            context,
            'oai:adminEmail/text()')
        earliestDatestamp = xpathEvalBytes(
            context, 
            'string(oai:earliestDatestamp/text())')
        deletedRecord = xpathEvalBytes(
            context,
            'string(oai:deletedRecord/text())')
        granularity = xpathEvalBytes(
            context,
            'string(oai:granularity/text())')
        compression = xpathEvalBytesList(
            context,
            'oai:compression/text()')
        # XXX description
        identify = ServerIdentify(repositoryName, baseURL, protocolVersion,
                                  adminEmails, earliestDatestamp,
                                  deletedRecord, granularity, compression)
    finally:
        dom.freeDoc()
        context.xpathFreeContext()
    return identify

class ResumptionList:
    def __init__(self, server, args, xml):
        self._server = server
        self._namespaces = server.getNamespaces()
        self._result, self._token = self.firstBatch(xml)
        self._max = args.get('max', None)
        
    def __getitem__(self, i):
        if type(i) == SliceType:
            raise Error, "Can't handle slices for resumption list"
        # XXX wish we could use generators/iterators but need to stay
        # python 2.1 compatible
        if i < len(self._result):
            return self._result[i]
        elif self._token is not None:
            # stop if we're over max already
            # this is not very accurate but max is only to have
            # an approximate cap anyway, and this performs better than
            # doing check each time (XXX or is this worth it?)
            if self._max is not None and i >= self._max:
                raise IndexError
            entries, self._token = self.nextBatch(self._token) 
            self._result.extend(entries)
            return self._result[i]
        else:
            raise IndexError

class ListIdentifiers(ResumptionList):
    def firstBatch(self, xml):
        return buildIdentifiers(self._server, self._namespaces, xml)

    def nextBatch(self, token):
        xml = self._server.makeRequest(
            verb= 'ListIdentifiers',
            resumptionToken=self._token)
        return buildIdentifiers(self._server, self._namespaces, xml)

def ListMetadataFormats(server, args, xml):
    #XXX args always thrown away?
    dom = server.parse(xml)
    context = dom.xpathNewContext()
    try:
        for prefix, uri in server.getNamespaces().items():
            context.xpathRegisterNs(prefix, uri)
        metadataFormat_nodes = context.xpathEval(
            '/oai:OAI-PMH/oai:ListMetadataFormats/oai:metadataFormat')
        metadataFormats = []
        for metadataFormat_node in metadataFormat_nodes:
            context.setContextNode(metadataFormat_node)
            metadataPrefix = xpathEvalBytes(
                context,
                'string(oai:metadataPrefix/text())')
            schema = xpathEvalBytes(
                context,                 
                'string(oai:schema/text())')
            metadataNamespace = xpathEvalBytes(
                context,
                'string(oai:metadataNamespace/text())')
            metadataFormat = (metadataPrefix, schema, metadataNamespace)
            metadataFormats.append(metadataFormat)
    finally:
        dom.freeDoc()
        context.xpathFreeContext()
    return metadataFormats

class ListRecords(ResumptionList):
    def __init__(self, server, args, xml):
        # have to do this before calling init..
        self._metadata_prefix = args['metadataPrefix']
        self._metadata_schema_registry = server.getMetadataSchemaRegistry()
        ResumptionList.__init__(self, server, args, xml)

    def firstBatch(self, xml):
        return buildRecords(
            self._server,
            self._metadata_prefix, self._namespaces,
            self._metadata_schema_registry,
            xml)

    def nextBatch(self, token):
        xml = self._server.makeRequest(
            verb= 'ListRecords',
            resumptionToken=self._token)
        return buildRecords(
            self._server,
            self._metadata_prefix, self._namespaces,
            self._metadata_schema_registry, xml)

class ListSets(ResumptionList):
    def firstBatch(self, xml):
        return buildSets(self._server, self._namespaces, xml)

    def nextBatch(self, token):
        xml = self._server.makeRequest(
            verb='ListSets',
            resumptionToken=self._token)
        return buildSets(self._server, self._namespaces, xml)

class OAIMethodError(Exception):
    pass

class OAIMethodImpl:
    def __init__(self, verb, argspec, factory):
        self._factory = factory
        self._verb = verb
        self._argspec = argspec
        self._exclusive = None
        for arg_name, arg_type in argspec.items():
            if arg_type == 'exclusive':
                self._exclusive = arg_name

    def _processArguments(self, dict):
        argspec = self._argspec
        # first filter out any local arguments, which will be returned
        local = {}
        for key, value in argspec.items():
            if value == 'local' and dict.has_key(key):
                local[key] = dict[key]
                del dict[key]
        # check if we have unknown arguments
        for key, value in dict.items():
            if not argspec.has_key(key):
                raise OAIMethodError, "Unknown argument: %s" % key
        # first investigate if we have exclusive argument
        if dict.has_key(self._exclusive):
            if len(dict) > 1:
                raise OAIMethodError, "Exclusive argument %s is used but other arguments found." % self._exclusive
            return
        # if not exclusive, check for required
        for arg_name, arg_type in argspec.items(): 
            if arg_type == 'required':
                if not dict.has_key(arg_name):
                    raise OAIMethodError, "Argument required but not found: %s" % arg_name
        return local
    
    def __call__(self, bound_self, **kw):
        # deal with 'from' (python keyword)
        if kw.has_key('from_'):
            kw['from'] = kw['from_']
            del kw['from_']
        local = self._processArguments(kw)
        kw['verb'] = self._verb
        # reconstruct all arguments XXX hack
        args = kw.copy()
        args.update(local)
        return self._factory(bound_self, args, bound_self.makeRequest(**kw))

def OAIMethod(verb, argspec, factory):
    obj = OAIMethodImpl(verb, argspec, factory)
    def method(self, **kw):
        return obj(self, **kw)
    return method

class BaseServerProxy:
    def __init__(self, metadataSchemaRegistry=None):
        self._metadata_schema_registry = (
            metadataSchemaRegistry or globalMetadataSchemaRegistry)
        self._ignore_bad_character_hack = 0

    def addMetadataSchema(self, schema):
        """Add metadata schema to registry.
        XXX what if this is the global registry, do we want to allow that?
        """
        self._metadata_schema_registry.addMetadataSchema(schema)

    def getNamespaces(self):
        """Get OAI namespaces.
        """
        return {'oai': 'http://www.openarchives.org/OAI/2.0/'}

    def getMetadataSchemaRegistry(self):
        """Get metadata schema registry for this server.
        """
        return self._metadata_schema_registry
    
    def makeRequest(self, **kw):
        """Actually retrieve XML from the server.
        """
        raise NotImplementedError

    def ignoreBadCharacters(self, true_or_false): 	 
        """Set to ignore bad characters in UTF-8 input. 	 
        This is a hack to get around well-formedness errors of 	 
        input sources which *should* be in UTF-8 but for some reason 	 
        aren't completely. 	 
        """ 	 
        self._ignore_bad_character_hack = true_or_false 	 

    def parse(self, xml): 	 
        """Parse the XML to a libxml2 DOM tree. 	 
        """ 	 
        # XXX this is only safe for UTF-8 encoded content, 	 
        # and we're basically hacking around non-wellformedness anyway, 	 
        # but oh well 	 
        if self._ignore_bad_character_hack: 	 
            xml = unicode(xml, 'UTF-8', 'replace') 	 
            xml = xml.encode('UTF-8') 	 
            # also get rid of character code 12 	 
            xml = xml.replace(chr(12), '?') 	 
        return libxml2.parseMemory(xml, len(xml))
        
    getRecord = OAIMethod(
        'GetRecord',
        {'identifier':'required',
        'metadataPrefix':'required'},
        GetRecord
        )
    
    identify = OAIMethod(
        'Identify',
        {},
        Identify
        )

    listIdentifiers = OAIMethod(
        'ListIdentifiers',
        {'from':'optional',
         'until':'optional',
         'metadataPrefix':'required',
         'set':'optional',
         'resumptionToken':'exclusive',
         'max':'local'
         },
        ListIdentifiers
        )

    listMetadataFormats = OAIMethod(
        'ListMetadataFormats',
        {'identifier':'optional'},
        ListMetadataFormats
        )

    listRecords = OAIMethod(
        'ListRecords',
        {'from':'optional',
         'until':'optional',
         'set':'optional',
         'resumptionToken':'exclusive',
         'metadataPrefix':'required',
         'max':'local',
         },
        ListRecords
        )

    listSets = OAIMethod(
        'ListSets',
        {'resumptionToken':'exclusive',
         'max':'local'
         },
        ListSets
        )
    
class ServerProxy(BaseServerProxy):
    def __init__(self, base_url, metadataSchemaRegistry=None):
        BaseServerProxy.__init__(self, metadataSchemaRegistry)
        self._base_url = base_url
        
    def makeRequest(self, **kw):
        """Actually retrieve XML from the server.
        """
        # XXX include From header?
        headers = {'User-Agent': 'Infrae-oaipmh',
                   }
        request = urllib2.Request(self._base_url, urlencode(kw), headers)
        return retrieveFromUrlWaiting(request)

def retrieveFromUrlWaiting(request,
                           wait_max=WAIT_MAX, wait_default=WAIT_DEFAULT):
    """Get text from URL, handling 503 Retry-After.
    """
    for i in range(wait_max):
        try:
            f = urllib2.urlopen(request)
            text = f.read()
            f.close()
            # we successfully opened without having to wait
            break
        except urllib2.HTTPError, e:
            if e.code == 503:
                try:
                    retryAfter = int(e.hdrs.get('Retry-After'))
                except ValueError:
                    retryAfter = None
                if retryAfter is None:
                    time.sleep(wait_default)
                else:
                    time.sleep(retryAfter)
            else:
                # reraise any other HTTP error
                raise
    else:
        raise Error, "Waited too often (more than %s times)" % wait_max
    return text

def register_oai_dc(server_proxy=None):
    """Register OAI DC with server proxy, or global registry.
    """
    oai_dc = MetadataSchema(
        'oai_dc',
        {'oai_dc': 'http://www.openarchives.org/OAI/2.0/oai_dc/',
         'dc' : 'http://purl.org/dc/elements/1.1/'}
        )
    
    fields = [
        ('title', 'oai_dc:dc/dc:title/text()'),
        ('creator', 'oai_dc:dc/dc:creator/text()'),
        ('subject', 'oai_dc:dc/dc:subject/text()'),
        ('description', 'oai_dc:dc/dc:description/text()'),
        ('publisher', 'oai_dc:dc/dc:publisher/text()'),
        ('contributor', 'oai_dc:dc/dc:contributor/text()'),
        ('date', 'oai_dc:dc/dc:date/text()'),
        ('type', 'oai_dc:dc/dc:type/text()'),
        ('format', 'oai_dc:dc/dc:format/text()'),
        ('identifier', 'oai_dc:dc/dc:identifier/text()'),
        ('source', 'oai_dc:dc/dc:source/text()'),
        ('language', 'oai_dc:dc/dc:language/text()'),
        ('relation', 'oai_dc:dc/dc:relation/text()'),
        ('coverage', 'oai_dc:dc/dc:coverage/text()'),
        ('rights', 'oai_dc:dc/dc:rights/text()')
        ]
    
    for name, xpath in fields:
        oai_dc.addFieldDescription(name, 'textList', xpath)

    if server_proxy is not None:
        server_proxy.addMetadataSchema(oai_dc)
    else:
        addMetadatSchema(oai_dc)

