#!/usr/bin/env python

# cacheclient.py - example cache client implementation

#

# Copyright 2013 Facebook, Inc.

#

# This software may be used and distributed according to the terms of the

# GNU General Public License version 2 or any later version.

# The remotefilelog extension can optionally use a caching layer to serve

# file revision requests. This is an example implementation that uses

# the python-memcached library: https://pypi.python.org/pypi/python-memcached/

# A better implementation would make all of the requests non-blocking.

from __future__ import absolute_import

import os

import sys

import memcache

stdin = sys.stdin

stdout = sys.stdout

stderr = sys.stderr

mc = None

keyprefix = None

cachepath = None

# Max number of keys per request

batchsize = 1000

# Max value size per key (in bytes)

maxsize = 512 * 1024

def readfile(path):

    f = open(path, "r")

    try:

        return f.read()

    finally:

        f.close()

def writefile(path, content):

    dirname = os.path.dirname(path)

    if not os.path.exists(dirname):

        os.makedirs(dirname)

    f = open(path, "w")

    try:

        f.write(content)

    finally:

        f.close()

def compress(value):

    # Real world implementations will want to compress values.

    # Insert your favorite compression here, ex:

    # return lz4wrapper.lzcompresshc(value)

    return value

def decompress(value):

    # Real world implementations will want to compress values.

    # Insert your favorite compression here, ex:

    # return lz4wrapper.lz4decompress(value)

    return value

def generateKey(id):

    return keyprefix + id

def generateId(key):

    return key[len(keyprefix):]

def getKeys():

    raw = stdin.readline()[:-1]

    keycount = int(raw)

    keys = []

    for i in range(keycount):

        id = stdin.readline()[:-1]

        keys.append(generateKey(id))

    results = mc.get_multi(keys)

    hits = 0

    for i, key in enumerate(keys):

        value = results.get(key)

        id = generateId(key)

        # On hit, write to disk

        if value:

            # Integer hit indicates a large file

            if isinstance(value, int):

                largekeys = list([key + str(i) for i in range(value)])

                largevalues = mc.get_multi(largekeys)

                if len(largevalues) == value:

                    value = ""

                    for largekey in largekeys:

                        value += largevalues[largekey]

                else:

                    # A chunk is missing, give up

                    stdout.write(id + "\n")

                    stdout.flush()

                    continue

            path = os.path.join(cachepath, id)

            value = decompress(value)

            writefile(path, value)

            hits += 1

        else:

            # On miss, report to caller

            stdout.write(id + "\n")

            stdout.flush()

        if i % 500 == 0:

            stdout.write("_hits_%s_\n" % hits)

            stdout.flush()

    # done signal

    stdout.write("0\n")

    stdout.flush()

def setKeys():

    raw = stdin.readline()[:-1]

    keycount = int(raw)

    values = {}

    for i in range(keycount):

        id = stdin.readline()[:-1]

        path = os.path.join(cachepath, id)

        value = readfile(path)

        value = compress(value)

        key = generateKey(id)

        if len(value) > maxsize:

            # split up large files

            start = 0

            i = 0

            while start < len(value):

                end = min(len(value), start + maxsize)

                values[key + str(i)] = value[start:end]

                start += maxsize

                i += 1

            # Large files are stored as an integer representing how many

            # chunks it's broken into.

            value = i

        values[key] = value

        if len(values) == batchsize:

            mc.set_multi(values)

            values = {}

    if values:

        mc.set_multi(values)

def main(argv=None):

    """

    remotefilelog uses this cacheclient by setting it in the repo config:

    [remotefilelog]

    cacheprocess = cacheclient <ip address:port> <memcache prefix>

    When memcache requests need to be made, it will execute this process

    with the following arguments:

    cacheclient <ip address:port> <memcache prefix><internal prefix> <cachepath>

    Communication happens via stdin and stdout. To make a get request,

    the following is written to stdin:

    get\n

    <key count>\n

    <key1>\n

    <key...>\n

    <keyN>\n

    The results of any cache hits will be written directly to <cachepath>/<key>.

    Any cache misses will be written to stdout in the form <key>\n. Once all

    hits and misses are finished 0\n will be written to stdout to signal

    completion.

    During the request, progress may be reported via stdout with the format

    _hits_###_\n where ### is an integer representing the number of hits so

    far. remotefilelog uses this to display a progress bar.

    A single cacheclient process may be used for multiple requests (though

    not in parallel), so it stays open until it receives exit\n via stdin.

    """

    if argv is None:

        argv = sys.argv

    global cachepath

    global keyprefix

    global mc

    ip = argv[1]

    keyprefix = argv[2]

    cachepath = argv[3]

    mc = memcache.Client([ip], debug=0)

    while True:

        cmd = stdin.readline()[:-1]

        if cmd == "get":

            getKeys()

        elif cmd == "set":

            setKeys()

        elif cmd == "exit":

            return 0

        else:

            stderr.write("Invalid Command %s\n" % cmd)

            return 1

if __name__ == "__main__":

    sys.exit(main())