package storage

import (
	"net"
	"path"

	"github.com/bradfitz/gomemcache/memcache"
	"github.com/buraksezer/consistent"
	"github.com/cespare/xxhash"
)

type MemcacheCluster struct {
	db *memcache.Client
}

type serverSelector struct {
	hash *consistent.Consistent
}

func (ss *serverSelector) PickServer(key string) (net.Addr, error) {
	return &netAddr{
		network: "tcp",
		addr:    ss.hash.LocateKey([]byte(key)).String(),
	}, nil
}

func (ss *serverSelector) Each(each func(net.Addr) error) error {
	for _, member := range ss.hash.GetMembers() {
		if err := each(&netAddr{
			network: "tcp",
			addr:    member.String(),
		}); err != nil {
			return err
		}
	}
	return nil
}

type hasher struct{}

func (h hasher) Sum64(data []byte) uint64 {
	return xxhash.Sum64(data)
}

func NewMemcacheCluster(addr string, addrs ...string) (*MemcacheCluster, error) {
	cfg := consistent.Config{
		PartitionCount:    71,
		ReplicationFactor: 20,
		Load:              1.25,
		Hasher:            hasher{},
	}
	hash := consistent.New(nil, cfg)
	for _, addr := range append(addrs, addr) {
		hash.Add(&netAddr{addr: addr})
	}
	ss := &serverSelector{
		hash: hash,
	}
	if err := ss.Each(func(addr net.Addr) error {
		conn, err := net.Dial("tcp", addr.String())
		if err != nil {
			return err
		}
		return conn.Close()
	}); err != nil {
		return nil, err
	}
	db := memcache.NewFromSelector(ss)
	return &MemcacheCluster{db: db}, nil
}

func (mc *MemcacheCluster) Get(key string, ns ...string) ([]byte, error) {
	namespace := resolveNamespace(ns)
	v, err := mc.db.Get(path.Join(namespace, key))
	return v.Value, err
}

func (mc *MemcacheCluster) Set(key string, value []byte, ns ...string) error {
	namespace := resolveNamespace(ns)
	return mc.db.Set(&memcache.Item{
		Key:   path.Join(namespace, key),
		Value: value,
	})
}

func (mc *MemcacheCluster) Close() error {
	return mc.db.FlushAll()
}