// Copyright (C) 2019 Marius Schellenberger

package store

import (
	"encoding/json"
	"github.com/boltdb/bolt"
	"io"
)

const (
	defaultBoltBucket = "default"
	fileMode          = 0640
)

type BoltStore struct {
	Marshaler
	db *bolt.DB
}

func NewBoltStore(file string, m Marshaler) (Store, error) {
	if m == nil {
		m = NewGOB()
	}
	db, err := bolt.Open(file, fileMode, nil)
	if err != nil {
		return nil, err
	}
	if err = db.Update(func(tx *bolt.Tx) (err error) {
		_, err = tx.CreateBucketIfNotExists([]byte(defaultBoltBucket))
		return
	}); err != nil {
		return nil, err
	}
	return &BoltStore{m, db}, nil
}

type dump struct {
	K string `json:"k"`
	V []byte `json:"v"`
}

func (bs *BoltStore) Dump(w io.Writer) (err error) {
	if w == nil {
		return ErrWriterIsNil
	}
	var d []dump
	err = bs.ForEach(func(k string, v []byte) error {
		b := make([]byte, len(v))
		copy(b, v)
		d = append(d, dump{k, b})
		return nil
	})
	if err != nil {
		return
	}
	return json.NewEncoder(w).Encode(d)
}

func (bs *BoltStore) Restore(r io.Reader) (err error) {
	if r == nil {
		return ErrReaderIsNil
	}
	var d []dump
	err = json.NewDecoder(r).Decode(&d)
	if err != nil {
		return
	}
	if err = bs.db.Update(func(tx *bolt.Tx) (err error) {
		if err = tx.DeleteBucket([]byte(defaultBoltBucket)); err != nil {
			return
		}
		_, err = tx.CreateBucket([]byte(defaultBoltBucket))
		return
	}); err != nil {
		return
	}
	return bs.db.Update(func(tx *bolt.Tx) (err error) {
		b := tx.Bucket([]byte(defaultBoltBucket))
		for _, v := range d {
			b.Put([]byte(v.K), v.V)
		}
		return nil
	})
}

func (bs *BoltStore) Get(key string, v interface{}) (err error) {
	err = bs.db.View(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(defaultBoltBucket)).Get([]byte(key))
		if b == nil {
			return ErrKeyNotFound
		}
		if v != nil {
			return bs.Unmarshal(b, v)
		}
		return nil
	})
	return
}

func (bs *BoltStore) Set(key string, v interface{}) error {
	return bs.db.Update(func(tx *bolt.Tx) (err error) {
		var b []byte
		if v != nil {
			b, err = bs.Marshal(v)
			if err != nil {
				return
			}
		}
		return tx.Bucket([]byte(defaultBoltBucket)).Put([]byte(key), b)
	})
}

func (bs *BoltStore) ForEach(f func(string, []byte) error) error {
	return bs.db.View(func(tx *bolt.Tx) error {
		return tx.Bucket([]byte(defaultBoltBucket)).ForEach(func(k, v []byte) error {
			return f(string(k), v)
		})
	})
}

func (bs *BoltStore) ForEachPrefix(prefix string, f func(string, []byte) error) error {
	return bs.ForEach(func(k string, v []byte) error {
		if trim, ok := hasTrimPrefix(k, prefix); ok {
			return f(trim, v)
		}
		return nil
	})
}

func (bs *BoltStore) Delete(key string) error {
	return bs.db.Update(func(tx *bolt.Tx) error {
		return tx.Bucket([]byte(defaultBoltBucket)).Delete([]byte(key))
	})
}

func (bs *BoltStore) Close() error {
	return bs.db.Close()
}
