/*
 *
 * Copyright 2020 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 clusterimpl

import (
	"context"

	v3orcapb "github.com/cncf/xds/go/xds/data/orca/v3"
	"google.golang.org/grpc/balancer"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/connectivity"
	"google.golang.org/grpc/internal/envconfig"
	"google.golang.org/grpc/internal/stats"
	"google.golang.org/grpc/internal/wrr"
	xdsinternal "google.golang.org/grpc/internal/xds"
	"google.golang.org/grpc/internal/xds/clients"
	"google.golang.org/grpc/internal/xds/xdsclient"
	"google.golang.org/grpc/internal/xds/xdsclient/xdsresource"
	"google.golang.org/grpc/metadata"
	"google.golang.org/grpc/status"
)

// NewRandomWRR is used when calculating drops. It's exported so that tests can
// override it.
var NewRandomWRR = wrr.NewRandom

const million = 1000000

type dropper struct {
	category string
	w        wrr.WRR
}

// greatest common divisor (GCD) via Euclidean algorithm
func gcd(a, b uint32) uint32 {
	for b != 0 {
		t := b
		b = a % b
		a = t
	}
	return a
}

func newDropper(c DropConfig) *dropper {
	w := NewRandomWRR()
	gcdv := gcd(c.RequestsPerMillion, million)
	// Return true for RequestPerMillion, false for the rest.
	w.Add(true, int64(c.RequestsPerMillion/gcdv))
	w.Add(false, int64((million-c.RequestsPerMillion)/gcdv))

	return &dropper{
		category: c.Category,
		w:        w,
	}
}

func (d *dropper) drop() (ret bool) {
	return d.w.Next().(bool)
}

// loadReporter wraps the methods from the loadStore that are used here.
type loadReporter interface {
	CallStarted(locality clients.Locality)
	CallFinished(locality clients.Locality, err error)
	CallServerLoad(locality clients.Locality, name string, val float64)
	CallDropped(category string)
}

// Picker implements RPC drop, circuit breaking drop and load reporting.
type picker struct {
	drops           []*dropper
	s               balancer.State
	loadStore       loadReporter
	counter         *xdsclient.ClusterRequestsCounter
	countMax        uint32
	telemetryLabels map[string]string
	clusterName     string
	metrics         *xdsresource.LRSReportEndpointMetricsConfig
}

func (d *picker) Pick(info balancer.PickInfo) (balancer.PickResult, error) {
	// Unconditionally update labels if present, even dropped or queued RPC's can
	// use these labels.
	stats.UpdateLabels(info.Ctx, d.telemetryLabels)

	// Don't drop unless the inner picker is READY. Similar to
	// https://github.com/grpc/grpc-go/issues/2622.
	if d.s.ConnectivityState == connectivity.Ready {
		// Check if this RPC should be dropped by category.
		for _, dp := range d.drops {
			if dp.drop() {
				if d.loadStore != nil {
					d.loadStore.CallDropped(dp.category)
				}
				return balancer.PickResult{}, status.Errorf(codes.Unavailable, "RPC is dropped")
			}
		}
	}

	// Check if this RPC should be dropped by circuit breaking.
	if d.counter != nil {
		if err := d.counter.StartRequest(d.countMax); err != nil {
			// Drops by circuit breaking are reported with empty category. They
			// will be reported only in total drops, but not in per category.
			if d.loadStore != nil {
				d.loadStore.CallDropped("")
			}
			return balancer.PickResult{}, status.Error(codes.Unavailable, err.Error())
		}
	}

	var lID clients.Locality
	pr, err := d.s.Picker.Pick(info)
	if scw, ok := pr.SubConn.(*scWrapper); ok {
		// This OK check also covers the case err!=nil, because SubConn will be
		// nil.
		pr.SubConn = scw.SubConn
		// If locality ID isn't found in the wrapper, an empty locality ID will
		// be used.
		lID = scw.localityID

		if scw.hostname != "" && autoHostRewriteEnabled(info.Ctx) {
			if pr.Metadata == nil {
				pr.Metadata = metadata.Pairs(":authority", scw.hostname)
			} else {
				pr.Metadata.Set(":authority", scw.hostname)
			}
		}
	}

	if err != nil {
		if d.counter != nil {
			// Release one request count if this pick fails.
			d.counter.EndRequest()
		}
		return pr, err
	}

	stats.UpdateLabels(
		info.Ctx,
		map[string]string{
			"grpc.lb.locality":        xdsinternal.LocalityString(lID),
			"grpc.lb.backend_service": d.clusterName,
		},
	)

	if d.loadStore != nil {
		locality := clients.Locality{Region: lID.Region, Zone: lID.Zone, SubZone: lID.SubZone}
		d.loadStore.CallStarted(locality)
		oldDone := pr.Done
		pr.Done = func(info balancer.DoneInfo) {
			if oldDone != nil {
				oldDone(info)
			}
			d.loadStore.CallFinished(locality, info.Err)

			load, ok := info.ServerLoad.(*v3orcapb.OrcaLoadReport)
			if !ok || load == nil {
				return
			}

			if envconfig.XDSORCAToLRSPropEnabled {
				if d.metrics != nil {
					if d.metrics.CPUUtilization {
						d.loadStore.CallServerLoad(locality, "cpu_utilization", load.CpuUtilization)
					}
					if d.metrics.MemUtilization {
						d.loadStore.CallServerLoad(locality, "mem_utilization", load.MemUtilization)
					}
					if d.metrics.ApplicationUtilization {
						d.loadStore.CallServerLoad(locality, "application_utilization", load.ApplicationUtilization)
					}
					for n, c := range load.NamedMetrics {
						_, ok := d.metrics.NamedMetrics[n]
						if d.metrics.NamedMetricsAll || ok {
							d.loadStore.CallServerLoad(locality, "named_metrics."+n, c)
						}
					}
				}
			} else {
				for n, c := range load.NamedMetrics {
					d.loadStore.CallServerLoad(locality, n, c)
				}
			}
		}
	}

	if d.counter != nil {
		// Update Done() so that when the RPC finishes, the request count will
		// be released.
		oldDone := pr.Done
		pr.Done = func(doneInfo balancer.DoneInfo) {
			d.counter.EndRequest()
			if oldDone != nil {
				oldDone(doneInfo)
			}
		}
	}

	return pr, err
}

// autoHostRewriteKey is the context key used to store the value of
// route's autoHostRewrite in the RPC context.
type autoHostRewriteKey struct{}

// autoHostRewriteEnabled retrieves the autoHostRewrite value from the provided context.
func autoHostRewriteEnabled(ctx context.Context) bool {
	v, _ := ctx.Value(autoHostRewriteKey{}).(bool)
	return v
}

// AutoHostRewriteEnabledForTesting returns the value of autoHostRewrite field;
// to be used for testing only.
func AutoHostRewriteEnabledForTesting(ctx context.Context) bool {
	return autoHostRewriteEnabled(ctx)
}

// EnableAutoHostRewrite adds the autoHostRewrite value to the context for
// the xds_cluster_impl LB policy to pick.
func EnableAutoHostRewrite(ctx context.Context) context.Context {
	return context.WithValue(ctx, autoHostRewriteKey{}, true)
}