package query

import (
	"context"
	"encoding/json"
	"errors"
	"strings"

	"github.com/pingcap/log"
	"github.com/tidbcloud/serverless-metering/pkg/store"
	"go.uber.org/zap"
)

type IntermediateQueryEngine struct {
	storage store.Storage
	ts      int64
	offset  int
}

func NewIntermediateQueryEngine(ctx context.Context, storage store.Storage, ts int64, offset int) (QueryEngine, error) {
	return &IntermediateQueryEngine{storage: storage, ts: ts, offset: offset}, nil
}

func (e *IntermediateQueryEngine) Close() error {
	return nil
}

func (e *IntermediateQueryEngine) CalLoadDataRu(ctx context.Context) ([]*Usage, error) {
	return nil, errors.New("not implemented")
}

func (e *IntermediateQueryEngine) CalPublicEgress(ctx context.Context) ([]*Usage, error) {
	return nil, errors.New("not implemented")
}

func (e *IntermediateQueryEngine) CalPrivatelinkEgress(ctx context.Context) ([]*Usage, error) {
	return nil, errors.New("not implemented")
}

func (e *IntermediateQueryEngine) CalTidbRU(ctx context.Context) ([]*Usage, error) {
	return nil, errors.New("not implemented")
}

func (e *IntermediateQueryEngine) CalTiflashRU(ctx context.Context) ([]*Usage, error) {
	return nil, errors.New("not implemented")
}

func (e *IntermediateQueryEngine) CalTidbRUV2(ctx context.Context) ([]*Usage, error) {
	return nil, errors.New("not implemented")
}

func (e *IntermediateQueryEngine) CalTiflashRUV2(ctx context.Context) ([]*Usage, error) {
	return nil, errors.New("not implemented")
}

func (e *IntermediateQueryEngine) CalTikvStorageMedium(ctx context.Context) ([]*Usage, error) {
	return e.aggregateIntermediateResult(ctx, MetricTypeTikvStorage)
}

func (e *IntermediateQueryEngine) CalTikvStorageIAMedium(ctx context.Context) ([]*Usage, error) {
	return e.aggregateIntermediateResult(ctx, MetricTypeTikvStorageIA)
}

func (e *IntermediateQueryEngine) CalTikvStorageIALarge(ctx context.Context) ([]*Usage, error) {
	return nil, errors.New("not implemented")
}

func (e *IntermediateQueryEngine) CalTiflashStorageMedium(ctx context.Context) ([]*Usage, error) {
	return e.aggregateIntermediateResult(ctx, MetricTypeTiflashStorage)
}

func (e *IntermediateQueryEngine) CalTikvStorageLarge(ctx context.Context) ([]*Usage, error) {
	return nil, errors.New("not implemented")
}

func (e *IntermediateQueryEngine) CalTiflashStorageLarge(ctx context.Context) ([]*Usage, error) {
	return nil, errors.New("not implemented")
}

func (e *IntermediateQueryEngine) CalIntermediateTikvStorageMedium(ctx context.Context) ([]*Usage, error) {
	return nil, errors.New("not implemented")
}

func (e *IntermediateQueryEngine) CalIntermediateTikvStorageIAMedium(ctx context.Context) ([]*Usage, error) {
	return nil, errors.New("not implemented")
}

func (e *IntermediateQueryEngine) CalIntermediateTiflashStorageMedium(ctx context.Context) ([]*Usage, error) {
	return nil, errors.New("not implemented")
}

func (e *IntermediateQueryEngine) CalRCU(ctx context.Context) ([]*Usage, error) {
	return nil, errors.New("not implemented")
}

func (e *IntermediateQueryEngine) CalExecCount(ctx context.Context) (ExecCount, error) {
	return nil, errors.New("not implemented")
}

// NOTE: Currently this func only support storage medium. It's determined by s3 file name and format.
// When more data need use intermediate result. Remember to double check s3 file name.
func (e *IntermediateQueryEngine) aggregateIntermediateResult(ctx context.Context, mt MetricType) ([]*Usage, error) {
	tmp := make(map[string]int64)
	for i := 0; i < e.offset; i++ {
		scanTs := e.ts + int64(i*60*60)
		data, err := e.storage.GetIntermediateResult(ctx, scanTs, string(mt))
		log.Debug("get intermediate result", zap.Int64("scanTs", scanTs))
		if err != nil {
			if strings.Contains(err.Error(), "NoSuchKey") {
				log.Debug("skip intermediate result", zap.Int64("scanTs", scanTs))
				continue
			}
			return nil, err
		}
		usages := make([]*Usage, 0)
		err = json.Unmarshal(data, &usages)
		if err != nil {
			return nil, err
		}
		log.Debug("get intermediate result success", zap.Int64("scanTs", scanTs), zap.Int("len", len(usages)))
		for _, usage := range usages {
			v, ok := tmp[usage.KeyspaceId]
			if !ok {
				tmp[usage.KeyspaceId] = usage.Value
			} else {
				tmp[usage.KeyspaceId] = v + usage.Value
			}
		}
	}

	result := make([]*Usage, 0, len(tmp))
	for keyspaceId, value := range tmp {
		result = append(result, &Usage{
			KeyspaceId: keyspaceId,
			Value:      value,
			MetricType: mt,
		})
	}
	return result, nil
}