package report

import (
	"context"
	"encoding/json"
	"fmt"
	"sort"
	"strings"
	"time"

	"github.com/pingcap/log"
	"github.com/tidbcloud/serverless-metering/pkg/config"
	"github.com/tidbcloud/serverless-metering/pkg/monitor"
	"github.com/tidbcloud/serverless-metering/pkg/query"
	"github.com/tidbcloud/serverless-metering/pkg/repository"
	"github.com/tidbcloud/serverless-metering/pkg/store"
	"github.com/tidbcloud/serverless-metering/pkg/util"
	"github.com/tidbcloud/tiinf/sdk/measure"
	"go.uber.org/zap"
)

const (
	BillingLabelTenantID         = "tenantID"
	BillingLabelProjectID        = "projectID"
	BillingLabelProjectName      = "project"
	BillingLabelClusterID        = "clusterID"
	BillingLabelClusterName      = "cluster"
	BillingLabelBranchID         = "branchID"
	BillingLabelBranchName       = "branch"
	BillingLabelRegion           = "region"
	BillingLabelServerless       = "serverless"
	BillingLabelMetricName       = "metric"
	BillingLabelSkuPlan          = "skuPlan"
	BillingLabelWorker           = "worker"
	BillingLabelComponent        = "Component"
	BillingComponentTiFlash      = "tiflash"
	BillingLabelBetaKey          = "beta" /* used for beta phase, removal or 'disable' upon official release */
	BillingBetaCategoryBranching = "Branching"
	BillingLabelBetaEnable       = "true"
	BillingLabelBetaCategoryKey  = "beta-category" /* used for beta phase, removal upon official release */
	BillingLabelKeyspaceId       = "keyspaceID"
	BillingBetaCategoryTraffic   = "Data Transfer Discount"
	BillingBetaCategoryPoc       = "Internal Usage"
	BillingBetaCategoryPreview   = "Preview"
	BillingLabelDatabaseKey      = "database"
	BillingLabelRUVersion        = "ruVersion"
	BillingLabelStorageClass     = "storageClass"

	WorkerNameDefault             = "worker-default"
	WorkerNameStorageMedium       = "worker-storage-medium"
	WorkerNameStorageMediumOneDay = "worker-storage-medium-one-day"

	IntervalDefault             = 1
	IntervalStorageMedium       = 2
	IntervalStorageMediumOneDay = 24
)

type meteringWorker struct {
	name                 string
	reportIntervalHour   int
	storage              store.Storage
	repo                 repository.ClusterBranchRepo
	finishedTs           int64
	cloudProvider        string
	region               string
	storeCheckPointNames []string
	enableReporter       bool
	billingReporter      *measure.MeasureSQSImpl
	dryRun               bool
	pocTenantIDs         map[uint64]struct{}
	pocFreePercent       int
	dbRuTenantIDs        map[uint64]struct{}
}

func newMeteringWorker(
	storage store.Storage,
	repo repository.ClusterBranchRepo,
	name string,
	reportIntervalHour int,
	dryRun bool,
) (*meteringWorker, error) {
	cfg := config.Global()
	worker := &meteringWorker{
		name:                 name,
		reportIntervalHour:   reportIntervalHour,
		storage:              storage,
		repo:                 repo,
		cloudProvider:        cfg.CloudProvider,
		region:               cfg.Region,
		storeCheckPointNames: make([]string, len(cfg.Report.StoreMembers)),
		enableReporter:       cfg.Report.EnableReporter,
		billingReporter:      &measure.MeasureSQSImpl{},
		dryRun:               dryRun,
		pocTenantIDs:         make(map[uint64]struct{}),
		pocFreePercent:       cfg.Report.PocFreePercent,
	}
	for _, id := range cfg.Report.PocTenantIDs {
		worker.pocTenantIDs[id] = struct{}{}
	}
	if cfg.DBLevelRU.Enable && len(cfg.DBLevelRU.TenantIDs) > 0 {
		worker.dbRuTenantIDs = make(map[uint64]struct{})
		for _, id := range cfg.DBLevelRU.TenantIDs {
			worker.dbRuTenantIDs[id] = struct{}{}
		}
	}

	if worker.enableReporter {
		if err := worker.billingReporter.InitMeasureSqsHandler(cfg.Report.ReporterEnvType, ""); err != nil {
			return nil, err
		}
	}

	ctx := context.TODO()
	checkpoint, err := storage.GetCheckPoint(ctx, worker.name)
	if err != nil {
		if strings.Contains(err.Error(), "NoSuchKey") {
			lastTs := util.LocateLatestOnTheHour(time.Now().Unix() - int64(reportIntervalHour*60*60))
			err = storage.PutCheckPoint(ctx, worker.name, &store.CheckPoint{Value: lastTs})
			if err != nil {
				return nil, err
			}
			worker.finishedTs = lastTs
		} else {
			return nil, err
		}
	} else {
		worker.finishedTs = checkpoint.Value
	}
	monitor.ReportProgressGauge.WithLabelValues(worker.name).Set(float64(worker.finishedTs))

	for i, member := range cfg.Report.StoreMembers {
		worker.storeCheckPointNames[i] = store.GetCheckPointName(member)
	}

	return worker, nil
}

// default worker meter ru and storage large
func newDefaultMeteringWorker(storage store.Storage, repo repository.ClusterBranchRepo) (*meteringWorker, error) {
	return newMeteringWorker(storage, repo, WorkerNameDefault, IntervalDefault, false)
}

func newStorageMediumWorker(storage store.Storage, repo repository.ClusterBranchRepo) (*meteringWorker, error) {
	return newMeteringWorker(storage, repo, WorkerNameStorageMedium, IntervalStorageMedium, false)
}

func newStorageMediumWorkerOneDay(storage store.Storage, repo repository.ClusterBranchRepo) (*meteringWorker, error) {
	return newMeteringWorker(storage, repo, WorkerNameStorageMediumOneDay, IntervalStorageMediumOneDay, false)
}

func (m *meteringWorker) Run(ctx context.Context, debug bool) {
	nextTs := m.NextTs()
	nowTs := time.Now().Unix()
	logger := log.L().With(zap.String("name", m.name), zap.Int64("nextTs", nextTs), zap.Int64("nowTs", nowTs), zap.Int64("finishedTs", m.finishedTs))
	if !debug && !m.shouldRun(ctx, logger, nowTs, nextTs) {
		return
	}
	m.run(ctx, logger, nextTs, debug)
}

func (m *meteringWorker) shouldRun(ctx context.Context, logger *zap.Logger, nowTs, nextTs int64) bool {
	if nowTs < nextTs {
		logger.Info("current round not started")
		return false
	}

	ft := store.GetFileType()
	for _, name := range m.storeCheckPointNames {
		checkpoint, err := m.storage.GetCheckPoint(ctx, name)
		if err != nil {
			logger.Error("get store checkpoint name failed", zap.Error(err))
			return false
		}
		var value int64
		switch ft {
		case store.FileTypeParquet:
			value = checkpoint.Value + 60*60
		case store.FileTypeJson:
			value = checkpoint.Value
		default:
			log.Error("unknown file type")
			return false
		}

		if value < nextTs {
			logger.Info("store checkpoint less than next ts, waitting", zap.String("checkPointName", name), zap.Int64("checkPointValue", checkpoint.Value))
			return false
		}
	}

	return true
}

func (m *meteringWorker) run(ctx context.Context, logger *zap.Logger, nextTs int64, debug bool) {
	logger.Info("start running")
	reports, err := m.genBillingReports(ctx, logger)
	if err != nil {
		logger.Error("failed to generate billing reports", zap.Error(err))
		return
	}
	logger.Info("billing reports length", zap.Int("len", len(reports)))
	for _, report := range reports {
		if !report.ShouldLog {
			continue
		}
		text, _ := json.Marshal(report)
		logger.Info("billing records for analysis", zap.String("version", "2"), zap.String("record", string(text)), zap.Bool("dryRun", m.dryRun))
	}

	if debug {
		return
	}

	// TODO: add metrics
	if m.enableReporter && !m.dryRun {
		for _, report := range reports {
			bm := measure.BillingMeasure{
				TenantID:      report.TenantID,
				CostID:        report.CostID,
				Cost:          report.Cost,
				CostTimestamp: report.CostTimeStamp,
				Labels:        report.Labels,
				Unit:          report.Sku,
			}
			if err := m.billingReporter.CheckMeasure(bm); err != nil {
				logger.Warn("skip invalid measure", zap.Error(err), zap.String("CostID", bm.CostID))
				continue
			}
			if err := m.billingReporter.SendMeasure(bm); err != nil {
				logger.Error("send measure to billing failed", zap.Error(err), zap.String("CostID", bm.CostID))
				return
			}
			logger.Debug("send measure to billing success", zap.String("CostID", bm.CostID))
		}
	}

	if err = m.storage.PutCheckPoint(ctx, m.name, &store.CheckPoint{Value: nextTs}); err != nil {
		logger.Error("store metering manager checkpoint failed", zap.Error(err))
		return
	}
	m.finishedTs = nextTs
	monitor.ReportProgressGauge.WithLabelValues(m.name).Set(float64(m.finishedTs))
}

func (m *meteringWorker) genBillingReports(ctx context.Context, logger *zap.Logger) ([]*BillingReport, error) {
	usages, execCount, err := m.queryUsages(ctx)
	if err != nil {
		return nil, err
	}
	logger.Debug("usages length", zap.Int("len", len(usages)), zap.Int("execCount", len(execCount)))

	reports := make([]*BillingReport, 0)
	for _, usage := range usages {
		err = usage.FillCluster(ctx, m.repo)
		if err != nil {
			if util.IsNotFoundError(err) {
				logger.Debug("usage key not found", zap.Error(err), zap.Any("usage", usage))
				continue
			}
			logger.Error("failed to fill cluster to usage", zap.Error(err), zap.Any("usage", usage))
			return nil, err
		}

		cluster := usage.GetMetaCluster()

		switch usage.MetricType {
		case query.MetricTypeTidbRUV2, query.MetricTypeTiflashRUV2:
			logger.Debug("ru usage v2", zap.Any("usage", usage), zap.Any("cluster", cluster))
			if !cluster.UseRUV2() {
				logger.Debug("skip ruv2 usage for non-ru_version=2 cluster", zap.Any("cluster", cluster))
				continue
			}
		case query.MetricTypeTidbRU, query.MetricTypeTiflashRU:
			logger.Debug("ru usage v1", zap.Any("usage", usage), zap.Any("cluster", cluster))
			if cluster.UseRUV2() {
				logger.Debug("skip ruv1 usage for ru_version=2 cluster", zap.Any("cluster", cluster))
				continue
			}
		}

		report := m.convertUsageToBillingReport(cluster, usage)
		if _, ok := m.pocTenantIDs[cluster.TenantID]; ok {
			percentMap := map[string]int64{
				"poc": int64(m.pocFreePercent),
				"":    int64(100 - m.pocFreePercent),
			}
			costMap := m.splitCost(report.Cost, percentMap)
			for k, v := range costMap {
				r := report.DeepCopy()
				r.Cost = v
				r.CostID += fmt.Sprintf("/%s", k)
				m.applyRUVersionLabels(r, cluster, usage)
				if k == "poc" {
					r.Labels[BillingLabelBetaKey] = BillingLabelBetaEnable
					r.Labels[BillingLabelBetaCategoryKey] = BillingBetaCategoryPreview
				}
				reports = append(reports, r)
			}
		} else if _, ok := m.dbRuTenantIDs[cluster.TenantID]; ok && usage.IsRu() {
			percentMap := execCount.GetExecCountOfEachDB(fmt.Sprint(cluster.ClusterID))
			costMap := m.splitCost(report.Cost, percentMap)
			for k, v := range costMap {
				r := report.DeepCopy()
				r.Cost = v
				// Truncate database name to 100 chars in CostID, but keep full name in label
				dbNameForCostID := k
				if len(dbNameForCostID) > 100 {
					dbNameForCostID = dbNameForCostID[:100]
				}
				r.CostID += fmt.Sprintf("/db_%s", dbNameForCostID)
				r.Labels[BillingLabelDatabaseKey] = k
				m.applyRUVersionLabels(r, cluster, usage)
				reports = append(reports, r)
			}
		} else {
			reports = append(reports, report)
		}
	}
	return reports, nil
}

func (m *meteringWorker) convertUsageToBillingReport(cluster *repository.ClusterBranch, usage *query.Usage) *BillingReport {
	report := &BillingReport{
		CostID:   fmt.Sprintf("%s/%s/%s/%d", m.region, cluster.PrimaryID, usage.GetMetricTypeString(), m.finishedTs),
		TenantID: cluster.TenantID,
		Sku:      usage.GetSkuName(m.cloudProvider, m.region),
		Cost:     usage.GetCost(),
		// 20250707 change cost timestamp from finishedTs to nextTs - 1.
		// finishedTs may be ahead of cluster's creation time and will cause billing credits can't cover part of the usage.
		CostTimeStamp: m.NextTs() - 1,
		Labels:        m.buildLabels(cluster, usage),
		ShouldLog:     true,
	}
	if usage.MetricType == query.MetricTypeTikvStorageIA {
		report.Labels[BillingLabelStorageClass] = "ia"
	}
	m.applyRUVersionLabels(report, cluster, usage)
	return report
}

func (m *meteringWorker) queryUsages(ctx context.Context) ([]*query.Usage, query.ExecCount, error) {
	var engine query.QueryEngine
	var err error
	switch m.name {
	case WorkerNameStorageMediumOneDay:
		engine, err = query.NewIntermediateQueryEngine(ctx, m.storage, m.finishedTs, m.reportIntervalHour)
	default:
		engine, err = query.NewDuckDBQueryEngine(ctx, m.storage, m.finishedTs, m.reportIntervalHour)
	}
	if err != nil {
		return nil, nil, fmt.Errorf("failed to build query engine: %s", err)
	}
	defer engine.Close()

	fns := make([]func(ctx context.Context) ([]*query.Usage, error), 0, 2)
	switch m.name {
	case WorkerNameDefault:
		fns = append(fns, engine.CalLoadDataRu)
		fns = append(fns, engine.CalPublicEgress)
		fns = append(fns, engine.CalPrivatelinkEgress)
		fns = append(fns, engine.CalTidbRU)
		fns = append(fns, engine.CalTiflashRU)
		fns = append(fns, engine.CalTidbRUV2)
		fns = append(fns, engine.CalTiflashRUV2)
		fns = append(fns, engine.CalTikvStorageLarge)
		fns = append(fns, engine.CalTiflashStorageLarge)
		fns = append(fns, engine.CalIntermediateTikvStorageMedium)
		fns = append(fns, engine.CalIntermediateTiflashStorageMedium)
		fns = append(fns, engine.CalTikvStorageIALarge)
		fns = append(fns, engine.CalIntermediateTikvStorageIAMedium)
		fns = append(fns, engine.CalRCU)
	case WorkerNameStorageMedium:
		fns = append(fns, engine.CalTikvStorageMedium)
		fns = append(fns, engine.CalTiflashStorageMedium)
	case WorkerNameStorageMediumOneDay:
		fns = append(fns, engine.CalTikvStorageMedium)
		fns = append(fns, engine.CalTiflashStorageMedium)
		fns = append(fns, engine.CalTikvStorageIAMedium)
	default:
		return nil, nil, fmt.Errorf("unknown worker name: %s", m.name)
	}

	result := make([]*query.Usage, 0)
	for i, fn := range fns {
		usages, err := fn(ctx)
		if err != nil {
			return nil, nil, fmt.Errorf("failed to calculate: index %d, err %s", i, err)
		}
		result = append(result, usages...)
	}

	if len(m.dbRuTenantIDs) > 0 && m.name == WorkerNameDefault {
		execCount, err := engine.CalExecCount(ctx)
		if err != nil {
			return nil, nil, fmt.Errorf("failed to calculate exec count: %s", err)
		}
		return result, execCount, nil
	}

	return result, nil, nil
}

func (m *meteringWorker) applyRUVersionLabels(report *BillingReport, cluster *repository.ClusterBranch, usage *query.Usage) {
	if !usage.MetricType.IsRUV2() {
		return
	}
	report.Labels[BillingLabelRUVersion] = cluster.RUVersion
}

func (m *meteringWorker) NextTs() int64 {
	return util.LocateLatestOnTheHour(m.finishedTs + int64(60*60*m.reportIntervalHour))
}

func (m *meteringWorker) buildClusterLabels(clusterBranch *repository.ClusterBranch, usage *query.Usage) map[string]string {
	labels := map[string]string{
		BillingLabelTenantID:   fmt.Sprint(clusterBranch.TenantID),
		BillingLabelProjectID:  fmt.Sprint(clusterBranch.ProjectID),
		BillingLabelClusterID:  fmt.Sprint(clusterBranch.ClusterID),
		BillingLabelRegion:     m.region,
		BillingLabelServerless: "true",
		BillingLabelWorker:     m.name,
		BillingLabelKeyspaceId: fmt.Sprint(clusterBranch.KeyspaceID),
	}
	if usage.IsTiflashRU() {
		labels[BillingLabelComponent] = BillingComponentTiFlash
	}
	if clusterBranch.TenantID == 1372813089197511287 {
		// Bai's org, used for internal test
		labels[BillingLabelBetaKey] = BillingLabelBetaEnable
		labels[BillingLabelBetaCategoryKey] = BillingBetaCategoryPoc
	}
	return labels
}

func (m *meteringWorker) buildBranchLabels(clusterBranch *repository.ClusterBranch, usage *query.Usage) map[string]string {
	labels := map[string]string{
		BillingLabelTenantID:   fmt.Sprint(clusterBranch.TenantID),
		BillingLabelProjectID:  fmt.Sprint(clusterBranch.ProjectID),
		BillingLabelClusterID:  fmt.Sprint(clusterBranch.ClusterID),
		BillingLabelBranchID:   clusterBranch.PrimaryID,
		BillingLabelBranchName: clusterBranch.Name,
		BillingLabelRegion:     m.region,
		BillingLabelServerless: "true",
		/* used for beta phase, removal upon official release */
		BillingLabelBetaKey:         BillingLabelBetaEnable,
		BillingLabelBetaCategoryKey: BillingBetaCategoryBranching,
		BillingLabelWorker:          m.name,
		BillingLabelKeyspaceId:      fmt.Sprint(clusterBranch.KeyspaceID),
	}
	if usage.IsTiflashRU() {
		labels[BillingLabelComponent] = BillingComponentTiFlash
	}
	return labels
}

func (m *meteringWorker) buildLabels(clusterBranch *repository.ClusterBranch, usage *query.Usage) map[string]string {
	switch clusterBranch.Type {
	case repository.CBTypeCluster:
		return m.buildClusterLabels(clusterBranch, usage)
	case repository.CBTypeBranch:
		return m.buildBranchLabels(clusterBranch, usage)
	default:
		return map[string]string{}
	}
}

func (m *meteringWorker) splitCost(cost int64, percentMap map[string]int64) map[string]int64 {
	result := make(map[string]int64)
	if len(percentMap) == 0 {
		result[""] = cost
		return result
	}

	var totalPercent int64 = 0
	keys := make([]string, 0, len(percentMap))
	for k, v := range percentMap {
		keys = append(keys, k)
		totalPercent += v
	}

	// Sort keys to make results deterministic
	sort.Strings(keys)

	var total int64 = 0
	// Calculate cost for all but the last
	for i := 0; i < len(keys)-1; i++ {
		key := keys[i]
		percent := percentMap[key]
		splitAmount := cost * percent / totalPercent
		result[key] = splitAmount
		total += splitAmount
	}

	// Last one gets the rest
	lastKey := keys[len(keys)-1]
	result[lastKey] = cost - total

	return result
}