package bloom

import (
	"encoding/binary"
	"errors"
	"math"
	"os"
	"sort"
	"sync"

	"github.com/RoaringBitmap/roaring/roaring64"
	"github.com/klauspost/compress/zstd"
	"google.golang.org/protobuf/proto"
)

/*
主体代码来源于https://github.com/bits-and-blooms/bloom
在此文件中，将bitmap的管理改用roaring bitmap实现
algotao 2022-08-29
*/

type BloomFilter struct {
	m                 uint64            // 存贮空间上限
	k                 uint64            // hash函数个数
	elementsMax       uint64            // 元素数量上限
	elementsAdded     uint64            // 已加入的元素数量
	falsePositiveRate float64           // 假阳率
	rb                *roaring64.Bitmap // 位图
	chOne             chan []uint64     // 接收每个插入的hash索引
	chInsert          chan []uint64     // 接收排好序的hash索引进Bitmap
	chSortJobQuota    chan int          // 排序工作的配额控制
	buf               []uint64          // 缓冲
	wgJobs            sync.WaitGroup
}

type BloomFilterStat struct {
	M                 uint64
	K                 uint64
	ElementsMax       uint64
	ElementsAdded     uint64
	FalsePositiveRate float64
}

const (
	headerVersion1 = 1
)

// bitmapFileHeader 存贮文件头
type bitmapFileHeader struct {
	Size uint64 // Header Protobuf size
}

// NewWithEstimates 创建一个BloomFilter，并期望有n个元素，<fp的误匹配率
func NewWithEstimates(e uint64, fr float64) *BloomFilter {
	m, k := EstimateParameters(e, fr)
	return newBloomFilter(m, k, e, fr)
}

// new 创建一个新的BloomFilter，具有 _m_ bits 和 _k_ hashing 函数
func newBloomFilter(m uint64, k uint64, e uint64, fr float64) *BloomFilter {
	b := &BloomFilter{
		m:                 max(1, m),
		k:                 max(1, k),
		elementsMax:       e,
		falsePositiveRate: fr,
		rb:                roaring64.New(),
		chOne:             make(chan []uint64, 1024*1024), // 索引缓冲区chan
		chInsert:          make(chan []uint64, 2),         // 插入队列
		chSortJobQuota:    make(chan int, 8),              // 排序队列
	}

	for i := 0; i < cap(b.chSortJobQuota); i++ {
		b.chSortJobQuota <- 0
	}

	//log.Printf("Init quota len(%v), cap(%v)\n", len(b.chSortJobQuota), cap(b.chSortJobQuota))

	go b.consumeOne()
	go b.consumeInsert()

	return b
}

// 生成 m 和 k
func EstimateParameters(n uint64, p float64) (m uint64, k uint64) {
	m = uint64(math.Ceil(-1 * float64(n) * math.Log(p) / math.Pow(math.Log(2), 2)))
	k = uint64(math.Ceil(math.Log(2) * float64(m) / float64(n)))
	return
}

// location 返回当前位置的hash值
func (b *BloomFilter) location(h [4]uint64, i uint64) uint64 {
	return (h[i%2] + i*h[2+(((i+(i%2))%4)/2)]) % b.m
	//return 0
}

// baseHashes 生成4个hash值，用于生产key
func (b *BloomFilter) baseHashes(data []byte) [4]uint64 {
	h := New128()
	h.Write(data)

	h1, h2 := h.Sum128()

	h.Write([]byte{1})
	h3, h4 := h.Sum128()

	return [4]uint64{
		h1, h2, h3, h4,
	}
}

// 消费一个计算好的bloom bits对象，并填入缓冲。当缓冲半满时发送给排序队列处理
func (b *BloomFilter) consumeOne() {
	batchSize := cap(b.chOne) * int(b.k) / 2 //一半buffer满了就开始处理，即一半个数的uint64。或是遇到Flush标志（bits长度0），则刷缓冲

	for bits := range b.chOne {
		if len(bits) != 0 {
			b.elementsAdded++
		}

		b.buf = append(b.buf, bits...)
		if len(b.buf) >= batchSize || len(bits) == 0 {

			buf := b.buf[:]
			b.buf = []uint64{}
			b.wgJobs.Add(1)

			// 如果接收到了Flush标志，则在处理最后buffer后，减一次waitgroup
			if len(bits) == 0 {
				b.wgJobs.Done()
			}

			//等待有可用排序配额，如成功则消耗一个配额
			<-b.chSortJobQuota

			go func() {
				sort.Slice(buf, func(i, j int) bool { return buf[i] < buf[j] })

				//提交至插入任务
				b.chInsert <- buf

				//恢复1个配额
				b.chSortJobQuota <- 1
			}()

		}
	}
}

// 将批量bits写到bitmap
func (b *BloomFilter) consumeInsert() {
	for bitsBatch := range b.chInsert {
		b.rb.AddMany(bitsBatch)
		b.wgJobs.Done()
	}
}

// Add 添加数据的Hash位图
func (b *BloomFilter) Add(data []byte) *BloomFilter {
	h := b.baseHashes(data)
	bits := make([]uint64, b.k)
	for i := uint64(0); i < b.k; i++ {
		bits[i] = b.location(h, i)
	}
	b.chOne <- bits //将一个计算好的bloom bits发送到待处理队列

	return b
}

// AddString 添加字符串的Hash位图
func (b *BloomFilter) AddString(data string) *BloomFilter {
	return b.Add([]byte(data))
}

// Test 如果命中Hash位图，则返回真 (有误匹配率)
func (b *BloomFilter) Test(data []byte) bool {
	h := b.baseHashes(data)
	for i := uint64(0); i < b.k; i++ {
		if !b.rb.Contains(b.location(h, i)) {
			return false
		}
	}
	return true
}

// TestString 如果命中字符串Hash位图，则返回真 (有误匹配率）
func (b *BloomFilter) TestString(data string) bool {
	return b.Test([]byte(data))
}

// Flush 将缓冲中的待处理Bit写入Bitmap
func (b *BloomFilter) Flush() {
	b.wgJobs.Add(1)

	//发出Flush指令
	b.chOne <- []uint64{}

	b.wgJobs.Wait()
}

// free 将缓冲中的待处理Bit写入Bitmap
func (b *BloomFilter) free() {
	close(b.chOne)
	close(b.chInsert)
	b.rb.Clear()
}

// Iterator 位图遍历器
func (b *BloomFilter) Iterator() roaring64.IntPeekable64 {
	return b.rb.Iterator()
}

// GetSizeInBytes 返回位图大小
func (b *BloomFilter) GetSizeInBytes() uint64 {
	return b.rb.GetSerializedSizeInBytes()
}

// 获得统计信息，主要用于运行期间获取状态
func (b *BloomFilter) GetStat() BloomFilterStat {
	return BloomFilterStat{
		M:                 b.m,
		K:                 b.k,
		ElementsMax:       b.elementsMax,
		ElementsAdded:     b.elementsAdded,
		FalsePositiveRate: b.falsePositiveRate,
	}
}

// SaveToFile 写入到文件
func (b *BloomFilter) SaveToFile(filename string) (err error) {
	b.Flush()

	headerPB := &Header{
		Version:           headerVersion1,
		M:                 b.m,
		K:                 b.k,
		ElementsMax:       b.elementsMax,
		ElementsAdded:     b.elementsAdded,
		FalsePositiveRate: b.falsePositiveRate,
	}
	headerData, err := proto.Marshal(headerPB)
	if err != nil {
		return err
	}

	fi, fe := os.Create(filename)
	if fe != nil {
		return fe
	}

	defer fi.Close()

	fh := bitmapFileHeader{
		Size: uint64(len(headerData)),
	}

	//写入文件头（字节数）
	fe = binary.Write(fi, binary.BigEndian, fh)
	if fe != nil {
		return fe
	}

	//写入文件头（PB详细信息）
	fe = binary.Write(fi, binary.BigEndian, headerData)
	if fe != nil {
		return fe
	}

	b.rb.RunOptimize()

	zw, err := zstd.NewWriter(fi)
	if err != nil {
		return err
	}
	defer zw.Close()

	_, fe = b.rb.WriteTo(zw)

	b.free()

	return fe
}

// LoadFromFile 从文件中读取
func LoadFromFile(filename string, headerOnly bool) (bft *BloomFilter, err error) {
	fi, fe := os.Open(filename)

	if fe != nil {
		return nil, fe
	}

	defer fi.Close()

	fh := bitmapFileHeader{}
	fe = binary.Read(fi, binary.BigEndian, &fh)
	if fe != nil {
		return nil, fe
	}

	headerData := make([]byte, fh.Size)

	n, err := fi.Read(headerData)
	if err != nil {
		return nil, err
	}
	if n != len(headerData) {
		return nil, errors.New("unknown file format")
	}
	headerPB := &Header{}

	err = proto.Unmarshal(headerData, headerPB)
	if err != nil {
		return nil, err
	}

	if headerPB.Version != headerVersion1 {
		return nil, errors.New("unsupported version")
	}

	bft = &BloomFilter{}
	bft.m = headerPB.GetM()
	bft.k = headerPB.GetK()
	bft.elementsMax = headerPB.GetElementsMax()
	bft.elementsAdded = headerPB.GetElementsAdded()
	bft.falsePositiveRate = headerPB.GetFalsePositiveRate()

	if headerOnly {
		return bft, nil
	}

	bft.rb = roaring64.New()

	zr, err := zstd.NewReader(fi)
	if err != nil {
		return nil, err
	}
	defer zr.Close()

	_, fe = bft.rb.ReadFrom(zr)
	if fe != nil {
		return nil, fe
	}

	return bft, nil
}