import os
import multiprocessing
import logging
import marshal
import cPickle
import shutil
import struct 
import urllib

import msgpack

from dpark.env import env
from dpark.tracker import GetValueMessage, AddItemMessage, RemoveItemMessage

logger = logging.getLogger("cache")

class Cache:
    data = {}

    def get(self, key): 
        return self.data.get(key)
    
    def put(self, key, value, is_iterator=False):
        if value is not None:
            if is_iterator:
                value = list(value)
            self.data[key] = value
            return value
        else:
            self.data.pop(key, None)

    def clear(self):
        self.data.clear()

class DiskCache(Cache):
    def __init__(self, tracker, path):
        if not os.path.exists(path):
            try: os.makedirs(path)
            except: pass
        self.tracker = tracker
        self.root = path

    def get_path(self, key):
        return os.path.join(self.root, '%s_%s' % key)

    def get(self, key):
        p = self.get_path(key)
        if os.path.exists(p):
            return self.load(open(p, 'rb'))

        # load from other node
        if not env.get('SERVER_URI'):
            return
        rdd_id, index = key
        locs = self.tracker.getCacheUri(rdd_id, index)
        if not locs:
            return

        serve_uri = locs[-1]
        uri = '%s/cache/%s' % (serve_uri, os.path.basename(p))
        f = urllib.urlopen(uri)