#!/usr/bin/env python

import datetime, os
from sqlalchemy import *
from sqlalchemy import exceptions, sql
from sqlalchemy.orm import *
from sqlalchemy.orm.shard import ShardedSession
from sqlalchemy.sql import operators

from sqlalchemy import create_engine

from blog_engine import *
from lookup import Lookup
from post import Post

from elixir import *
import md5

BUCKET = 2
DATA = 800
SHARD = 2

lookup_dict = {
    '0':'lookup1',
    '1':'lookup2'
}

shard_lookup_dict={
    '0':'shard1',
    '1':'shard2'
}

def determine_lookup_shard(username):
    hash = md5.new()
    hash.update("%s" % username)
    value = hash.digest()
    rem = long(value.encode("hex"), 16)% BUCKET
    return lookup_dict['%s' % rem]

def shard_chooser_lookup(mapper, instance, clause=None):
    hash = md5.new()
    hash.update("%s" % instance.username)
    value = hash.digest()
    rem = long(value.encode("hex"), 16)% BUCKET
    return lookup_dict['%s' % rem]

def id_chooser_lookup(query, ident):
    hash = md5.new()
    hash.update("%s" % ident)
    value = hash.digest()
    rem = long(value.encode("hex"), 16)% BUCKET
    return lookup_dict['%s' % rem]

def query_chooser_lookup(query):
    ids = []

    class FindLookupShard(sql.ClauseVisitor):
        def visit_binary(self, binary):
            if binary.left is Lookup.username:
                if binary.operator == operators.eq:
                    ids.append(determine_lookup_shard[binary.right.value])
                elif binary.operator == operators.in_op:
                    for bind in binary.right.clauses:
                        ids.append(determine_lookup_shard[bind.value])
                    
    FindLookupShard().traverse(query._criterion)
    if len(ids) == 0:
        return ['lookup1','lookup2'] #['0', '1'] 
    else:
        return ids


create_session_lookup = sessionmaker(class_=ShardedSession, autoflush=True, transactional=True)

create_session_lookup.configure(shards={
            'lookup1':lookup1,
            'lookup2':lookup2
        }, shard_chooser=shard_chooser_lookup, id_chooser=id_chooser_lookup, query_chooser=query_chooser_lookup)


#sesslk = create_session_lookup()