package main

import (
	"bytes"
	"compress/flate"
	"encoding/gob"
	fmt "fmt"
	"log"
	"math/rand"
	"os"
	"reflect"

	"github.com/tecbot/gorocksdb"
)

type KeyKind struct {
	Key       string
	FieldName string
}

func NewKeyKind(key string, fieldname string) KeyKind {
	return KeyKind{Key: key, FieldName: fieldname}
}

type KeyList struct {
	Keys []interface{}
}

func (kl *KeyList) GetKeys() []interface{} {
	return kl.Keys
}

func (kl *KeyList) AppendKey(key interface{}) {
	kl.Keys = append(kl.Keys, key)
}

func (kl *KeyList) GetLength() int {
	return len(kl.Keys)
}

// CheckErrorPanic 查错误 存在就panic操作
func CheckErrorPanic(err error) {
	if err != nil {
		log.Panic(err)
	}
}

// OpenDB 打开rocksdb
func OpenDB() *gorocksdb.DB {
	bbto := gorocksdb.NewDefaultBlockBasedTableOptions()

	bbto.SetBlockCache(gorocksdb.NewLRUCache(2 << 30))
	bbto.SetCacheIndexAndFilterBlocksWithHighPriority(true)

	opts := gorocksdb.NewDefaultOptions()

	bbto.SetFilterPolicy(gorocksdb.NewBloomFilter(16))

	opts.SetBlockBasedTableFactory(bbto)
	opts.SetCreateIfMissing(true)
	opts.SetCreateIfMissingColumnFamilies(true)
	opts.SetCompression(gorocksdb.LZ4Compression)

	db, err := gorocksdb.OpenDb(opts, ".rocksdb")
	CheckErrorPanic(err)
	return db
}

var db = OpenDB()
var wopts = gorocksdb.NewDefaultWriteOptions()
var ropts = gorocksdb.NewDefaultReadOptions()

// IKeyList key list interface
type IKeyList interface {
	AppendKey(key interface{})
	GetKeys() []interface{}
	GetLength() int
}

// SaveData gob序列化存数据
func SaveData(fname string, v interface{}) {
	var keybuf = &bytes.Buffer{}

	enc := gob.NewEncoder(keybuf)
	enc.Encode(v)

	f, err := os.OpenFile(fname, os.O_CREATE|os.O_RDWR|os.O_TRUNC, os.ModePerm)
	CheckErrorPanic(err)

	gw, err := flate.NewWriter(f, 7)
	CheckErrorPanic(err)
	defer gw.Close()

	gw.Write(keybuf.Bytes())
}

// SaveGob 用于解析rocksdb源数据, 爬取数据存储在rocksdb后取出 序列化存储
func SaveGob(fname string, datatype reflect.Type, kfs ...KeyKind) {

	// cl := &CountryList{}
	kl := &KeyList{}

	i := 0
	iter := db.NewIterator(ropts)
	bkey := []byte(kfs[0].Key)
	iter.Seek(bkey)

	for ; iter.ValidForPrefix(bkey); iter.Next() {
		i++
		first := iter.Value().Data()
		data := reflect.New(datatype).Elem()

		data.FieldByName(kfs[0].FieldName).Set(reflect.ValueOf(first))

		// country := &Country{}
		// country.Name = name
		for n := 1; n < len(kfs); n++ {
			v, err := db.Get(ropts, []byte(fmt.Sprintf(kfs[n].Key, first)))
			if err != nil {
				panic(err)
			}
			fv := data.FieldByName(kfs[n].FieldName)
			fv.Set(reflect.ValueOf(v.Data()))
		}
		kl.AppendKey(data.Interface())
	}
	SaveData(fname, kl)
}

// LoadGob load gob from file
func LoadGob(fname string, v IKeyList) {

	f, err := os.Open(fname)
	log.Println("load data: ", fname)
	CheckErrorPanic(err)
	reader := flate.NewReader(f)
	dec := gob.NewDecoder(reader)
	CheckErrorPanic(dec.Decode(v))

}

// GetRandomKey get ikeylist key by random
func GetRandomKey(v IKeyList) interface{} {
	ret := v.GetKeys()[rand.Intn(v.GetLength())]
	return ret
}

// GetRandomKeyBySlice get ikeylist key by random
func GetRandomKeyBySlice(v interface{}) interface{} {
	sli := reflect.ValueOf(v)
	return sli.Index(rand.Intn(sli.Len())).Interface()
}