package debugportal

import (
	"context"
	"encoding/hex"
	"errors"
	"sort"

	topsqlmodel "github.com/pingcap/dbaas/central/model/debug/topsql"
	"github.com/pingcap/log"
	proto "github.com/tidbcloud/idl/pbgen/proto/central"
	"go.uber.org/zap"
	"go.uber.org/zap/zapcore"
)

const (
	// othersKey is the key representing aggregated other top sql digests, i.e. others
	othersKey = ""
)

func (service *debugPortalService) TopSQLGetInstances(_ context.Context, req *proto.TopSQLGetInstancesReq) (*proto.TopSQLGetInstanceResp, error) {
	rows, err := service.topsqlModel.ListInstances(req.ClusterId, req.BeginTs, req.EndTs)
	if err != nil {
		log.Error("failed to call ListInstances", zap.Error(err))
		return &proto.TopSQLGetInstanceResp{}, err
	}
	instances := make([]*proto.TopSQLInstance, 0, len(rows))
	for _, row := range rows {
		instances = append(instances, &proto.TopSQLInstance{Id: row.ID, Name: row.Name, Component: row.Component})
	}
	resp := &proto.TopSQLGetInstanceResp{
		Instances: instances,
	}
	return resp, nil
}

type DigestRecordMap map[string]*proto.TopSQLDigestCPUTimeRecord

func (m DigestRecordMap) MarshalLogObject(enc zapcore.ObjectEncoder) error {
	for k, v := range m {
		if err := enc.AddReflected(hex.EncodeToString([]byte(k)), v.Data); err != nil {
			log.Error("failed to encode data for DigestRecordMap",
				zap.Error(err),
				zap.String("key", k),
				zap.Any("value", v),
			)
		}
	}
	return nil
}

type TopNSqlDigests map[string]struct{}

func (m TopNSqlDigests) MarshalLogObject(enc zapcore.ObjectEncoder) error {
	for k := range m {
		enc.AddString("digest", hex.EncodeToString([]byte(k)))
	}
	return nil
}

// TopSQLGetTopDigests returns top N and a summed "others" digests measured in CPU time ms.
//
// Consider there's W time windows in the query result, each has N_i(N_i <= req.TopN) digests and an optional special others_i, then there's possibly
// at most sum(N_1, ..., N_W) different sql statements in the query result, and W others, i.e. others_1, ..., others_W.
// Because we cannot tell from others_1, ..., others_W, which has empty sql digest, we should first add them up into a single `others`.
// Then, suppose there are U unique sql statements from these W time windows, there are 2 possible situations:
// 1. req.TopN >= U, then just return U sql statements and `others`, which actually is less than the request's expectation
// 2. req.TopN < U, we should count all the sql statements except for the topN as `others`
func (service *debugPortalService) TopSQLGetTopDigests(_ context.Context, req *proto.TopSQLGetTopDigestsReq) (*proto.TopSQLGetTopDigestResp, error) {
	// check parameters
	if req.By != "cpu" {
		return nil, errors.New("invalid parameter by, should be cpu")
	}
	if req.TopN <= 0 {
		return nil, errors.New("invalid parameter top_n, should be > 0")
	}
	if req.WindowSizeSec <= 0 {
		return nil, errors.New("invalid parameter window_size_sec, should be > 0 ")
	}

	// get data from DB
	log.Debug("calling topsqlModel.GetTopCPUTimeDigestsWindowed", zap.Any("req", req))
	rows, err := service.topsqlModel.GetTopCPUTimeDigestsWindowed(req.ClusterId, req.WindowSizeSec, req.TopN, req.BeginTs, req.EndTs, req.InstanceId)
	if err != nil {
		log.Error("failed to call GetTopCPUTimeDigestsWindowed", zap.Error(err))
		return &proto.TopSQLGetTopDigestResp{}, err
	}
	if len(rows) == 0 {
		return &proto.TopSQLGetTopDigestResp{}, nil
	}

	// group the records by sql digest, across all the time windows
	digestRecordMap := groupRowsBySqlDigest(rows, uint64(req.WindowSizeSec))
	log.Debug("groupRowsBySqlDigest()", zap.Any("map", digestRecordMap))

	// remove the possible `others` from query results, which is collected by the agent
	// we need the `others` in API response to be calculated dynamically according to the req.topN
	delete(digestRecordMap, othersKey)

	// sort and find the top n sql digests
	sqlCpuTimeList := calcSqlDigestCpuTimeSum(digestRecordMap)
	log.Debug("calcSqlDigestCpuTimeSum()", zap.Any("sqlTimeList", sqlCpuTimeList))
	topNSqlDigests := calcTopNSqlDigestByCpuTime(sqlCpuTimeList, int(req.TopN))
	log.Debug("calcTopNSqlDigestByCpuTime()", zap.Any("topNSqlDigests", topNSqlDigests))
	// remove sql digests other than top n
	for digest := range digestRecordMap {
		if _, ok := topNSqlDigests[digest]; !ok {
			delete(digestRecordMap, digest)
		}
	}

	// find the number of records in database query results in each time window
	windowSlots := calcWindowSlots(rows)
	log.Debug("calcWindowSlots()", zap.Ints("windowSlots", windowSlots))
	// calculate `others` dynamically according to request
	others := calcOthers(rows, windowSlots, topNSqlDigests, req.WindowSizeSec)
	log.Debug("calcOthers()", zap.Any("others", others))

	resp := &proto.TopSQLGetTopDigestResp{
		Records: make([]*proto.TopSQLDigestCPUTimeRecord, 0),
	}
	for _, v := range digestRecordMap {
		resp.Records = append(resp.Records, v)
	}
	if len(others) > 0 {
		othersRecord := &proto.TopSQLDigestCPUTimeRecord{
			IsOthers:  true,
			SqlDigest: []byte(""),
			SqlText:   "",
			Data:      others,
		}
		resp.Records = append(resp.Records, othersRecord)
	}

	// select SQL text from database for all the SQL digests
	sqlDigestList := make([]string, 0, len(resp.Records))
	for _, r := range resp.Records {
		sqlDigestList = append(sqlDigestList, string(r.SqlDigest))
	}
	rowsSqlText, err := service.topsqlModel.GetNormalizedSQLText(req.ClusterId, sqlDigestList)
	if err != nil {
		log.Error("failed to call GetNormalizedSQLText", zap.Error(err))
	}
	// join sql digests with texts
	digestsToJoin := make(map[string]string, len(sqlDigestList))
	for _, row := range rowsSqlText {
		digestsToJoin[string(row.SqlDigest)] = row.NormalizedSql
	}
	for _, record := range resp.Records {
		record.SqlText = digestsToJoin[string(record.SqlDigest)]
	}

	return resp, nil
}

// calcOthers calculate others in each time window, and append the (timestamp, cpu time) tuple
func calcOthers(rows []topsqlmodel.TopCPUTimeDigestsWindowedRow, windowSlots []int, topNSqlDigests TopNSqlDigests, windowSizeSec uint32) []*proto.TopSQLTimestampCPUMsTuple {
	others := make([]*proto.TopSQLTimestampCPUMsTuple, 0)
	i := 0
	for _, w := range windowSlots {
		topNCpuTimeSum := uint64(0)
		for j := 0; j < w; j++ {
			row := rows[i+j]
			if _, ok := topNSqlDigests[string(row.SqlDigest)]; ok {
				// this row is in the top n digest set
				topNCpuTimeSum += row.CpuTimeAgg
			}
		}
		row := rows[i]
		othersCpuTime := row.CpuTimeSumInWindow - topNCpuTimeSum
		if othersCpuTime > 0 {
			others = append(others, &proto.TopSQLTimestampCPUMsTuple{
				Timestamp: row.TimeWindow * uint64(windowSizeSec),
				CpuTimeMs: uint32(othersCpuTime),
			})
		}
		i += w
	}
	return others
}

func calcTopNSqlDigestByCpuTime(sqlCpuTimeList []sqlDigestCPUTimeTotal, topN int) TopNSqlDigests {
	if topN > len(sqlCpuTimeList) {
		topN = len(sqlCpuTimeList)
	}
	sort.Sort(ByCPUTimeTotalDescSqlDigestDesc(sqlCpuTimeList))
	// topNSqlDigests is a set of the top N sql digests, for calculating cpu time for `others`
	topNSqlDigests := make(TopNSqlDigests)
	for _, sqlTime := range sqlCpuTimeList[:topN] {
		topNSqlDigests[sqlTime.SqlDigest] = struct{}{}
	}
	return topNSqlDigests
}

func calcSqlDigestCpuTimeSum(digestRecordMap DigestRecordMap) []sqlDigestCPUTimeTotal {
	sqlTimeList := make([]sqlDigestCPUTimeTotal, 0, len(digestRecordMap))
	for digest, record := range digestRecordMap {
		sqlTime := sqlDigestCPUTimeTotal{
			SqlDigest:    digest,
			CpuTimeTotal: 0,
		}
		// sum cpu time for the sql digest
		for _, tuple := range record.Data {
			sqlTime.CpuTimeTotal += uint64(tuple.CpuTimeMs)
		}
		sqlTimeList = append(sqlTimeList, sqlTime)
	}
	return sqlTimeList
}

func groupRowsBySqlDigest(rows []topsqlmodel.TopCPUTimeDigestsWindowedRow, windowSizeSec uint64) DigestRecordMap {
	groupByDigest := make(DigestRecordMap)
	for _, row := range rows {
		// if sql digest is "", it's others in this time window, and will be accumulated through all the time windows
		sqlDigest := row.SqlDigest
		record, ok := groupByDigest[string(sqlDigest)]
		if !ok {
			record = &proto.TopSQLDigestCPUTimeRecord{
				IsOthers:  false,
				SqlDigest: sqlDigest,
				SqlText:   "",
				Data:      make([]*proto.TopSQLTimestampCPUMsTuple, 0),
			}
			groupByDigest[string(sqlDigest)] = record
		}
		record.Data = append(record.Data, &proto.TopSQLTimestampCPUMsTuple{
			// time window is floor(timestamp/window_size), which means the Timestamp field is the start of its time window
			Timestamp: row.TimeWindow * windowSizeSec,
			// averaged over window_size_sec
			CpuTimeMs: uint32(row.CpuTimeAgg),
		})
	}
	return groupByDigest
}

// calcWindowSlots calculates the number of records in each time windows in rows
func calcWindowSlots(rows []topsqlmodel.TopCPUTimeDigestsWindowedRow) []int {
	if len(rows) == 0 {
		return []int{}
	}
	if len(rows) == 1 {
		return []int{1}
	}
	windowSlots := make([]int, 0)
	for i := 0; i < len(rows)-1; {
		for j := i; j < len(rows)-1; j++ {
			if rows[j].TimeWindow != rows[j+1].TimeWindow {
				windowSlots = append(windowSlots, j-i+1)
				i = j + 1
				break
			} else if j == len(rows)-2 {
				windowSlots = append(windowSlots, len(rows)-i)
				i = len(rows) - 1 // terminates the outer loop
				break
			}
		}
	}
	return windowSlots
}

type sqlDigestCPUTimeTotal struct {
	SqlDigest    string
	CpuTimeTotal uint64
}

type ByTimestampAsc []*proto.TopSQLTimestampCPUMsTuple

func (x ByTimestampAsc) Len() int           { return len(x) }
func (x ByTimestampAsc) Less(i, j int) bool { return x[i].Timestamp < x[j].Timestamp }
func (x ByTimestampAsc) Swap(i, j int)      { x[i], x[j] = x[j], x[i] }

type ByCPUTimeTotalDescSqlDigestDesc []sqlDigestCPUTimeTotal

func (x ByCPUTimeTotalDescSqlDigestDesc) Len() int { return len(x) }
func (x ByCPUTimeTotalDescSqlDigestDesc) Less(i, j int) bool {
	if x[i].CpuTimeTotal == x[j].CpuTimeTotal {
		return x[i].SqlDigest > x[j].SqlDigest
	}
	return x[i].CpuTimeTotal > x[j].CpuTimeTotal
}
func (x ByCPUTimeTotalDescSqlDigestDesc) Swap(i, j int) { x[i], x[j] = x[j], x[i] }