package main

import (
	"context"
	"io/ioutil"
	"path"
	"sync"
	"time"

	"go.etcd.io/bbolt"
)

type Driver interface {
	Close() error
	ForEach(context.Context, string, func(string, []byte) error) error
	Get(context.Context, string, string) ([]byte, error)
	Set(context.Context, string, string, []byte) error
}

type RAM struct {
	m    map[string]map[string][]byte
	lock *sync.RWMutex
}

func NewRAM() RAM {
	return RAM{
		m:    make(map[string]map[string][]byte),
		lock: &sync.RWMutex{},
	}
}

func (ram RAM) Close() error {
	return nil
}

func (ram RAM) ForEach(ctx context.Context, ns string, cb func(string, []byte) error) error {
	ram.lock.RLock()
	defer ram.lock.RUnlock()

	for k, v := range ram.m[ns] {
		if ctx.Err() != nil {
			break
		}
		if err := cb(k, v); err != nil {
			return err
		}
	}
	return ctx.Err()
}

func (ram RAM) Get(_ context.Context, ns, id string) ([]byte, error) {
	ram.lock.RLock()
	defer ram.lock.RUnlock()

	if _, ok := ram.m[ns]; !ok {
		return nil, nil
	}

	return ram.m[ns][id], nil
}

func (ram RAM) Set(_ context.Context, ns, id string, v []byte) error {
	ram.lock.Lock()
	defer ram.lock.Unlock()

	if _, ok := ram.m[ns]; !ok {
		ram.m[ns] = map[string][]byte{}
	}
	ram.m[ns][id] = v

	return nil
}

type BBolt struct {
	db *bbolt.DB
}

func NewTestDBIn(d string) BBolt {
	d, err := ioutil.TempDir(d, "test-db-*")
	if err != nil {
		panic(err)
	}
	db, err := NewDB(path.Join(d, "bb"))
	if err != nil {
		panic(err)
	}
	return db
}

func NewDB(p string) (BBolt, error) {
	db, err := bbolt.Open(p, 0600, &bbolt.Options{
		Timeout: time.Second,
	})
	return BBolt{db: db}, err
}

func (bb BBolt) Close() error {
	return bb.db.Close()
}

func (bb BBolt) ForEach(ctx context.Context, db string, cb func(string, []byte) error) error {
	return bb.db.View(func(tx *bbolt.Tx) error {
		bkt := tx.Bucket([]byte(db))
		if bkt == nil {
			return nil
		}

		c := bkt.Cursor()
		for k, v := c.First(); k != nil && ctx.Err() == nil; k, v = c.Next() {
			if err := cb(string(k), v); err != nil {
				return err
			}
		}

		return ctx.Err()
	})
}

func (bb BBolt) Get(_ context.Context, db, id string) ([]byte, error) {
	var b []byte
	err := bb.db.View(func(tx *bbolt.Tx) error {
		bkt := tx.Bucket([]byte(db))
		if bkt == nil {
			return nil
		}

		b = bkt.Get([]byte(id))
		return nil
	})
	return b, err
}

func (bb BBolt) Set(_ context.Context, db, id string, value []byte) error {
	return bb.db.Update(func(tx *bbolt.Tx) error {
		bkt := tx.Bucket([]byte(db))
		if bkt == nil {
			var err error
			bkt, err = tx.CreateBucket([]byte(db))
			if err != nil {
				return err
			}
		}

		if value == nil {
			return bkt.Delete([]byte(id))
		}
		return bkt.Put([]byte(id), value)
	})
}