package vm

import (
	"bytes"
	"context"
	"encoding/base64"
	"encoding/json"
	"sort"
	"strconv"
	"strings"
	"sync"

	"github.com/golang/snappy"
	"github.com/pingcap/log"
	"github.com/pingcap/tidb/util/plancodec"
	"go.uber.org/zap"

	"github.com/pingcap-inc/clinic/pkg/ngm/topsql/api"
)

const (
	queryTypeTopk    = "topk"
	queryTypeCPUTime = "cpu_time"
	queryTypeSummary = "summary"
	queryTypeMeta    = "meta"

	planMetaFetchLimit = 100
)

func (vq *VMQuery) Summary(ctx context.Context, params *api.SummaryParams) (res []api.SummaryItem, err error) {
	startSecs := params.StartSecs
	endSecs := params.EndSecs
	windowSecs := params.WindowSecs

	if startSecs > endSecs {
		return []api.SummaryItem{}, nil
	}

	// adjust start to make result align to end
	startSecs = endSecs - (endSecs-startSecs)/windowSecs*windowSecs
	rangeSecs := endSecs - startSecs + windowSecs

	args := metricsQueryArgs{
		Provider:  params.Provider,
		Region:    params.Region,
		TenantID:  params.TenantID,
		ProjectID: params.ProjectID,
		ClusterID: params.ClusterID,

		Instance:     params.Instance,
		InstanceType: params.InstanceType,
		Topk:         params.Top,

		StartSecs:  startSecs,
		EndSecs:    endSecs,
		WindowSecs: windowSecs,
		RangeSecs:  rangeSecs,
	}

	buf := bytes.NewBuffer(nil)

	var topkSqlDigests []string
	if topkSqlDigests, err = vq.fetchTopSQL(ctx, buf, &args); err != nil {
		log.Warn("failed to fetch topsql from vm", zap.Error(err))
		return nil, err
	}

	var cmap cpuTimesMap
	var smap summaryMap
	args.TopkRegex = strings.Join(topkSqlDigests, "|")
	if cmap, smap, err = vq.fetchCPUTimeAndSummary(ctx, buf, &args); err != nil {
		log.Warn("failed to fetch cpu time and summary from vm", zap.Error(err))
		return nil, err
	}

	var meta metaMap
	args.PlanDigestRegex = buildPlanDigestRegex(cmap, smap)
	if meta, err = vq.fetchMeta(ctx, buf, &args); err != nil {
		log.Warn("failed to fetch meta from vm", zap.Error(err))
		return nil, err
	}

	return buildResp(cmap, smap, meta), nil
}

// fetchTopSQL will fetch top SQL from VM, and update the args.TopkRegex
func (vq *VMQuery) fetchTopSQL(ctx context.Context, buf *bytes.Buffer, args *metricsQueryArgs) ([]string, error) {
	buf.Reset()
	if err := topkQueryTmp.Execute(buf, args); err != nil {
		return nil, err
	}
	body, err := vq.vmQuery(ctx, queryTypeTopk, buf.String(), args.EndSecs, args.toAddrArgs())
	if err != nil {
		return nil, err
	}

	return extractTopk(body)
}

// fetch CPU Time & summary, and update cpuTimeMap and summaryMap
func (vq *VMQuery) fetchCPUTimeAndSummary(
	ctx context.Context, buf *bytes.Buffer, args *metricsQueryArgs,
) (cmap cpuTimesMap, smap summaryMap, err error) {
	// 1. build queries
	buf.Reset()
	if err = cpuTimeQueryTmp.Execute(buf, args); err != nil {
		return
	}
	cpuTimeQuery := buf.String()

	buf.Reset()
	if err = summaryQueryTmp.Execute(buf, args); err != nil {
		return
	}
	summaryQuery := buf.String()

	// 2. fetch data in parallel
	var wg sync.WaitGroup
	wg.Add(2)
	go func() {
		defer wg.Done()
		body, queryErr := vq.vmQueryRange(ctx, queryTypeCPUTime, cpuTimeQuery, args.StartSecs, args.EndSecs, args.WindowSecs, args.toAddrArgs())
		if queryErr != nil {
			err = queryErr
			return
		}
		cmap, err = extractCPUTime(body)
	}()
	go func() {
		defer wg.Done()
		body, queryErr := vq.vmQuery(ctx, queryTypeSummary, summaryQuery, args.EndSecs, args.toAddrArgs())
		if queryErr != nil {
			err = queryErr
			return
		}
		smap, err = extractSummary(body)
	}()

	wg.Wait()
	return cmap, smap, err
}

// fetch meta from VM and update metaMap
func (vq *VMQuery) fetchMeta(ctx context.Context, buf *bytes.Buffer, args *metricsQueryArgs) (metaMap, error) {
	buf.Reset()
	if err := metaQueryTmp.Execute(buf, args); err != nil {
		return metaMap{}, err
	}
	body, err := vq.vmQuery(ctx, queryTypeMeta, buf.String(), args.EndSecs, args.toAddrArgs())
	if err != nil {
		return metaMap{}, err
	}

	return extractMeta(body)
}

// consume cpuTimeMap, summaryMap, metaMap to build response
func buildResp(cpuTimeMap cpuTimesMap, summaryMap summaryMap, meta metaMap) []api.SummaryItem {
	var res []api.SummaryItem

	// If existing SQL with cpu time but without summary, we fill it with empty SQL summary,
	// then the summary map is only need be to iterated once.
	for sqlDigest := range cpuTimeMap {
		if _, ok := summaryMap[sqlDigest]; !ok {
			summaryMap[sqlDigest] = &sqlSummary{PlanSummaryMap: map[planDigest]planSummary{}}
		}
	}

	for sqlDigest, sqlSummary := range summaryMap {
		item := api.SummaryItem{}
		item.SQLDigest = sqlDigest.Digest
		item.SQLText = meta.SQLMeta[sqlDigest.Digest]
		item.IsOther = sqlDigest.IsOther
		item.CPUTimeMs = sqlSummary.CPUTimeMs
		item.ExecCountPerSec = sqlSummary.ExecCountPerSec
		item.DurationPerExecMs = sqlSummary.DurationPerExecMs
		item.ScanRecordsPerSec = sqlSummary.ScanRecordsPerSec
		item.ScanIndexesPerSec = sqlSummary.ScanIndexesPerSec

		if _, ok := cpuTimeMap[sqlDigest]; !ok {
			cpuTimeMap[sqlDigest] = map[planDigest]cpuTimes{}
		}

		// If existing plan with cpu time but without summary, we fill it with empty plan summary,
		// then the plan summary map is only need be to iterated once.
		cpuMap := cpuTimeMap[sqlDigest]
		for planDigest := range cpuMap {
			if _, ok := sqlSummary.PlanSummaryMap[planDigest]; !ok {
				sqlSummary.PlanSummaryMap[planDigest] = planSummary{}
			}
		}

		for planDigest, planSummary := range sqlSummary.PlanSummaryMap {
			// avoid marshaling TimestampSec and CPUTimeMs to null
			var planExists bool
			var planCPU cpuTimes
			if planCPU, planExists = cpuMap[planDigest]; !planExists {
				planCPU = cpuTimes{TimestampSec: []uint64{}, CPUTimeMs: []uint64{}}
			}

			planItem := api.SummaryPlanItem{}
			planItem.PlanDigest = string(planDigest)
			planItem.PlanText = meta.PlanMeta[string(planDigest)]
			planItem.TimestampSec = planCPU.TimestampSec
			planItem.CPUTimeMs = planCPU.CPUTimeMs
			planItem.ExecCountPerSec = planSummary.ExecCountPerSec
			planItem.DurationPerExecMs = planSummary.DurationPerExecMs
			planItem.ScanRecordsPerSec = planSummary.ScanRecordsPerSec
			planItem.ScanIndexesPerSec = planSummary.ScanIndexesPerSec
			item.Plans = append(item.Plans, planItem)
		}
		res = append(res, item)
	}

	return res
}

// extractTopk extracts topk sql digests from response body
func extractTopk(body []byte) ([]string, error) {
	var resp topkResponse
	if err := json.Unmarshal(body, &resp); err != nil {
		return nil, err
	}

	var sqlDigests []string
	for _, item := range resp.Data.Result {
		sqlDigests = append(sqlDigests, item.Metric.SQLDigest)
	}

	return sqlDigests, nil
}

func extractCPUTime(body []byte) (cpuTimesMap, error) {
	var resp cpuTimeResponse
	if err := json.Unmarshal(body, &resp); err != nil {
		return nil, err
	}

	cmap := cpuTimesMap{}
	for _, item := range resp.Data.Result {
		ct := cpuTimes{}
		item.readReponseValues(func(t, c uint64) {
			ct.TimestampSec = append(ct.TimestampSec, t)
			ct.CPUTimeMs = append(ct.CPUTimeMs, c)
		})

		other, _ := strconv.ParseBool(item.Metric.IsOther)
		sqlDigest := sqlDigest{
			Digest:  item.Metric.SQLDigest,
			IsOther: other,
		}
		if _, ok := cmap[sqlDigest]; !ok {
			cmap[sqlDigest] = map[planDigest]cpuTimes{}
		}
		cmap[sqlDigest][item.Metric.PlanDigest] = ct
	}

	return cmap, nil
}

func extractSummary(body []byte) (summaryMap, error) {
	var resp summaryResponse
	if err := json.Unmarshal(body, &resp); err != nil {
		return nil, err
	}

	smap := summaryMap{}
	for _, item := range resp.Data.Result {
		grouped, _ := strconv.ParseBool(item.Metric.IsGrouped)

		sqlD := sqlDigest{Digest: item.Metric.SQLDigest, IsOther: false}
		if _, ok := smap[sqlD]; !ok {
			smap[sqlD] = &sqlSummary{PlanSummaryMap: map[planDigest]planSummary{}}
		}
		s := smap[sqlD]

		if grouped {
			switch item.Metric.SummaryType {
			case summaryTypeCPUTime:
				s.CPUTimeMs = uint64(item.readReponseValue())
			case summaryTypeExecCount:
				s.ExecCountPerSec = item.readReponseValue()
			case summaryTypeDuration:
				s.DurationPerExecMs = item.readReponseValue()
			case summaryTypeScanRecords:
				s.ScanRecordsPerSec = item.readReponseValue()
			case summaryTypeScanIndexes:
				s.ScanIndexesPerSec = item.readReponseValue()
			}
		} else {
			p := s.PlanSummaryMap[item.Metric.PlanDigest]
			switch item.Metric.SummaryType {
			case summaryTypeExecCount:
				p.ExecCountPerSec = item.readReponseValue()
			case summaryTypeDuration:
				p.DurationPerExecMs = item.readReponseValue()
			case summaryTypeScanRecords:
				p.ScanRecordsPerSec = item.readReponseValue()
			case summaryTypeScanIndexes:
				p.ScanIndexesPerSec = item.readReponseValue()
			}
			s.PlanSummaryMap[item.Metric.PlanDigest] = p
		}
	}

	// if there is no `other`, give it one
	if smap[sqlDigest{IsOther: true}] == nil {
		smap[sqlDigest{IsOther: true}] = &sqlSummary{PlanSummaryMap: map[planDigest]planSummary{"": {}}}
	}
	return smap, nil
}

// buildPlanDigestRegex collects plan digests from cpuTimesMap and summaryMap. If the count of plan digests is
// greater than planMetaFetchLimit, it will return top planMetaFetchLimit plan digests with the most cpu time.
func buildPlanDigestRegex(cpuTime cpuTimesMap, summary summaryMap) string {
	planDigests := make(map[string]struct{})
	var planCPUTimes []planCPUTime

	for _, planMap := range cpuTime {
		for planDigest, cpu := range planMap {
			if planDigest != "" {
				planDigests[string(planDigest)] = struct{}{}

				cpuTimeMs := uint64(0)
				for _, v := range cpu.CPUTimeMs {
					cpuTimeMs += v
				}
				planCPUTimes = append(planCPUTimes, planCPUTime{PlanDigest: string(planDigest), CPUTimeMs: cpuTimeMs})
			}
		}
	}

	for _, sqlSummary := range summary {
		for planDigest := range sqlSummary.PlanSummaryMap {
			if planDigest != "" {
				if _, ok := planDigests[string(planDigest)]; !ok {
					planCPUTimes = append(planCPUTimes, planCPUTime{PlanDigest: string(planDigest), CPUTimeMs: 0})
				}
			}
		}
	}

	sort.Slice(planCPUTimes, func(i, j int) bool {
		return planCPUTimes[i].CPUTimeMs > planCPUTimes[j].CPUTimeMs
	})
	if len(planCPUTimes) > planMetaFetchLimit {
		planCPUTimes = planCPUTimes[:planMetaFetchLimit]
	}

	var planSlice []string
	for _, plan := range planCPUTimes {
		planSlice = append(planSlice, plan.PlanDigest)
	}

	return strings.Join(planSlice, "|")
}

func extractMeta(body []byte) (metaMap, error) {
	var resp metaResponse
	if err := json.Unmarshal(body, &resp); err != nil {
		return metaMap{}, err
	}

	meta := metaMap{
		SQLMeta:  map[string]string{},
		PlanMeta: map[string]string{},
	}
	for _, item := range resp.Data.Result {
		if item.Metric.NormalizedSQL != "" {
			meta.SQLMeta[item.Metric.SQLDigest] = item.Metric.NormalizedSQL
		}
		if item.Metric.NormalizedPlan != "" {
			meta.PlanMeta[item.Metric.PlanDigest] = item.Metric.NormalizedPlan
		}
		if item.Metric.EncodedPlan != "" {
			decodedPlan, _ := decodePlan(item.Metric.EncodedPlan)
			meta.PlanMeta[item.Metric.PlanDigest] = decodedPlan
		}
	}

	return meta, nil
}

func decodePlan(planString string) (string, error) {
	binaryPlan, err := decompress(planString)
	if err != nil {
		return "", err
	}

	return plancodec.DecodeNormalizedPlan(binaryPlan)
}

func decompress(str string) (string, error) {
	decodeBytes, err := base64.StdEncoding.DecodeString(str)
	if err != nil {
		return "", err
	}

	bs, err := snappy.Decode(nil, decodeBytes)
	if err != nil {
		return "", err
	}
	return string(bs), nil
}

type metricsQueryArgs struct {
	Provider  string
	Region    string
	TenantID  uint64
	ProjectID uint64
	ClusterID uint64

	Instance     string
	InstanceType string
	Topk         int
	StartSecs    int
	EndSecs      int
	WindowSecs   int
	RangeSecs    int

	TopkRegex       string
	PlanDigestRegex string
}

func (a *metricsQueryArgs) toAddrArgs() *addrParams {
	return &addrParams{
		Provider:  a.Provider,
		Region:    a.Region,
		ClusterID: a.ClusterID,
	}
}

/// Aggregation structs

type cpuTimes struct {
	TimestampSec []uint64
	CPUTimeMs    []uint64
}

type sqlDigest struct {
	Digest  string
	IsOther bool
}

type planDigest string

type cpuTimesMap map[sqlDigest]map[planDigest]cpuTimes

type summaryMap map[sqlDigest]*sqlSummary

type sqlSummary struct {
	CPUTimeMs         uint64
	ExecCountPerSec   float64
	DurationPerExecMs float64
	ScanRecordsPerSec float64
	ScanIndexesPerSec float64
	PlanSummaryMap    map[planDigest]planSummary
}

type planSummary struct {
	ExecCountPerSec   float64
	DurationPerExecMs float64
	ScanRecordsPerSec float64
	ScanIndexesPerSec float64
}

type metaMap struct {
	SQLMeta  map[string]string
	PlanMeta map[string]string
}

type planCPUTime struct {
	PlanDigest string
	CPUTimeMs  uint64
}

/// Response structs

type topkResponse struct {
	Data topkResponseData `json:"data"`
}

type topkResponseData struct {
	Result []topkResponseResult `json:"result"`
}

type topkResponseResult struct {
	Metric topkResponseMetric `json:"metric"`
}

type topkResponseMetric struct {
	SQLDigest string `json:"sql_digest"`
}

type cpuTimeResponse struct {
	Data cpuTimeResponseData `json:"data"`
}

type cpuTimeResponseData struct {
	Result []cpuTimeResponseResult `json:"result"`
}

type cpuTimeResponseResult struct {
	Metric cpuTimeResponseMetric `json:"metric"`
	// [[1664269812,"30"],[1664269821,"10"],[1664269839,"20"]]
	Values [][]any `json:"values"`
}

func (c *cpuTimeResponseResult) readReponseValues(consumer func(uint64, uint64)) {
	for _, value := range c.Values {
		ts := value[0].(float64)
		v, _ := strconv.ParseFloat(value[1].(string), 64)
		consumer(uint64(ts), uint64(v))
	}
}

type cpuTimeResponseMetric struct {
	SQLDigest  string     `json:"sql_digest"`
	PlanDigest planDigest `json:"plan_digest"`
	IsOther    string     `json:"is_other"`
}

type summaryResponse struct {
	Data summaryResponseData `json:"data"`
}

type summaryResponseData struct {
	Result []summaryResponseResult `json:"result"`
}

type summaryResponseResult struct {
	Metric summaryResponseMetric `json:"metric"`
	// [1664269812,"30"]
	Value []any `json:"value"`
}

func (s *summaryResponseResult) readReponseValue() float64 {
	if len(s.Value) != 2 {
		return 0
	}
	v, _ := strconv.ParseFloat(s.Value[1].(string), 64)
	return v
}

type summaryResponseMetric struct {
	SQLDigest   string     `json:"sql_digest"`
	SummaryType string     `json:"summary_type"`
	PlanDigest  planDigest `json:"plan_digest"`
	IsGrouped   string     `json:"is_grouped"`
}

const (
	summaryTypeCPUTime     = "cpu_time_ms"
	summaryTypeExecCount   = "exec_count_per_sec"
	summaryTypeDuration    = "duration_per_exec_ms"
	summaryTypeScanRecords = "scan_records_per_sec"
	summaryTypeScanIndexes = "scan_indexes_per_sec"
)

type metaResponse struct {
	Data metaResponseData `json:"data"`
}

type metaResponseData struct {
	Result []metaResponseResult `json:"result"`
}

type metaResponseResult struct {
	Metric metaResponseMetric `json:"metric"`
}

type metaResponseMetric struct {
	SQLDigest      string `json:"sql_digest"`
	NormalizedSQL  string `json:"normalized_sql"`
	PlanDigest     string `json:"plan_digest"`
	NormalizedPlan string `json:"normalized_plan"`
	EncodedPlan    string `json:"encoded_normalized_plan"`
}