/*
 *
 * 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 resolver

import (
	"encoding/json"
	"fmt"
	"math/bits"
	rand "math/rand/v2"
	"strings"
	"sync"
	"time"

	xxhash "github.com/cespare/xxhash/v2"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/internal/grpcsync"
	"google.golang.org/grpc/internal/grpcutil"
	iresolver "google.golang.org/grpc/internal/resolver"
	iringhash "google.golang.org/grpc/internal/ringhash"
	"google.golang.org/grpc/internal/serviceconfig"
	"google.golang.org/grpc/internal/wrr"
	"google.golang.org/grpc/internal/xds/balancer/clusterimpl"
	"google.golang.org/grpc/internal/xds/balancer/clustermanager"
	"google.golang.org/grpc/internal/xds/httpfilter"
	"google.golang.org/grpc/internal/xds/xdsclient/xdsresource"
	"google.golang.org/grpc/metadata"
	"google.golang.org/grpc/status"
)

const (
	cdsName                      = "cds_experimental"
	xdsClusterManagerName        = "xds_cluster_manager_experimental"
	clusterPrefix                = "cluster:"
	clusterSpecifierPluginPrefix = "cluster_specifier_plugin:"
)

type serviceConfig struct {
	LoadBalancingConfig balancerConfig `json:"loadBalancingConfig"`
}

type balancerConfig []map[string]any

func newBalancerConfig(name string, config any) balancerConfig {
	return []map[string]any{{name: config}}
}

type cdsBalancerConfig struct {
	Cluster string `json:"cluster"`
}

type xdsChildConfig struct {
	ChildPolicy balancerConfig `json:"childPolicy"`
}

type xdsClusterManagerConfig struct {
	Children map[string]xdsChildConfig `json:"children"`
}

// serviceConfigJSON produces a service config in JSON format that contains LB
// policy config for the "xds_cluster_manager" LB policy, with entries in the
// children map for all active clusters.
func serviceConfigJSON(activeClusters map[string]*clusterInfo, activePlugins map[string]*clusterInfo) []byte {
	// Generate children (all entries in activeClusters).
	children := make(map[string]xdsChildConfig)
	for cluster, ci := range activeClusters {
		children[cluster] = ci.cfg
	}
	for plugin, ci := range activePlugins {
		children[plugin] = ci.cfg
	}

	sc := serviceConfig{
		LoadBalancingConfig: newBalancerConfig(
			xdsClusterManagerName, xdsClusterManagerConfig{Children: children},
		),
	}

	// This is not expected to fail as we have constructed the service config by
	// hand right above, and therefore ok to panic.
	bs, err := json.Marshal(sc)
	if err != nil {
		panic(fmt.Sprintf("failed to marshal service config %+v: %v", sc, err))
	}
	return bs
}

type virtualHost struct {
	// retry policy present in virtual host
	retryConfig *xdsresource.RetryConfig
}

// routeCluster holds information about a cluster as referenced by a route.
type routeCluster struct {
	name        string                       // Name of the cluster.
	interceptor httpfilter.ClientInterceptor // HTTP filters to run for RPCs matching this route.
}

type route struct {
	m                 *xdsresource.CompositeMatcher         // converted from route matchers
	actionType        xdsresource.RouteActionType           // holds route action type
	clusters          wrr.WRR                               // holds *routeCluster entries
	routeClusters     []*grpcsync.RefCounted[*routeCluster] // Route clusters belonging to this route
	maxStreamDuration time.Duration
	retryConfig       *xdsresource.RetryConfig
	hashPolicies      []*xdsresource.HashPolicy
	autoHostRewrite   bool
}

func (r route) String() string {
	return fmt.Sprintf("%s -> { clusters: %v, maxStreamDuration: %v }", r.m.String(), r.clusters, r.maxStreamDuration)
}

// stoppableConfigSelector extends the iresolver.ConfigSelector interface with a
// stop() method. This makes it possible to swap the current config selector
// with an erroring config selector when the LDS or RDS resource is not found on
// the management server.
type stoppableConfigSelector interface {
	iresolver.ConfigSelector
	stop()
}

// erroringConfigSelector always returns an error, with the xDS node ID included
// in the error message. It is used to swap out the current config selector
// when the LDS or RDS resource is not found on the management server.
type erroringConfigSelector struct {
	err error
}

func newErroringConfigSelector(err error, xdsNodeID string) *erroringConfigSelector {
	return &erroringConfigSelector{err: annotateErrorWithNodeID(status.Error(codes.Unavailable, err.Error()), xdsNodeID)}
}

func (cs *erroringConfigSelector) SelectConfig(iresolver.RPCInfo) (*iresolver.RPCConfig, error) {
	return nil, cs.err
}
func (cs *erroringConfigSelector) stop() {}

type configSelector struct {
	channelID            uint64 // Static hash when hash policy is HashPolicyTypeChannelID
	xdsNodeID            string // xDS node ID, for annotating errors.
	sendNewServiceConfig func() // Function to send a new service config to gRPC.

	// Configuration received from the xDS management server.
	virtualHost      virtualHost
	routes           []route
	clusters         map[string]*clusterInfo
	plugins          map[string]*clusterInfo
	httpFilterConfig []xdsresource.HTTPFilter
	xdsConfig        *xdsresource.XDSConfig
}

var errNoMatchedRouteFound = status.Errorf(codes.Unavailable, "no matched route was found")
var errUnsupportedClientRouteAction = status.Errorf(codes.Unavailable, "matched route does not have a supported route action type")

// annotateErrorWithNodeID annotates the given error with the provided xDS node
// ID. This is used by the real config selector when it runs into errors, and
// also by the erroring config selector.
func annotateErrorWithNodeID(err error, nodeID string) error {
	return fmt.Errorf("[xDS node id: %s]: %w", nodeID, err)
}

func (cs *configSelector) SelectConfig(rpcInfo iresolver.RPCInfo) (*iresolver.RPCConfig, error) {
	var rt *route
	md, _ := metadata.FromOutgoingContext(rpcInfo.Context)
	if extraMD, ok := grpcutil.ExtraMetadata(rpcInfo.Context); ok {
		md = metadata.Join(md, extraMD)
		// Remove all binary headers. They are hard to match with. May need
		// to add back if asked by users.
		for k := range md {
			if strings.HasSuffix(k, "-bin") {
				delete(md, k)
			}
		}
	}
	// Loop through routes in order and select first match.
	for _, r := range cs.routes {
		if r.m.Match(rpcInfo.Method, md) {
			rt = &r
			break
		}
	}

	if rt == nil || rt.clusters == nil {
		return nil, annotateErrorWithNodeID(errNoMatchedRouteFound, cs.xdsNodeID)
	}

	if rt.actionType != xdsresource.RouteActionRoute {
		return nil, annotateErrorWithNodeID(errUnsupportedClientRouteAction, cs.xdsNodeID)
	}

	rc, ok := rt.clusters.Next().(*grpcsync.RefCounted[*routeCluster])
	if !ok {
		return nil, annotateErrorWithNodeID(status.Errorf(codes.Internal, "error retrieving cluster for match: %v (%T)", rc, rc), cs.xdsNodeID)
	}
	cluster := rc.Value()
	lbCtx := clustermanager.SetPickedCluster(rpcInfo.Context, cluster.name)
	lbCtx = xdsresource.NewContextWithXDSConfig(lbCtx, cs.xdsConfig)
	lbCtx = iringhash.SetXDSRequestHash(lbCtx, cs.generateHash(rpcInfo, rt.hashPolicies))
	if rt.autoHostRewrite {
		lbCtx = clusterimpl.EnableAutoHostRewrite(lbCtx)
	}

	config := &iresolver.RPCConfig{
		Context:     lbCtx,
		Interceptor: cluster.interceptor,
	}
	// Add a ref to the selected cluster to keep the interceptors alive until RPC
	// is committed.
	rc.Increment()
	if info, ok := cs.clusters[cluster.name]; ok {
		// Add a ref to the selected cluster, as this RPC needs this
		// cluster until it is committed.
		info.refCount.Add(1)
		config.OnCommitted = sync.OnceFunc(func() {
			if v := info.refCount.Add(-1); v == 0 {
				// We call unsubscribe rather than sendNewServiceConfig to
				// prevent redundant updates. If the reference count in the
				// dependency manager drops to zero, it will automatically
				// trigger a service config update with this cluster
				// removed. Calling unsubscribe allows the dependency
				// manager to handle the update flow once and for all.
				info.unsubscribe()
			}
			// Decrement the refcount of the route cluster and close the interceptor
			// if refcount goes to zero.
			rc.Decrement()
		})
	} else if info, ok := cs.plugins[cluster.name]; ok {
		// Add a ref to the selected plugin, as this RPC needs this
		// plugin until it is committed.
		info.refCount.Add(1)
		config.OnCommitted = sync.OnceFunc(func() {
			if v := info.refCount.Add(-1); v == 0 {
				// This entry will be removed from activePlugins when
				// producing a new service config update.
				cs.sendNewServiceConfig()
			}
			// Decrement the refcount of the route cluster and close the interceptor
			// if refcount goes to zero.
			rc.Decrement()
		})
	} else {
		// This should be unreachable because all route clusters are normalized
		// into cs.clusters or cs.plugins during config selector creation.
		panic(fmt.Sprintf("matched cluster %q not found in ConfigSelector", cluster.name))
	}

	if rt.maxStreamDuration != 0 {
		config.MethodConfig.Timeout = &rt.maxStreamDuration
	}
	if rt.retryConfig != nil {
		config.MethodConfig.RetryPolicy = retryConfigToPolicy(rt.retryConfig)
	} else if cs.virtualHost.retryConfig != nil {
		config.MethodConfig.RetryPolicy = retryConfigToPolicy(cs.virtualHost.retryConfig)
	}

	return config, nil
}

func retryConfigToPolicy(config *xdsresource.RetryConfig) *serviceconfig.RetryPolicy {
	return &serviceconfig.RetryPolicy{
		MaxAttempts:          int(config.NumRetries) + 1,
		InitialBackoff:       config.RetryBackoff.BaseInterval,
		MaxBackoff:           config.RetryBackoff.MaxInterval,
		BackoffMultiplier:    2,
		RetryableStatusCodes: config.RetryOn,
	}
}

func (cs *configSelector) generateHash(rpcInfo iresolver.RPCInfo, hashPolicies []*xdsresource.HashPolicy) uint64 {
	var hash uint64
	var generatedHash bool
	var md, emd metadata.MD
	var mdRead bool
	for _, policy := range hashPolicies {
		var policyHash uint64
		var generatedPolicyHash bool
		switch policy.HashPolicyType {
		case xdsresource.HashPolicyTypeHeader:
			if strings.HasSuffix(policy.HeaderName, "-bin") {
				continue
			}
			if !mdRead {
				md, _ = metadata.FromOutgoingContext(rpcInfo.Context)
				emd, _ = grpcutil.ExtraMetadata(rpcInfo.Context)
				mdRead = true
			}
			values := emd.Get(policy.HeaderName)
			if len(values) == 0 {
				// Extra metadata (e.g. the "content-type" header) takes
				// precedence over the user's metadata.
				values = md.Get(policy.HeaderName)
				if len(values) == 0 {
					// If the header isn't present at all, this policy is a no-op.
					continue
				}
			}
			joinedValues := strings.Join(values, ",")
			if policy.Regex != nil {
				joinedValues = policy.Regex.ReplaceAllString(joinedValues, policy.RegexSubstitution)
			}
			policyHash = xxhash.Sum64String(joinedValues)
			generatedHash = true
			generatedPolicyHash = true
		case xdsresource.HashPolicyTypeChannelID:
			// Use the static channel ID as the hash for this policy.
			policyHash = cs.channelID
			generatedHash = true
			generatedPolicyHash = true
		}

		// Deterministically combine the hash policies. Rotating prevents
		// duplicate hash policies from cancelling each other out and preserves
		// the 64 bits of entropy.
		if generatedPolicyHash {
			hash = bits.RotateLeft64(hash, 1)
			hash = hash ^ policyHash
		}

		// If terminal policy and a hash has already been generated, ignore the
		// rest of the policies and use that hash already generated.
		if policy.Terminal && generatedHash {
			break
		}
	}

	if generatedHash {
		return hash
	}
	// If no generated hash return a random long. In the grand scheme of things
	// this logically will map to choosing a random backend to route request to.
	return rand.Uint64()
}

// stop decrements refs of all clusters referenced by this config selector.
func (cs *configSelector) stop() {
	// The resolver's old configSelector may be nil.  Handle that here.
	if cs == nil {
		return
	}

	// Decrement the refcount of all the route clusters associated with this
	// config selector and close the interceptors of the route cluster if it's
	// refcount goes to zero.
	for _, r := range cs.routes {
		for _, rc := range r.routeClusters {
			rc.Decrement()
		}
	}

	// If any reference counts drop to zero, a service config update is required
	// to remove the clusters. Since the old config selector is stopped
	// after a new one is active, we must trigger a subsequent update to delete
	// the now-unused clusters.
	for _, ci := range cs.clusters {
		if v := ci.refCount.Add(-1); v == 0 {
			ci.unsubscribe()
		}
	}
	for _, ci := range cs.plugins {
		if v := ci.refCount.Add(-1); v == 0 {
			cs.sendNewServiceConfig()
		}
	}
}