// Copyright 2024 PingCAP, Inc. Licensed under Apache-2.0.

package k

import (
	"encoding/hex"
	"fmt"
	"sort"
	"time"

	"github.com/pingcap-inc/clinic/pkg/ngm/keyviz/model"
)

const heatmapsMaxDisplayY = 1536

type RegionsInfo struct {
	Time    time.Time    `json:"-"`
	Count   int          `json:"count"`
	Regions []RegionInfo `json:"regions"`
}

type RegionInfo struct {
	StartKey     string `json:"start_key"`
	EndKey       string `json:"end_key"`
	WrittenBytes uint64 `json:"written_bytes"`
	ReadBytes    uint64 `json:"read_bytes"`
	WrittenKeys  uint64 `json:"written_keys"`
	ReadKeys     uint64 `json:"read_keys"`
}

type DBTablesInfo struct {
	DB     model.DBInfo    `json:"db"`
	Tables model.TableInfo `json:"tables"`
}

type Matrix struct {
	DataMap  map[string][][]uint64 `json:"data" binding:"required"`
	KeyAxis  []LabelKey            `json:"keyAxis" binding:"required"`
	TimeAxis []int64               `json:"timeAxis" binding:"required"`
}

type LabelKey struct {
	Key    string   `json:"key" binding:"required"`
	Labels []string `json:"labels" binding:"required"`
}

type KeyRange struct {
	StartKey string
	EndKey   string
}

type ValuesAxis struct {
	Sum    uint64
	Counts []uint64
	Values []uint64
}

type KeyRangeValues struct {
	KeyRange
	ValuesAxis
}

type KeyRangeValuesList []KeyRangeValues

func (k KeyRangeValuesList) Len() int {
	return len(k)
}

func (k KeyRangeValuesList) Less(i, j int) bool {
	if k[i].StartKey == "" {
		if k[j].StartKey != "" {
			return true
		}
		return k[i].EndKey < k[j].EndKey
	}
	if k[i].EndKey == "" {
		if k[j].EndKey != "" {
			return false
		}
		return k[i].StartKey < k[j].StartKey
	}
	if k[i].StartKey == k[j].StartKey {
		return k[i].EndKey < k[j].EndKey
	}
	if k[i].EndKey == k[j].EndKey {
		return k[i].StartKey < k[j].StartKey
	}
	return k[i].EndKey < k[j].EndKey
}

func (k KeyRangeValuesList) Swap(i, j int) {
	k[i], k[j] = k[j], k[i]
}

type Merger struct {
	typ    string
	times  []int64
	data   map[KeyRange]*ValuesAxis
	tables map[int64]model.DBTableInfo
}

func NewMerger(typ string, begin, end int64, tables map[int64]model.DBTableInfo) *Merger {
	begin -= begin % 60
	end -= end % 60
	if begin >= end {
		panic("invalid begin end")
	}
	times := calcTimes(begin, end)
	switch typ {
	case "":
		typ = "integration"
	case "integration", "read_keys", "written_keys", "read_bytes":
		// pass
	default:
		typ = "written_bytes"
	}
	return &Merger{
		typ:    typ,
		times:  times,
		data:   make(map[KeyRange]*ValuesAxis),
		tables: tables,
	}
}

func (m *Merger) Merge(ts int64, regions []RegionInfo) {
	idx := findIndexFromTimes(m.times, ts)
	if idx == -1 {
		return
	}
	for _, region := range regions {
		kr := KeyRange{
			StartKey: region.StartKey,
			EndKey:   region.EndKey,
		}
		if _, ok := m.data[kr]; !ok {
			m.data[kr] = &ValuesAxis{
				Counts: make([]uint64, len(m.times)-1),
				Values: make([]uint64, len(m.times)-1),
			}
		}
		var value uint64
		switch m.typ {
		case "integration":
			value = region.WrittenBytes + region.ReadBytes
		case "read_keys":
			value = region.ReadKeys
		case "written_keys":
			value = region.WrittenKeys
		case "read_bytes":
			value = region.ReadBytes
		case "written_bytes":
			value = region.WrittenBytes
		}
		m.data[kr].Sum += value
		m.data[kr].Counts[idx]++
		m.data[kr].Values[idx] += value
	}
}

func (m *Merger) Matrix() Matrix {
	if len(m.data) == 0 {
		return Matrix{
			DataMap:  map[string][][]uint64{},
			KeyAxis:  []LabelKey{},
			TimeAxis: []int64{},
		}
	}
	kvs := make(KeyRangeValuesList, 0, len(m.data))
	for kr, va := range m.data {
		kvs = append(kvs, KeyRangeValues{
			KeyRange:   kr,
			ValuesAxis: *va,
		})
	}
	sort.Sort(kvs)
	kvs = divide(kvs, heatmapsMaxDisplayY)
	keys := make([]LabelKey, len(kvs)+1)
	if kvs[0].StartKey == "" {
		keys[0] = LabelKey{
			Key:    "",
			Labels: []string{"meta"},
		}
	} else {
		keys[0] = label(m.tables, kvs[0].StartKey)
	}
	values := make([][]uint64, len(m.times)-1)
	for n := 0; n < len(m.times)-1; n++ {
		values[n] = make([]uint64, len(kvs))
	}
	for i, kv := range kvs {
		if i == len(keys)-1 && kv.EndKey == "" {
			keys[i+1] = LabelKey{
				Key:    "",
				Labels: []string{},
			}
		} else {
			keys[i+1] = label(m.tables, kv.EndKey)
		}
		for j, v := range kv.ValuesAxis.Values {
			if kv.ValuesAxis.Counts[j] > 0 {
				values[j][i] = v / kv.ValuesAxis.Counts[j]
			}
		}
	}
	return Matrix{
		TimeAxis: m.times,
		KeyAxis:  keys,
		DataMap: map[string][][]uint64{
			m.typ: values,
		},
	}
}

func calcTimes(begin, end int64) (times []int64) {
	total := int64(0)
	step := int64(60)
	for end > begin {
		times = append(times, end)
		end -= step
		total += step
		if total == 60*60 {
			step = 2 * 60
		} else if total == 8*60*60 {
			step = 6 * 60
		} else if total == 24*60*60 {
			step = 30 * 60
		} else if total == 7*24*60*60 {
			step = 4 * 60 * 60
		}
	}
	times = append(times, begin)
	for i, j := 0, len(times)-1; i < j; i, j = i+1, j-1 { // reverse
		times[i], times[j] = times[j], times[i]
	}
	return
}

func findIndexFromTimes(times []int64, ts int64) int {
	for idx := range times {
		if idx+1 >= len(times) {
			break
		}
		if times[idx] <= ts && times[idx+1] > ts {
			return idx
		}
	}
	return -1
}

func divide(kvs KeyRangeValuesList, target int) KeyRangeValuesList {
	if len(kvs) < target {
		return kvs
	}
	var upper uint64
	for _, kv := range kvs {
		upper += kv.Sum
	}
	var lower uint64 = 1
	targetFocusRows := heatmapsMaxDisplayY * 2 / 3
	for lower < upper {
		mid := (lower + upper) >> 1
		if calcFocusRows(kvs, mid) > targetFocusRows {
			lower = mid + 1
		} else {
			upper = mid
		}
	}
	threshold := lower
	focusRows := calcFocusRows(kvs, threshold)
	ratio := len(kvs)/(target-focusRows) + 1
	return focus(kvs, threshold, ratio, target)
}

func calcFocusRows(kvs KeyRangeValuesList, threshold uint64) (count int) {
	start := 0
	var bucketSum uint64
	generateBucket := func(end int) {
		if end > start {
			count++
			start = end
			bucketSum = 0
		}
	}
	for i, kv := range kvs {
		if kv.Sum >= threshold || bucketSum >= threshold {
			generateBucket(i)
		}
		bucketSum += kv.Sum
	}
	generateBucket(len(kvs))
	return
}

func focus(kvs KeyRangeValuesList, threshold uint64, ratio, target int) KeyRangeValuesList {
	timesLen := len(kvs[0].Counts)
	newKvs := make(KeyRangeValuesList, 0, target)
	start := 0
	var bucketSum uint64
	generateBucket := func(end int) {
		if end <= start {
			return
		}
		newKv := KeyRangeValues{
			KeyRange: KeyRange{
				StartKey: kvs[start].StartKey,
				EndKey:   kvs[end-1].EndKey,
			},
			ValuesAxis: ValuesAxis{
				Counts: make([]uint64, timesLen),
				Values: make([]uint64, timesLen),
			},
		}
		for n := start; n < end; n++ {
			for i := 0; i < timesLen; i++ {
				newKv.Counts[i] += kvs[n].Counts[i]
				newKv.Values[i] += kvs[n].Values[i]
			}
		}
		newKvs = append(newKvs, newKv)
		start = end
		bucketSum = 0
	}
	for i, kv := range kvs {
		if kv.Sum >= threshold ||
			bucketSum >= threshold ||
			i-start >= ratio ||
			crossBorder(kv.KeyRange, kvs[start].KeyRange) {
			generateBucket(i)
		}
		bucketSum += kv.Sum
	}
	generateBucket(len(kvs))
	return newKvs
}

var keyBuffer model.KeyInfoBuffer

func label(tables map[int64]model.DBTableInfo, key string) LabelKey {
	lk := LabelKey{Key: key, Labels: []string{}}
	keyBytes, err := hex.DecodeString(key)
	if err != nil {
		return lk
	}
	keyInfo, _ := keyBuffer.DecodeKey(keyBytes)
	isMeta, tableID := keyInfo.MetaOrTable()
	if isMeta {
		lk.Labels = []string{"meta"}
		return lk
	}
	v, ok := tables[tableID]
	if ok {
		lk.Labels = append(lk.Labels, v.DB.Name.O, v.Table.Name.O)
	} else {
		lk.Labels = append(lk.Labels, fmt.Sprintf("table_%d", tableID))
	}
	if isCommonHandle, rowID := keyInfo.RowInfo(); isCommonHandle {
		lk.Labels = append(lk.Labels, "row")
	} else if rowID != 0 {
		lk.Labels = append(lk.Labels, fmt.Sprintf("row_%d", rowID))
	} else if indexID := keyInfo.IndexInfo(); indexID != 0 {
		if !ok {
			lk.Labels = append(lk.Labels, fmt.Sprintf("index_%d", indexID))
		} else if name, ok := v.Table.Indices[indexID]; ok {
			lk.Labels = append(lk.Labels, name)
		} else {
			lk.Labels = append(lk.Labels, fmt.Sprintf("index_%d", indexID))
		}
	}
	return lk
}

func crossBorder(kr1, kr2 KeyRange) bool {
	k1Bytes, err := hex.DecodeString(kr1.StartKey)
	if err != nil {
		return true
	}
	k2Bytes, err := hex.DecodeString(kr2.EndKey)
	if err != nil {
		return true
	}
	startInfo, _ := keyBuffer.DecodeKey(k1Bytes)
	startIsMeta, startTableID := startInfo.MetaOrTable()
	startIndex := startInfo.IndexInfo()
	endInfo, _ := keyBuffer.DecodeKey(k2Bytes)
	endIsMeta, endTableID := endInfo.MetaOrTable()
	endIndex := endInfo.IndexInfo()
	if startIsMeta || endIsMeta {
		return startIsMeta != endIsMeta
	}
	if startTableID != endTableID {
		return true
	}
	return startIndex != endIndex
}