#!/usr/bin/env python3
"""Bounded real Redis experiments using Python's standard library (RESP2).

No CONFIG SET, restarts, scans, FLUSH, persistence jobs or existing-key writes.
"""
import argparse
import json
import socket
import time
import uuid


class RedisError(Exception):
    pass


class Client:
    def __init__(self, host, port):
        self.sock = socket.create_connection((host, port), timeout=3)
        self.file = self.sock.makefile('rb')

    def read(self):
        line = self.file.readline()
        if not line.endswith(b'\r\n'):
            raise ConnectionError('incomplete RESP response')
        kind, value = line[:1], line[1:-2]
        if kind == b'+':
            return value.decode()
        if kind == b'-':
            return RedisError(value.decode())
        if kind == b':':
            return int(value)
        if kind == b'$':
            size = int(value)
            if size == -1:
                return None
            data = self.file.read(size)
            if len(data) != size or self.file.read(2) != b'\r\n':
                raise ConnectionError('incomplete RESP bulk string')
            return data.decode()
        if kind == b'*':
            size = int(value)
            return None if size == -1 else [self.read() for _ in range(size)]
        raise ValueError('unsupported RESP2 response')

    def call(self, *args):
        parts = [str(arg).encode() for arg in args]
        self.sock.sendall(b'*' + str(len(parts)).encode() + b'\r\n' + b''.join(
            b'$' + str(len(part)).encode() + b'\r\n' + part + b'\r\n' for part in parts))
        result = self.read()
        if isinstance(result, RedisError):
            raise result
        return result

    def close(self):
        self.file.close()
        self.sock.close()


def info(client, section):
    return dict(line.split(':', 1) for line in client.call('INFO', section).splitlines()
                if line and not line.startswith('#'))


class Experiment:
    def __init__(self, client):
        self.client = client
        self.prefix = 'study_t17_' + uuid.uuid4().hex + ':'
        self.created = []

    def key(self, suffix, empty=False):
        key = self.prefix + suffix
        # Collision fails before touching the key; record only successful creation.
        if self.client.call('SET', key, 'reserved', 'NX', 'EX', 120) != 'OK':
            raise RuntimeError('unique namespace collision; refusing existing key')
        self.created.append(key)
        if empty:
            self.client.call('DEL', key)  # only our own reservation
        return key

    def cleanup(self):
        # An unexpected queueing failure may leave this connection in MULTI.
        # Discard only its queue before issuing real cleanup commands.
        try:
            self.client.call('DISCARD')
        except RedisError:
            pass  # normal path: no active transaction
        # No prefix scan: delete exactly the keys this run actually reserved.
        removed = self.client.call('DEL', *self.created) if self.created else 0
        remaining = sum(self.client.call('EXISTS', key) for key in self.created)
        if remaining:
            raise RuntimeError('experiment cleanup incomplete')
        return {'created_keys': len(self.created), 'deleted_existing_keys': removed, 'remaining_keys': remaining}


def encoding(exp, config):
    c = exp.client
    limits = {name: int(config[name]) for name in (
        'hash-max-listpack-entries', 'hash-max-listpack-value', 'set-max-intset-entries',
        'set-max-listpack-entries', 'set-max-listpack-value', 'zset-max-listpack-entries',
        'zset-max-listpack-value')}
    # Never follow a user config into a large workload. The lessons use defaults.
    caps = {'hash-max-listpack-entries': 512, 'set-max-intset-entries': 512,
            'set-max-listpack-entries': 128, 'zset-max-listpack-entries': 128}
    if any(not 1 <= limits[k] <= cap for k, cap in caps.items()) or any(
            not 1 <= limits[k] <= 64 for k in limits if k.endswith('-value')):
        return {'status': 'skipped', 'reason': 'configured threshold exceeds bounded lab range', 'limits': limits}
    h = exp.key('hash-count', empty=True)
    n = limits['hash-max-listpack-entries']
    c.call('HSET', h, *[part for i in range(n) for part in (f'f{i}', 'x')])
    small_hash = c.call('OBJECT', 'ENCODING', h)
    small_bytes = c.call('MEMORY', 'USAGE', h)
    c.call('HSET', h, 'overflow', 'x')
    large_hash = c.call('OBJECT', 'ENCODING', h)
    c.call('EXPIRE', h, 120)
    hlong = exp.key('hash-long', empty=True)
    c.call('HSET', hlong, 'f', 'x' * (limits['hash-max-listpack-value'] + 1))
    long_hash = c.call('OBJECT', 'ENCODING', hlong)
    hdup = exp.key('hash-duplicate', empty=True)
    c.call('HSET', hdup, *['f', 'x'] * (n + 1))
    duplicate = {'fields': c.call('HLEN', hdup), 'encoding': c.call('OBJECT', 'ENCODING', hdup)}
    s = exp.key('set-mixed', empty=True)
    c.call('SADD', s, 1, 2)
    integer_set = c.call('OBJECT', 'ENCODING', s)
    c.call('SADD', s, 'a')
    mixed_set = c.call('OBJECT', 'ENCODING', s)
    c.call('SADD', s, 'x' * (limits['set-max-listpack-value'] + 1))
    large_set = c.call('OBJECT', 'ENCODING', s)
    si = exp.key('set-integer-count', empty=True)
    c.call('SADD', si, *range(limits['set-max-intset-entries'] + 1))
    many_integers = c.call('OBJECT', 'ENCODING', si)
    z = exp.key('zset-count', empty=True)
    zn = limits['zset-max-listpack-entries']
    c.call('ZADD', z, *[part for i in range(zn) for part in (i, f'm{i}')])
    small_zset = c.call('OBJECT', 'ENCODING', z)
    c.call('ZADD', z, zn, 'overflow')
    large_zset = c.call('OBJECT', 'ENCODING', z)
    for key in (hlong, hdup, s, si, z):
        c.call('EXPIRE', key, 120)
    assert [small_hash, large_hash, long_hash] == ['listpack', 'hashtable', 'hashtable']
    assert [integer_set, mixed_set, large_set] == ['intset', 'listpack', 'hashtable']
    assert [small_zset, large_zset] == ['listpack', 'skiplist']
    return {'status': 'passed', 'limits': limits, 'hash': [small_hash, large_hash, long_hash],
            'hash_small_memory_bytes': small_bytes, 'hash_duplicate_arguments': duplicate,
            'set': [integer_set, mixed_set, large_set], 'set_many_integers': many_integers,
            'zset': [small_zset, large_zset]}


def atomic(exp):
    c = exp.client
    script_key = exp.key('script')
    try:
        c.call('EVAL', "redis.call('SET',KEYS[1],'before'); return redis.call('INCR',KEYS[1])", 1, script_key)
        raise AssertionError('script must fail on non-integer value')
    except RedisError as error:
        script_error = str(error)
    script_value = c.call('GET', script_key)
    tx_key, after_key = exp.key('transaction'), exp.key('after')
    c.call('MULTI')
    c.call('SET', tx_key, 'before')
    c.call('LPUSH', tx_key, 'item')
    c.call('SET', after_key, 'after')
    replies = c.call('EXEC')
    assert replies[0] == 'OK' and isinstance(replies[1], RedisError) and replies[2] == 'OK'
    values = [c.call('GET', tx_key), c.call('GET', after_key)]
    assert script_value == 'before' and values == ['before', 'after']
    for key in (script_key, tx_key, after_key):
        c.call('EXPIRE', key, 120)
    return {'status': 'passed', 'script_error': script_error, 'script_value_after_error': script_value,
            'exec_replies': [str(reply) for reply in replies], 'transaction_values': values}


def ttl(exp):
    c = exp.client
    key = exp.key('ttl')
    c.call('SET', key, 'v', 'PX', 200)
    initial = c.call('PTTL', key)
    assert 0 <= initial <= 200
    deadline = time.monotonic() + 3
    while c.call('GET', key) is not None:
        if time.monotonic() >= deadline:
            raise AssertionError('TTL key still readable after bounded wait')
        time.sleep(0.02)
    expired = c.call('PTTL', key)
    assert expired == -2
    overwrite = exp.key('overwrite-ttl')
    c.call('SET', overwrite, 'old', 'PX', 60000)
    c.call('SET', overwrite, 'new', 'KEEPTTL')
    kept = c.call('PTTL', overwrite)
    c.call('SET', overwrite, 'replace')
    cleared = c.call('PTTL', overwrite)
    assert 0 < kept <= 60000 and cleared == -1
    c.call('EXPIRE', overwrite, 120)
    return {'status': 'passed', 'initial_pttl_ms': initial, 'expired_pttl': expired,
            'keepttl_ms': kept, 'plain_set_pttl': cleared}


def wait(exp, replication):
    c = exp.client
    key = exp.key('wait')
    c.call('SET', key, 'written')
    # Same Client/socket as SET. No pool checkout or second connection here.
    acknowledgements = c.call('WAIT', 1, 100)
    value = c.call('GET', key)
    assert value == 'written'
    if int(replication['connected_slaves']) == 0:
        assert acknowledgements == 0
    c.call('EXPIRE', key, 120)
    return {'status': 'passed', 'requested_replicas': 1, 'acknowledgements': acknowledgements,
            'value_after_wait': value, 'same_connection': True}


def run(args):
    c = Client(args.host, args.port)
    exp = Experiment(c)
    result = {'kind': 'real-redis', 'prefix': exp.prefix}
    try:
        server, memory, replication, persistence = [info(c, section) for section in (
            'server', 'memory', 'replication', 'persistence')]
        raw = c.call('CONFIG', 'GET', 'hash-max-listpack-*', 'set-max-*', 'zset-max-listpack-*',
                     'maxmemory', 'maxmemory-policy', 'appendfsync', 'io-threads', 'io-threads-do-reads')
        config = dict(zip(raw[::2], raw[1::2]))
        result['runtime'] = {name: server[name] for name in ('redis_version', 'redis_mode')}
        result['replication'] = {name: replication[name] for name in ('role', 'connected_slaves')}
        result['aof_enabled'] = persistence['aof_enabled']
        result['memory'] = {name: int(memory[name]) for name in ('used_memory', 'used_memory_rss', 'maxmemory')}
        result['config'] = config
        if args.mode != 'inspect':
            if replication['role'] != 'master' or server['redis_mode'] != 'standalone':
                raise RuntimeError('bounded local lab requires a writable standalone master')
            used, maximum = int(memory['used_memory']), int(memory['maxmemory'])
            if maximum and maximum - used < 16 * 1024 * 1024:
                raise RuntimeError('less than 16 MiB maxmemory headroom; refusing all writes')
            result['memory_guard'] = 'passed: small bounded objects, >=16MiB headroom when capped'
            for name, operation in (('encoding', lambda: encoding(exp, config)), ('atomic', lambda: atomic(exp)),
                                    ('ttl', lambda: ttl(exp)), ('wait', lambda: wait(exp, replication))):
                if args.mode in (name, 'all'):
                    result[name] = operation()
    except Exception as error:
        result['error'] = f'{type(error).__name__}: {error}'
    finally:
        try:
            result['cleanup'] = exp.cleanup()
        except Exception as error:
            result['cleanup_error'] = str(error)
        c.close()
    print(json.dumps(result, ensure_ascii=False, indent=2))
    return 1 if 'error' in result or 'cleanup_error' in result else 0


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--host', default='127.0.0.1')
    parser.add_argument('--port', type=int, default=6379)
    parser.add_argument('--mode', choices=('inspect', 'encoding', 'atomic', 'ttl', 'wait', 'all'), default='inspect')
    raise SystemExit(run(parser.parse_args()))