# -*- coding: utf-8 -*-
from functools import partial
from itertools import izip
import logging
import weakref

from tornado.ioloop import IOLoop
from tornado import gen

from datetime import datetime
from .exceptions import RequestError, ConnectionError, ResponseError
from .connection import Connection


log = logging.getLogger('tornadoredis.client')


class Message(object):
    ''' Wrapper Message object.
        kind = command
       channel = channel from which the message was received
        pattern = subscription pattern
        body = message body
    '''
    def __init__(self, *args):
        if len(args) == 3:
            (self.kind, self.channel, self.body) = args
            self.pattern = self.channel
        elif len(args) == 4:
            (self.kind, self.pattern, self.channel, self.body) = args
        else:
            raise ValueError('Invalid number of arguments')


class CmdLine(object):
    def __init__(self, cmd, *args, **kwargs):
        self.cmd = cmd
        self.args = args
        self.kwargs = kwargs

    def __repr__(self):
        return self.cmd + '(' + str(self.args) + ',' + str(self.kwargs) + ')'


def string_keys_to_dict(key_string, callback):
    return dict([(key, callback) for key in key_string.split()])


def dict_merge(*dicts):
    merged = {}
    for d in dicts:
        merged.update(d)
    return merged


def reply_to_bool(r, *args, **kwargs):
    return bool(r)


def make_reply_assert_msg(msg):
    def reply_assert_msg(r, *args, **kwargs):
        return r == msg
    return reply_assert_msg


def reply_set(r, *args, **kwargs):
    return set(r)


def reply_dict_from_pairs(r, *args, **kwargs):
    return dict(izip(r[::2], r[1::2]))


def reply_str(r, *args, **kwargs):
    return r or ''


def reply_int(r, *args, **kwargs):
    return int(r) if r is not None else None


def reply_float(r, *args, **kwargs):
    return float(r) if r is not None else None


def reply_datetime(r, *args, **kwargs):
    return datetime.fromtimestamp(int(r))


def reply_pubsub_message(r, *args, **kwargs):
    return Message(*r)


def reply_zset(r, *args, **kwargs):
    if (not r) or (not 'WITHSCORES' in args):
        return r
    return zip(r[::2], map(float, r[1::2]))


def reply_hmget(r, key, *fields, **kwargs):
    return dict(zip(fields, r))


def reply_info(response):
    info = {}

    def get_value(value):
        # Does this string contains subvalues?
        if (',' not in value) or ('=' not in value):
            return value
        sub_dict = {}
        for item in value.split(','):
            k, v = item.split('=')
            try:
                sub_dict[k] = int(v)
            except ValueError:
                sub_dict[k] = v
        return sub_dict
    for line in response.splitlines():
        line = line.strip()
        if line and not line.startswith('#'):
            key, value = line.split(':')
            try:
                info[key] = int(value)
            except ValueError:
                info[key] = get_value(value)
    return info


def reply_ttl(r, *args, **kwargs):
    return r != -1 and r or None


def to_list(source):
    if isinstance(source, basestring):
        return [source]
    else:
        return list(source)


PUB_SUB_COMMANDS = set([
    'SUBSCRIBE',
    'PSUBSCRIBE',
    'UNSUBSCRIBE',
    'PUNSUBSCRIBE',
    # Not a command at all
    'LISTEN',
])


REPLY_MAP = dict_merge(
    string_keys_to_dict('AUTH BGREWRITEAOF BGSAVE DEL EXISTS '
                        'EXPIRE HDEL HEXISTS '
                        'HMSET MOVE MSET MSETNX SAVE SETNX',
                        reply_to_bool),
    string_keys_to_dict('FLUSHALL FLUSHDB SELECT SET SETEX '
                        'SHUTDOWN RENAME RENAMENX WATCH UNWATCH',
                        make_reply_assert_msg('OK')),
    string_keys_to_dict('SMEMBERS SINTER SUNION SDIFF',
                        reply_set),
    string_keys_to_dict('HGETALL BRPOP BLPOP',
                        reply_dict_from_pairs),
    string_keys_to_dict('HGET',
                        reply_str),
    string_keys_to_dict('SUBSCRIBE UNSUBSCRIBE LISTEN '
                        'PSUBSCRIBE UNSUBSCRIBE',
                        reply_pubsub_message),
    string_keys_to_dict('ZRANK ZREVRANK',
                        reply_int),
    string_keys_to_dict('ZSCORE ZINCRBY ZCOUNT ZCARD',
                        reply_int),
    string_keys_to_dict('ZRANGE ZRANGEBYSCORE ZREVRANGE '
                        'ZREVRANGEBYSCORE',
                        reply_zset),
    {'HMGET': reply_hmget},
    {'PING': make_reply_assert_msg('PONG')},
    {'LASTSAVE': reply_datetime},
    {'TTL': reply_ttl},
    {'INFO': reply_info},
    {'MULTI_PART': make_reply_assert_msg('QUEUED')},
)


class Client(object):
#    __slots__ = ('_io_loop', '_connection_pool', 'connection', 'subscribed',
#                 'password', 'selected_db', '_pipeline', '_weak')

    def __init__(self, host='localhost', port=6379, password=None,
                 selected_db=None, io_loop=None, connection_pool=None):
        self._io_loop = io_loop or IOLoop.instance()
        self._connection_pool = connection_pool
        self._weak = weakref.proxy(self)
        if connection_pool:
            connection = (connection_pool
                          .get_connection(event_handler_proxy=self._weak))
        else:
            connection = Connection(host=host, port=port,
                                    weak_event_handler=self._weak,
                                    io_loop=self._io_loop)
        self.connection = connection
        self.subscribed = False
        self.password = password
        self.selected_db = selected_db
        self._pipeline = None

    def __del__(self):
        try:
            connection = self.connection
            pool = self._connection_pool
        except AttributeError:
            connection = None
            pool = None
        if connection:
            if pool:
                pool.release(connection)
                connection.wait_until_ready()
            else:
                connection.disconnect()

    def __repr__(self):
        return 'tornadoredis.Client (db=%s)' % (self.selected_db)

    def __enter__(self):
        return self

    def __exit__(self, *args, **kwargs):
        pass

    def __getattribute__(self, item):
        '''
        Bind methods to the weak proxy to avoid memory leaks
        when bound method is passed as argument to the gen.Task
        constructor.
        '''
        a = super(Client, self).__getattribute__(item)
        try:
            if callable(a) and a.__self__:
                try:
                    a = self.__class__.__dict__[item]
                except KeyError:
                    a = Client.__dict__[item]
                a = partial(a, self._weak)
        except AttributeError:
            pass
        return a

    def pipeline(self, transactional=False):
        '''
        Create the 'Pipeline' to pack multiple redis commands
        in a single request.

        Usage:
            pipe = self.client.pipeline()
            pipe.hset('foo', 'bar', 1)
            pipe.expire('foo', 60)

            yield gen.Task(pipe.execute)

        or:

            with self.client.pipeline() as pipe:
                pipe.hset('foo', 'bar', 1)
                pipe.expire('foo', 60)

                yield gen.Task(pipe.execute)
        '''
        if not self._pipeline:
            self._pipeline = Pipeline(
                selected_db=self.selected_db,
                io_loop=self._io_loop,
                transactional=transactional
            )
            self._pipeline.connection = self.connection
        return self._pipeline

    #### Connection event handlers
    def on_connect(self):
        if self.password:
            password = self.connection.info.get('password', None)
            if password != self.password:
                self.auth(self.password)
                self.connection.info['password'] = self.password
        if self.selected_db:
            db = self.connection.info.get('db', None)
            if db != self.selected_db:
                self.select(self.selected_db)
                self.connection.info['db'] = self.selected_db

    def on_disconnect(self):
        if self.subscribed:
            self.subscribed = False
        raise ConnectionError("Socket closed on remote end")

    #### connection
    def connect(self):
        if not self.connection.connected():
            pool = self._connection_pool
            if pool:
                old_conn = self.connection
                self.connection = pool.get_connection(event_handler=self)
                self.connection.ready_callbacks = old_conn.ready_callbacks
            else:
                self.connection.connect()

    @gen.engine
    def disconnect(self, callback=None):
        '''
        Disconnect from the Redis server.
        '''
        connection = self.connection
        if connection:
            pool = self._connection_pool
            if pool:
                pool.release(connection)
                yield gen.Task(connection.wait_until_ready)
                proxy = pool.make_proxy(client_proxy=self._weak,
                                        connected=False)
                self.connection = proxy
            else:
                self.connection.disconnect()
        if callback:
            callback(False)

    #### formatting
    def encode(self, value):
        if isinstance(value, str):
            return value
        elif isinstance(value, unicode):
            return value.encode('utf-8')
        return str(value)

    def format_command(self, *tokens, **kwargs):
        cmds = []
        for t in tokens:
            e_t = self.encode(t)
            cmds.append('$%s\r\n%s\r\n' % (len(e_t), e_t))
        return '*%s\r\n%s' % (len(tokens), ''.join(cmds))

    def format_reply(self, cmd_line, data):
        if cmd_line.cmd not in REPLY_MAP:
            return data
        try:
            res = REPLY_MAP[cmd_line.cmd](data,
                                          *cmd_line.args,
                                          **cmd_line.kwargs)
        except Exception, e:
            raise ResponseError(
                'failed to format reply to %s, raw data: %s; err message: %s'
                % (cmd_line, data, e), cmd_line
            )
        return res
    ####

    @gen.engine
    def execute_command(self, cmd, *args, **kwargs):
        callback = kwargs.get('callback', None)
        del kwargs['callback']
        cmd_line = CmdLine(cmd, *args, **kwargs)
        if self.subscribed and cmd not in PUB_SUB_COMMANDS:
            callback(RequestError(
                'Calling not pub/sub command during subscribed state',
                cmd_line))
            return

        n_tries = 2
        while n_tries > 0:
            n_tries -= 1
            if not self.connection.connected():
                self.connection.connect()

            if not self.connection.ready() and not self.subscribed:
                yield gen.Task(self.connection.wait_until_ready)

            with self.connection as c:
                command = self.format_command(cmd, *args, **kwargs)
                try:
                    yield gen.Task(self.connection.write, command)
                except Exception, e:
                    self.connection.disconnect()
                    if not n_tries:
                        raise e
                    else:
                        continue

                if ((cmd in PUB_SUB_COMMANDS) or
                    (self.subscribed and cmd == 'PUBLISH')):
                    result = True
                    break
                else:
                    result = None
                    data = yield gen.Task(c.readline)
                    if not data:
                        if not n_tries:
                            raise ConnectionError('no data received')
                    else:
                        resp = yield gen.Task(self.process_data,
                                              data,
                                              cmd_line)
                        result = self.format_reply(cmd_line, resp)
                        break

        if callback:
            callback(result)

    @gen.engine
    def process_data(self, data, cmd_line, callback=None):
        data = data[:-2]  # strip \r\n

        if data == '$-1':
            response = None
        elif data == '*0' or data == '*-1':
            response = []
        else:
            head, tail = data[0], data[1:]

            if head == '*':
                response = yield gen.Task(self.consume_multibulk,
                                          int(tail),
                                          cmd_line)
            elif head == '$':
                # Consume bulk reply
                response = yield gen.Task(self.connection.read, int(tail) + 2)
                if isinstance(response, Exception):
                    raise response
                if not response:
                    raise ResponseError('EmptyResponse')
                else:
                    response = response[:-2]
            elif head == '+':
                response = tail
            elif head == ':':
                response = int(tail)
            elif head == '-':
                if tail.startswith('ERR'):
                    tail = tail[4:]
                response = ResponseError(tail, cmd_line)
            else:
                raise ResponseError('Unknown response type %s' % head,
                                    cmd_line)
        callback(response)

    @gen.engine
    def consume_multibulk(self, length, cmd_line, callback=None):
        tokens = []
        while len(tokens) < length:
            data = yield gen.Task(self.connection.readline)
            if not data:
                raise ResponseError(
                    'Not enough data in response to %s, accumulated tokens: %s'
                    % (cmd_line, tokens),
                    cmd_line)
            token = yield gen.Task(self.process_data, data, cmd_line)
            tokens.append(token)

        callback(tokens)

    ### MAINTENANCE
    def bgrewriteaof(self, callback=None):
        self.execute_command('BGREWRITEAOF', callback=callback)

    def dbsize(self, callback=None):
        self.execute_command('DBSIZE', callback=callback)

    def flushall(self, callback=None):
        self.execute_command('FLUSHALL', callback=callback)

    def flushdb(self, callback=None):
        self.execute_command('FLUSHDB', callback=callback)

    def ping(self, callback=None):
        self.execute_command('PING', callback=callback)

    def info(self, callback=None):
        self.execute_command('INFO', callback=callback)

    def select(self, db, callback=None):
        self.selected_db = db
        self.execute_command('SELECT', db, callback=callback)

    def shutdown(self, callback=None):
        self.execute_command('SHUTDOWN', callback=callback)

    def save(self, callback=None):
        self.execute_command('SAVE', callback=callback)

    def bgsave(self, callback=None):
        self.execute_command('BGSAVE', callback=callback)

    def lastsave(self, callback=None):
        self.execute_command('LASTSAVE', callback=callback)

    def keys(self, pattern, callback=None):
        self.execute_command('KEYS', pattern, callback=callback)

    def auth(self, password, callback=None):
        self.execute_command('AUTH', password, callback=callback)

    ### BASIC KEY COMMANDS
    def append(self, key, value, callback=None):
        self.execute_command('APPEND', key, value, callback=callback)

    def expire(self, key, ttl, callback=None):
        self.execute_command('EXPIRE', key, ttl, callback=callback)

    def ttl(self, key, callback=None):
        self.execute_command('TTL', key, callback=callback)

    def type(self, key, callback=None):
        self.execute_command('TYPE', key, callback=callback)

    def randomkey(self, callback=None):
        self.execute_command('RANDOMKEY', callback=callback)

    def rename(self, src, dst, callback=None):
        self.execute_command('RENAME', src, dst, callback=callback)

    def renamenx(self, src, dst, callback=None):
        self.execute_command('RENAMENX', src, dst, callback=callback)

    def move(self, key, db, callback=None):
        self.execute_command('MOVE', key, db, callback=callback)

    def substr(self, key, start, end, callback=None):
        self.execute_command('SUBSTR', key, start, end, callback=callback)

    def delete(self, *keys, **kwargs):
        self.execute_command('DEL', *keys, callback=kwargs.get('callback'))

    def set(self, key, value, callback=None):
        self.execute_command('SET', key, value, callback=callback)

    def setex(self, key, ttl, value, callback=None):
        self.execute_command('SETEX', key, ttl, value, callback=callback)

    def setnx(self, key, value, callback=None):
        self.execute_command('SETNX', key, value, callback=callback)

    def mset(self, mapping, callback=None):
        items = []
        [items.extend(pair) for pair in mapping.iteritems()]
        self.execute_command('MSET', *items, callback=callback)

    def msetnx(self, mapping, callback=None):
        items = []
        [items.extend(pair) for pair in mapping.iteritems()]
        self.execute_command('MSETNX', *items, callback=callback)

    def get(self, key, callback=None):
        self.execute_command('GET', key, callback=callback)

    def mget(self, keys, callback=None):
        self.execute_command('MGET', *keys, callback=callback)

    def getset(self, key, value, callback=None):
        self.execute_command('GETSET', key, value, callback=callback)

    def exists(self, key, callback=None):
        self.execute_command('EXISTS', key, callback=callback)

    def sort(self, key, start=None, num=None, by=None, get=None, desc=False,
             alpha=False, store=None, callback=None):
        if (start is not None and num is None) \
        or (num is not None and start is None):
            raise ValueError("``start`` and ``num`` must both be specified")

        tokens = [key]
        if by is not None:
            tokens.append('BY')
            tokens.append(by)
        if start is not None and num is not None:
            tokens.append('LIMIT')
            tokens.append(start)
            tokens.append(num)
        if get is not None:
            tokens.append('GET')
            tokens.append(get)
        if desc:
            tokens.append('DESC')
        if alpha:
            tokens.append('ALPHA')
        if store is not None:
            tokens.append('STORE')
            tokens.append(store)
        return self.execute_command('SORT', *tokens, callback=callback)

    def getbit(self, key, offset, callback=None):
        self.execute_command('GETBIT', key, offset, callback=callback)

    def setbit(self, key, offset, value, callback=None):
        self.execute_command('SETBIT', key, offset, value, callback=callback)

    ### COUNTERS COMMANDS
    def incr(self, key, callback=None):
        self.execute_command('INCR', key, callback=callback)

    def decr(self, key, callback=None):
        self.execute_command('DECR', key, callback=callback)

    def incrby(self, key, amount, callback=None):
        self.execute_command('INCRBY', key, amount, callback=callback)

    def decrby(self, key, amount, callback=None):
        self.execute_command('DECRBY', key, amount, callback=callback)

    ### LIST COMMANDS
    def blpop(self, keys, timeout=0, callback=None):
        tokens = to_list(keys)
        tokens.append(timeout)
        self.execute_command('BLPOP', *tokens, callback=callback)

    def brpop(self, keys, timeout=0, callback=None):
        tokens = to_list(keys)
        tokens.append(timeout)
        self.execute_command('BRPOP', *tokens, callback=callback)

    def brpoplpush(self, src, dst, timeout=1, callback=None):
        tokens = [src, dst, timeout]
        self.execute_command('BRPOPLPUSH', *tokens, callback=callback)

    def lindex(self, key, index, callback=None):
        self.execute_command('LINDEX', key, index, callback=callback)

    def llen(self, key, callback=None):
        self.execute_command('LLEN', key, callback=callback)

    def lrange(self, key, start, end, callback=None):
        self.execute_command('LRANGE', key, start, end, callback=callback)

    def lrem(self, key, value, num=0, callback=None):
        self.execute_command('LREM', key, num, value, callback=callback)

    def lset(self, key, index, value, callback=None):
        self.execute_command('LSET', key, index, value, callback=callback)

    def ltrim(self, key, start, end, callback=None):
        self.execute_command('LTRIM', key, start, end, callback=callback)

    def lpush(self, key, value, callback=None):
        self.execute_command('LPUSH', key, value, callback=callback)

    def rpush(self, key, value, callback=None):
        self.execute_command('RPUSH', key, value, callback=callback)

    def lpop(self, key, callback=None):
        self.execute_command('LPOP', key, callback=callback)

    def rpop(self, key, callback=None):
        self.execute_command('RPOP', key, callback=callback)

    def rpoplpush(self, src, dst, callback=None):
        self.execute_command('RPOPLPUSH', src, dst, callback=callback)

    ### SET COMMANDS
    def sadd(self, key, value, callback=None):
        self.execute_command('SADD', key, value, callback=callback)

    def srem(self, key, value, callback=None):
        self.execute_command('SREM', key, value, callback=callback)

    def scard(self, key, callback=None):
        self.execute_command('SCARD', key, callback=callback)

    def spop(self, key, callback=None):
        self.execute_command('SPOP', key, callback=callback)

    def smove(self, src, dst, value, callback=None):
        self.execute_command('SMOVE', src, dst, value, callback=callback)

    def sismember(self, key, value, callback=None):
        self.execute_command('SISMEMBER', key, value, callback=callback)

    def smembers(self, key, callback=None):
        self.execute_command('SMEMBERS', key, callback=callback)

    def srandmember(self, key, callback=None):
        self.execute_command('SRANDMEMBER', key, callback=callback)

    def sinter(self, keys, callback=None):
        self.execute_command('SINTER', *keys, callback=callback)

    def sdiff(self, keys, callback=None):
        self.execute_command('SDIFF', *keys, callback=callback)

    def sunion(self, keys, callback=None):
        self.execute_command('SUNION', *keys, callback=callback)

    def sinterstore(self, keys, dst, callback=None):
        self.execute_command('SINTERSTORE', dst, *keys, callback=callback)

    def sunionstore(self, keys, dst, callback=None):
        self.execute_command('SUNIONSTORE', dst, *keys, callback=callback)

    def sdiffstore(self, keys, dst, callback=None):
        self.execute_command('SDIFFSTORE', dst, *keys, callback=callback)

    ### SORTED SET COMMANDS
    def zadd(self, key, score, value, callback=None):
        self.execute_command('ZADD', key, score, value, callback=callback)

    def zcard(self, key, callback=None):
        self.execute_command('ZCARD', key, callback=callback)

    def zincrby(self, key, value, amount, callback=None):
        self.execute_command('ZINCRBY', key, amount, value, callback=callback)

    def zrank(self, key, value, callback=None):
        self.execute_command('ZRANK', key, value, callback=callback)

    def zrevrank(self, key, value, callback=None):
        self.execute_command('ZREVRANK', key, value, callback=callback)

    def zrem(self, key, value, callback=None):
        self.execute_command('ZREM', key, value, callback=callback)

    def zcount(self, key, start, end, offset=None, limit=None,
               with_scores=None, callback=None):
        tokens = [key, start, end]
        if offset is not None:
            tokens.append('LIMIT')
            tokens.append(offset)
            tokens.append(limit)
        if with_scores:
            tokens.append('WITHSCORES')
        self.execute_command('ZCOUNT', *tokens, callback=callback)

    def zscore(self, key, value, callback=None):
        self.execute_command('ZSCORE', key, value, callback=callback)

    def zrange(self, key, start, num, with_scores, callback=None):
        tokens = [key, start, num]
        if with_scores:
            tokens.append('WITHSCORES')
        self.execute_command('ZRANGE', *tokens, callback=callback)

    def zrevrange(self, key, start, num, with_scores, callback=None):
        tokens = [key, start, num]
        if with_scores:
            tokens.append('WITHSCORES')
        self.execute_command('ZREVRANGE', *tokens, callback=callback)

    def zrangebyscore(self, key, start, end, offset=None, limit=None,
                      with_scores=False, callback=None):
        tokens = [key, start, end]
        if offset is not None:
            tokens.append('LIMIT')
            tokens.append(offset)
            tokens.append(limit)
        if with_scores:
            tokens.append('WITHSCORES')
        self.execute_command('ZRANGEBYSCORE', *tokens, callback=callback)

    def zrevrangebyscore(self, key, end, start, offset=None, limit=None,
                         with_scores=False, callback=None):
        tokens = [key, end, start]
        if offset is not None:
            tokens.append('LIMIT')
            tokens.append(offset)
            tokens.append(limit)
        if with_scores:
            tokens.append('WITHSCORES')
        self.execute_command('ZREVRANGEBYSCORE', *tokens, callback=callback)

    def zremrangebyrank(self, key, start, end, callback=None):
        self.execute_command('ZREMRANGEBYRANK', key, start, end,
                             callback=callback)

    def zremrangebyscore(self, key, start, end, callback=None):
        self.execute_command('ZREMRANGEBYSCORE', key, start, end,
                             callback=callback)

    def zinterstore(self, dest, keys, aggregate=None, callback=None):
        return self._zaggregate('ZINTERSTORE', dest, keys, aggregate, callback)

    def zunionstore(self, dest, keys, aggregate=None, callback=None):
        return self._zaggregate('ZUNIONSTORE', dest, keys, aggregate, callback)

    def _zaggregate(self, command, dest, keys, aggregate, callback):
        tokens = [dest, len(keys)]
        if isinstance(keys, dict):
            items = keys.items()
            keys = [i[0] for i in items]
            weights = [i[1] for i in items]
        else:
            weights = None
        tokens.extend(keys)
        if weights:
            tokens.append('WEIGHTS')
            tokens.extend(weights)
        if aggregate:
            tokens.append('AGGREGATE')
            tokens.append(aggregate)
        return self.execute_command(command, *tokens, callback=callback)

    ### HASH COMMANDS
    def hgetall(self, key, callback=None):
        self.execute_command('HGETALL', key, callback=callback)

    def hmset(self, key, mapping, callback=None):
        items = []
        map(items.extend, mapping.iteritems())
        self.execute_command('HMSET', key, *items, callback=callback)

    def hset(self, key, field, value, callback=None):
        self.execute_command('HSET', key, field, value, callback=callback)

    def hget(self, key, field, callback=None):
        self.execute_command('HGET', key, field, callback=callback)

    def hdel(self, key, field, callback=None):
        self.execute_command('HDEL', key, field, callback=callback)

    def hlen(self, key, callback=None):
        self.execute_command('HLEN', key, callback=callback)

    def hexists(self, key, field, callback=None):
        self.execute_command('HEXISTS', key, field, callback=callback)

    def hincrby(self, key, field, amount=1, callback=None):
        self.execute_command('HINCRBY', key, field, amount, callback=callback)

    def hkeys(self, key, callback=None):
        self.execute_command('HKEYS', key, callback=callback)

    def hmget(self, key, fields, callback=None):
        self.execute_command('HMGET', key, *fields, callback=callback)

    def hvals(self, key, callback=None):
        self.execute_command('HVALS', key, callback=callback)

    ### PUBSUB
    def subscribe(self, channels, callback=None):
        self._subscribe('SUBSCRIBE', channels, callback=callback)

    def psubscribe(self, channels, callback=None):
        self._subscribe('PSUBSCRIBE', channels, callback=callback)

    def _subscribe(self, cmd, channels, callback=None):
        if isinstance(channels, basestring):
            channels = [channels]
        if not self.subscribed:
            original_callback = callback

            def _cb(*args, **kwargs):
                self.on_subscribed(*args, **kwargs)
                original_callback(*args, **kwargs)

            callback = _cb if original_callback else self.on_subscribed
        self.execute_command(cmd, *channels, callback=callback)

    def on_subscribed(self, result):
        self.subscribed = True

    def unsubscribe(self, channels, callback=None):
        self._unsubscribe('UNSUBSCRIBE', channels, callback=callback)

    def punsubscribe(self, channels, callback=None):
        self._unsubscribe('PUNSUBSCRIBE', channels, callback=callback)

    def _unsubscribe(self, cmd, channels, callback=None):
        if isinstance(channels, basestring):
            channels = [channels]
        self.execute_command(cmd, *channels, callback=callback)

    def on_unsubscribed(self, *args, **kwargs):
        self.subscribed = False

    def publish(self, channel, message, callback=None):
        self.execute_command('PUBLISH', channel, message, callback=callback)

    @gen.engine
    def listen(self, callback=None):
        if callback:
            def error_wrapper(e):
                if isinstance(e, GeneratorExit):
                    return ConnectionError('Connection lost')
                else:
                    return e

            cmd_listen = CmdLine('LISTEN')
            with self.connection as c:
                while self.subscribed:
                    data = yield gen.Task(c.readline)
                    if isinstance(data, Exception):
                        raise data

                    response = yield gen.Task(self.process_data,
                                              data,
                                              cmd_listen)
                    if isinstance(response, Exception):
                        raise response

                    result = self.format_reply(cmd_listen, response)

                    callback(result)
                    if result.kind in ['unsubscribe', 'punsubscribe'] \
                    and result.body == 0:
                        self.on_unsubscribed()

    ### CAS
    def watch(self, key, callback=None):
        self.execute_command('WATCH', key, callback=callback)

    def unwatch(self, callback=None):
        self.execute_command('UNWATCH', callback=callback)

    ### SCRIPTING COMMANDS
    def eval(self, script, keys=None, args=None, callback=None):
        if keys is None:
            keys = []
        if args is None:
            args = []
        num_keys = len(keys)
        keys.extend(args)
        self.execute_command('EVAL', script, num_keys, *keys, callback=callback)

    def evalsha(self, shahash, keys=None, args=None, callback=None):
        if keys is None:
            keys = []
        if args is None:
            args = []
        num_keys = len(keys)
        keys.extend(args)
        self.execute_command('EVALSHA', shahash, num_keys, *keys, callback=callback)

    def script_exists(self, shahashes, callback=None):
        # not yet implemented in the redis protocol
        self.execute_command('SCRIPT EXISTS', *shahashes, callback=callback)

    def script_flush(self, callback=None):
        # not yet implemented in the redis protocol
        self.execute_command('SCRIPT FLUSH', callback=callback, verbose=True)

    def script_kill(self, callback=None):
        # not yet implemented in the redis protocol
        self.execute_command('SCRIPT KILL', callback=callback)

    def script_load(self, script, callback=None):
        # not yet implemented in the redis protocol
        self.execute_command('SCRIPT LOAD', script, callback=callback)


class Pipeline(Client):

    def __init__(self, transactional, *args, **kwargs):
        super(Pipeline, self).__init__(*args, **kwargs)
        self.transactional = transactional
        self.command_stack = []

    def execute_command(self, cmd, *args, **kwargs):
        if cmd in ('AUTH', 'SELECT'):
            super(Pipeline, self).execute_command(cmd, *args, **kwargs)
        elif cmd in PUB_SUB_COMMANDS:
            raise RequestError(
                'Client is not supposed to issue command %s in pipeline' % cmd)
        self.command_stack.append(CmdLine(cmd, *args, **kwargs))

    def discard(self):
        # actually do nothing with redis-server, just flush the command_stack
        self.command_stack = []

    def format_replies(self, cmd_lines, responses):
        results = []
        for cmd_line, response in zip(cmd_lines, responses):
            try:
                results.append(self.format_reply(cmd_line, response))
            except Exception, e:
                results.append(e)
        return results

    def format_pipeline_request(self, command_stack):
        return ''.join(self.format_command(c.cmd, *c.args, **c.kwargs)
                       for c in command_stack)

    @gen.engine
    def execute(self, callback=None):
        command_stack = self.command_stack
        self.command_stack = []

        if self.transactional:
            command_stack = [CmdLine('MULTI')] \
                          + command_stack \
                          + [CmdLine('EXEC')]

        request = self.format_pipeline_request(command_stack)

        if not self.connection.ready():
            yield gen.Task(self.connection.wait_until_ready)

        try:
            self.connection.write(request)
        except IOError:
            self.command_stack = []
            self.connection.disconnect()
            raise ConnectionError("Socket closed on remote end")
        except Exception, e:
            self.command_stack = []
            self.connection.disconnect()
            raise e

        responses = []
        total = len(command_stack)
        cmds = iter(command_stack)

        with self.connection:
            while len(responses) < total:
                data = yield gen.Task(self.connection.readline)
                if not data:
                    raise ResponseError('Not enough data after EXEC')
                try:
                    cmd_line = cmds.next()
                    if self.transactional and cmd_line.cmd != 'EXEC':
                        response = yield gen.Task(self.process_data,
                                                  data,
                                                  CmdLine('MULTI_PART'))
                    else:
                        response = yield gen.Task(self.process_data,
                                                  data,
                                                  cmd_line)
                    responses.append(response)
                except Exception, e:
                    responses.append(e)

        if self.transactional:
            command_stack = command_stack[:-1]
            responses = responses[-1]
            results = self.format_replies(command_stack[1:], responses)
        else:
            results = self.format_replies(command_stack, responses)

        callback(results)
