from pywb.utils.binsearch import iter_range

from pywb.utils.wbexception import AccessException, NotFoundException
from pywb.utils.wbexception import BadRequestException, WbException

from query import CDXQuery

import urllib
import urllib2
import itertools


#=================================================================
class CDXSource(object):
    """
    Represents any cdx index source
    """
    def load_cdx(self, query):  # pragma: no cover
        raise NotImplementedError('Implement in subclass')


#=================================================================
class CDXFile(CDXSource):
    """
    Represents a local plain-text .cdx file
    """
    def __init__(self, filename):
        self.filename = filename

    def load_cdx(self, query):
        def do_open():
            try:
                source = open(self.filename, 'rb')
                gen = iter_range(source, query.key, query.end_key)
                for line in gen:
                    yield line
            finally:
                source.close()

        return do_open()
        #return iter_range(do_open(), query.key, query.end_key)

    def __str__(self):
        return 'CDX File - ' + self.filename


#=================================================================
class RemoteCDXSource(CDXSource):
    """
    Represents a remote cdx server, to which requests will be proxied.

    Only ``url`` and ``match_type`` params are proxied at this time,
    the stream is passed through all other filters locally.
    """
    def __init__(self, filename, cookie=None, remote_processing=False):
        self.remote_url = filename
        self.cookie = cookie
        self.remote_processing = remote_processing

    def load_cdx(self, query):
        if self.remote_processing:
            remote_query = query
        else:
            # Only send url and matchType to remote
            remote_query = CDXQuery(url=query.url,
                                    match_type=query.match_type)

        urlparams = remote_query.urlencode()

        try:
            request = urllib2.Request(self.remote_url + '?' + urlparams)

            if self.cookie:
                request.add_header('Cookie', self.cookie)

            response = urllib2.urlopen(request)

        except urllib2.HTTPError as e:
            if e.code == 403:
                raise AccessException('Access Denied')
            elif e.code == 404:
                # return empty list for consistency with other cdx sources
                # will be converted to 404 if no other retry
                return []
            elif e.code == 400:
                raise BadRequestException()
            else:
                raise WbException('Invalid response from remote cdx server')

        return iter(response)

    def __str__(self):
        if self.remote_processing:
            return 'Remote CDX Server: ' + self.remote_url
        else:
            return 'Remote CDX Source: ' + self.remote_url


#=================================================================
class RedisCDXSource(CDXSource):
    DEFAULT_KEY_PREFIX = 'c:'

    def __init__(self, redis_url, config=None):
        import redis

        parts = redis_url.split('/')
        if len(parts) > 4:
            self.cdx_key = parts[4]
        else:
            self.cdx_key = None

        self.redis_url = redis_url
        self.redis = redis.StrictRedis.from_url(redis_url)

        self.key_prefix = self.DEFAULT_KEY_PREFIX

    def load_cdx(self, query):
        """
        Load cdx from redis cache, from an ordered list

        If cdx_key is set, treat it as cdx file and load use
        zrangebylex! (Supports all match types!)

        Otherwise, assume a key per-url and load all entries for that key.
        (Only exact match supported)
        """

        if self.cdx_key:
            return self.load_sorted_range(query)
        else:
            return self.load_single_key(query.key)

    def load_sorted_range(self, query):
        cdx_list = self.redis.zrangebylex(self.cdx_key,
                                          '[' + query.key,
                                          '(' + query.end_key)

        return cdx_list

    def load_single_key(self, key):
        # ensure only url/surt is part of key
        key = key.split(' ')[0]
        cdx_list = self.redis.zrange(self.key_prefix + key, 0, -1)

        # key is not part of list, so prepend to each line
        key += ' '
        cdx_list = itertools.imap(lambda x: key + x, cdx_list)
        return cdx_list

    def __str__(self):
        return 'Redis - ' + self.redis_url
