package main

import (
	"context"
	"io/ioutil"
	"path"
	"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 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)
	})
}