/*
 *
 * Copyright 2022 gRPC authors.
 *
 * 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 orca

import (
	"context"
	"sync"
	"sync/atomic"

	v3orcapb "github.com/cncf/xds/go/xds/data/orca/v3"
)

// CallMetricRecorder provides functionality to record per-RPC custom backend
// metrics. See CallMetricsServerOption() for more details.
//
// Safe for concurrent use.
type CallMetricRecorder struct {
	cpu    atomic.Value // float64
	memory atomic.Value // float64

	mu          sync.RWMutex
	requestCost map[string]float64
	utilization map[string]float64
}

func newCallMetricRecorder() *CallMetricRecorder {
	return &CallMetricRecorder{
		requestCost: make(map[string]float64),
		utilization: make(map[string]float64),
	}
}

// SetCPUUtilization records a measurement for the CPU utilization metric.
func (c *CallMetricRecorder) SetCPUUtilization(val float64) {
	c.cpu.Store(val)
}

// SetMemoryUtilization records a measurement for the memory utilization metric.
func (c *CallMetricRecorder) SetMemoryUtilization(val float64) {
	c.memory.Store(val)
}

// SetRequestCost records a measurement for a request cost metric,
// uniquely identifiable by name.
func (c *CallMetricRecorder) SetRequestCost(name string, val float64) {
	c.mu.Lock()
	c.requestCost[name] = val
	c.mu.Unlock()
}

// SetUtilization records a measurement for a utilization metric uniquely
// identifiable by name.
func (c *CallMetricRecorder) SetUtilization(name string, val float64) {
	c.mu.Lock()
	c.utilization[name] = val
	c.mu.Unlock()
}

// toLoadReportProto dumps the recorded measurements as an OrcaLoadReport proto.
func (c *CallMetricRecorder) toLoadReportProto() *v3orcapb.OrcaLoadReport {
	c.mu.RLock()
	defer c.mu.RUnlock()

	cost := make(map[string]float64, len(c.requestCost))
	for k, v := range c.requestCost {
		cost[k] = v
	}
	util := make(map[string]float64, len(c.utilization))
	for k, v := range c.utilization {
		util[k] = v
	}
	cpu, _ := c.cpu.Load().(float64)
	mem, _ := c.memory.Load().(float64)
	return &v3orcapb.OrcaLoadReport{
		CpuUtilization: cpu,
		MemUtilization: mem,
		RequestCost:    cost,
		Utilization:    util,
	}
}

type callMetricRecorderCtxKey struct{}

// CallMetricRecorderFromContext returns the RPC specific custom metrics
// recorder [CallMetricRecorder] embedded in the provided RPC context.
//
// Returns nil if no custom metrics recorder is found in the provided context,
// which will be the case when custom metrics reporting is not enabled.
func CallMetricRecorderFromContext(ctx context.Context) *CallMetricRecorder {
	rw, ok := ctx.Value(callMetricRecorderCtxKey{}).(*recorderWrapper)
	if !ok {
		return nil
	}
	return rw.recorder()
}

func newContextWithRecorderWrapper(ctx context.Context, r *recorderWrapper) context.Context {
	return context.WithValue(ctx, callMetricRecorderCtxKey{}, r)
}

// recorderWrapper is a wrapper around a CallMetricRecorder to ensures that
// concurrent calls to CallMetricRecorderFromContext() results in only one
// allocation of the underlying metric recorder.
type recorderWrapper struct {
	once sync.Once
	r    *CallMetricRecorder
}

func (rw *recorderWrapper) recorder() *CallMetricRecorder {
	rw.once.Do(func() {
		rw.r = newCallMetricRecorder()
	})
	return rw.r
}