// Copyright 2025 PingCAP, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package execdetails

import (
	"context"
	"sync/atomic"

	"github.com/tikv/client-go/v2/util"
)

// ContextWithMissingExecDetailsInitialized initializes any missing statement execution, execution,
// and resource usage details in the context while preserving existing objects.
func ContextWithMissingExecDetailsInitialized(ctx context.Context) context.Context {
	if ctx.Value(StmtExecDetailKey) == nil {
		ctx = context.WithValue(ctx, StmtExecDetailKey, &StmtExecDetails{})
	}
	if ctx.Value(util.ExecDetailsKey) == nil {
		ctx = context.WithValue(ctx, util.ExecDetailsKey, &util.ExecDetails{})
	}
	if ctx.Value(util.RUDetailsCtxKey) == nil {
		ctx = context.WithValue(ctx, util.RUDetailsCtxKey, util.NewRUDetails())
	}
	if ctx.Value(RUV2MetricsCtxKey) == nil {
		ctx = context.WithValue(ctx, RUV2MetricsCtxKey, NewRUV2Metrics())
	}
	return ctx
}

// ContextWithInitializedExecDetails initializes any missing statement execution, execution,
// and resource usage details in the context while preserving existing objects.
func ContextWithInitializedExecDetails(ctx context.Context) context.Context {
	return ContextWithMissingExecDetailsInitialized(ctx)
}

// ContextWithInheritedRUV2Details reuses statement-level RUDetails and RUv2 metrics
// from source when ctx does not already carry them.
func ContextWithInheritedRUV2Details(ctx, source context.Context) context.Context {
	if source == nil {
		return ctx
	}
	if ctx.Value(util.RUDetailsCtxKey) == nil {
		if ruDetails, _ := source.Value(util.RUDetailsCtxKey).(*util.RUDetails); ruDetails != nil {
			ctx = context.WithValue(ctx, util.RUDetailsCtxKey, ruDetails)
		}
	}
	if ctx.Value(RUV2MetricsCtxKey) == nil {
		if metrics := RUV2MetricsFromContext(source); metrics != nil {
			ctx = context.WithValue(ctx, RUV2MetricsCtxKey, metrics)
		}
	}
	return ctx
}

// SyncRUV2MetricsFromContext drains any raw RUv2 counters in ctx's RUDetails into
// the statement-level RUv2 metrics and returns the metrics object.
func SyncRUV2MetricsFromContext(ctx context.Context) *RUV2Metrics {
	metrics := RUV2MetricsFromContext(ctx)
	if metrics == nil {
		return nil
	}
	ruDetails, _ := ctx.Value(util.RUDetailsCtxKey).(*util.RUDetails)
	SyncRUV2MetricsFromRUDetails(metrics, ruDetails)
	return metrics
}

// GetExecDetailsFromContext gets stmt execution, execution and resource usage details from context.
func GetExecDetailsFromContext(ctx context.Context) (stmtDetail StmtExecDetails, tikvExecDetail util.ExecDetails, ruDetails *util.RUDetails) {
	stmtDetailRaw := ctx.Value(StmtExecDetailKey)
	if stmtDetailRaw != nil {
		stmtDetail = *(stmtDetailRaw.(*StmtExecDetails))
	}
	tikvExecDetailRaw := ctx.Value(util.ExecDetailsKey)
	if tikvExecDetailRaw != nil {
		tikvExecDetail = LoadTiKVExecDetails(tikvExecDetailRaw.(*util.ExecDetails))
	}
	if ruDetailsVal := ctx.Value(util.RUDetailsCtxKey); ruDetailsVal != nil {
		ruDetails = ruDetailsVal.(*util.RUDetails)
	} else {
		ruDetails = util.NewRUDetails()
	}

	return
}

// LoadTiKVExecDetails snapshots the fields in util.ExecDetails that are updated via atomic operations.
func LoadTiKVExecDetails(detail *util.ExecDetails) util.ExecDetails {
	if detail == nil {
		return util.ExecDetails{}
	}
	return util.ExecDetails{
		BackoffCount:       atomic.LoadInt64(&detail.BackoffCount),
		BackoffDuration:    atomic.LoadInt64(&detail.BackoffDuration),
		WaitKVRespDuration: atomic.LoadInt64(&detail.WaitKVRespDuration),
		WaitPDRespDuration: atomic.LoadInt64(&detail.WaitPDRespDuration),
	}
}