/*
 * Copyright 2024 CloudWeGo 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 retry

import (
	"context"
	"errors"
	"fmt"
	"sync"
	"sync/atomic"
	"time"

	"github.com/cloudwego/kitex/pkg/circuitbreak"
	"github.com/cloudwego/kitex/pkg/gofunc"
	"github.com/cloudwego/kitex/pkg/klog"
	"github.com/cloudwego/kitex/pkg/rpcinfo"
	"github.com/cloudwego/kitex/pkg/utils"
)

func newMixedRetryer(policy Policy, r *ShouldResultRetry, cbC *cbContainer) (Retryer, error) {
	mr := &mixedRetryer{failureCommon: &failureCommon{specifiedResultRetry: r, cbContainer: cbC}}
	if err := mr.UpdatePolicy(policy); err != nil {
		return nil, fmt.Errorf("newMixedRetryer failed, err=%w", err)
	}
	return mr, nil
}

type mixedRetryer struct {
	enable bool
	*failureCommon
	policy     *MixedPolicy
	retryDelay time.Duration
	sync.RWMutex
	errMsg string
}

// ShouldRetry to check if retry request can be called, it is checked in retryer.Do.
// If not satisfy will return the reason message
// Actually, the ShouldRetry logic is same with failureRetryer, because
func (r *mixedRetryer) ShouldRetry(ctx context.Context, err error, callTimes int, req interface{}, cbKey string) (string, bool) {
	r.RLock()
	defer r.RUnlock()
	if !r.enable {
		return "", false
	}
	return r.shouldRetry(ctx, callTimes, req, cbKey, &r.policy.FailurePolicy)
}

// AllowRetry implements the Retryer interface.
func (r *mixedRetryer) AllowRetry(ctx context.Context) (string, bool) {
	r.RLock()
	defer r.RUnlock()
	if !r.enable || r.policy.StopPolicy.MaxRetryTimes == 0 {
		return "", false
	}
	if stop, msg := chainStop(ctx, r.policy.StopPolicy); stop {
		return msg, false
	}
	if stop, msg := ddlStop(ctx, r.policy.StopPolicy); stop {
		return msg, false
	}
	return "", true
}

// Do implement the Retryer interface.
func (r *mixedRetryer) Do(ctx context.Context, rpcCall RPCCallFunc, firstRI rpcinfo.RPCInfo, req, resp interface{}) (lastRI rpcinfo.RPCInfo, recycleRI bool, err error) {
	r.RLock()
	var maxDuration time.Duration
	if r.policy.StopPolicy.MaxDurationMS > 0 {
		maxDuration = time.Duration(r.policy.StopPolicy.MaxDurationMS) * time.Millisecond
	}
	retryTimes := r.policy.StopPolicy.MaxRetryTimes
	retryDelay := r.retryDelay
	r.RUnlock()

	var callTimes int32
	var callCosts utils.StringBuilder
	callCosts.RawStringBuilder().Grow(32)

	var recordCostDoing int32 = 0
	var abort int32 = 0
	doneCount := 0

	// notice: buff num of chan is very important here, it cannot less than call times, or the below chan receive will block
	callDone := make(chan resultWrapper, retryTimes+1)
	timer := time.NewTimer(retryDelay)
	cbKey, _ := r.cbContainer.cbCtl.GetKey(ctx, req)
	defer func() {
		if panicInfo := recover(); panicInfo != nil {
			err = panicToErr(ctx, panicInfo, firstRI)
		}
		timer.Stop()
	}()
	startTime := time.Now()
	newRespFunc := getNewRespFunc(firstRI.Invocation().MethodInfo())
	// include first call, max loop is retryTimes + 1
	doCall := true
	for callCount := 0; ; {
		if doCall {
			doCall = false
			if callCount > 0 {
				if ret, e := isExceedMaxDuration(ctx, startTime, maxDuration, atomic.LoadInt32(&callTimes)); ret {
					return firstRI, false, e
				}
			}
			callCount++
			gofunc.GoFunc(ctx, func() {
				if atomic.LoadInt32(&abort) == 1 {
					return
				}
				var (
					e       error
					cRI     rpcinfo.RPCInfo
					curResp = newRespFunc()
				)
				defer func() {
					if panicInfo := recover(); panicInfo != nil {
						e = panicToErr(ctx, panicInfo, firstRI)
					}
					callDone <- resultWrapper{cRI, curResp, e}
				}()
				ct := atomic.AddInt32(&callTimes, 1)
				callStart := time.Now()
				if r.cbContainer.enablePercentageLimit {
					// record stat before call since requests may be slow, making the limiter more accurate
					recordRetryStat(cbKey, r.cbContainer.cbPanel, ct)
				}
				cRI, e = rpcCall(ctx, r, req, curResp)
				recordCost(ct, callStart, &recordCostDoing, &callCosts, &abort, e)
				if !r.cbContainer.enablePercentageLimit && r.cbContainer.cbStat {
					circuitbreak.RecordStat(ctx, req, nil, e, cbKey, r.cbContainer.cbCtl, r.cbContainer.cbPanel)
				}
			})
		}
		select {
		case <-timer.C:
			// backup retry
			if _, ok := r.ShouldRetry(ctx, nil, callCount, req, cbKey); ok && callCount < retryTimes+1 {
				doCall = true
				timer.Reset(retryDelay)
			}
		case res := <-callDone:
			// result retry
			doneCount++
			if doneCount < retryTimes+1 {
				if callCount < retryTimes+1 {
					if msg, ok := r.ShouldRetry(ctx, nil, callCount, req, cbKey); ok {
						if r.isRetryResult(ctx, res.ri, res.resp, res.err, &r.policy.FailurePolicy) {
							doCall = true
							timer.Reset(retryDelay)
							continue
						}
					} else if msg != "" {
						appendMsg := fmt.Sprintf("retried %d, %s", callCount-1, msg)
						appendErrMsg(res.err, appendMsg)
					}
				} else if r.isRetryResult(ctx, res.ri, res.resp, res.err, &r.policy.FailurePolicy) {
					continue
				}
			}
			atomic.StoreInt32(&abort, 1)
			recordRetryInfo(res.ri, atomic.LoadInt32(&callTimes), callCosts.String())
			shallowCopyResults(res.resp, resp)
			return res.ri, false, res.err
		}
	}
}

// UpdatePolicy implements the Retryer interface.
func (r *mixedRetryer) UpdatePolicy(rp Policy) (err error) {
	r.Lock()
	defer r.Unlock()
	if !rp.Enable {
		r.enable = rp.Enable
		return nil
	}
	if rp.MixedPolicy == nil || rp.Type != MixedType {
		err = fmt.Errorf("MixedPolicy is nil or retry type not match(type=%v), cannot do update in mixedRetryer", rp.Type)
		r.errMsg = err.Error()
		return err
	}
	if rp.MixedPolicy.RetryDelayMS == 0 {
		err = errors.New("invalid retry delay duration in mixedRetryer")
		r.errMsg = err.Error()
		return err
	}
	if err = checkStopPolicy(&rp.MixedPolicy.StopPolicy, maxMixRetryTimes, r); err != nil {
		r.errMsg = err.Error()
		return err
	}
	r.enable = rp.Enable
	r.policy = rp.MixedPolicy
	r.retryDelay = time.Duration(rp.MixedPolicy.RetryDelayMS) * time.Millisecond
	r.setSpecifiedResultRetryIfNeeded(r.specifiedResultRetry, &r.policy.FailurePolicy)
	if bo, e := initBackOff(rp.MixedPolicy.BackOffPolicy); e != nil {
		r.errMsg = fmt.Sprintf("mixedRetryer update BackOffPolicy failed, err=%s", e.Error())
		klog.Warnf("KITEX: %s", r.errMsg)
	} else {
		r.backOff = bo
	}
	return nil
}

// AppendErrMsgIfNeeded implements the Retryer interface.
func (r *mixedRetryer) AppendErrMsgIfNeeded(ctx context.Context, err error, ri rpcinfo.RPCInfo, msg string) {
	if r.isRetryErr(ctx, err, ri, &r.policy.FailurePolicy) {
		// Add additional reason when retry is not applied.
		appendErrMsg(err, msg)
	}
}

// Prepare implements the Retryer interface.
func (r *mixedRetryer) Prepare(ctx context.Context, prevRI, retryRI rpcinfo.RPCInfo) {
	handleRetryInstance(r.policy.RetrySameNode, prevRI, retryRI)
}

// Type implements the Retryer interface.
func (r *mixedRetryer) Type() Type {
	return MixedType
}

// Dump implements the Retryer interface.
func (r *mixedRetryer) Dump() map[string]interface{} {
	r.RLock()
	defer r.RUnlock()
	dm := make(map[string]interface{})
	dm["enable"] = r.enable
	dm["mixed_retry"] = r.policy
	if r.policy != nil {
		dm["specified_result_retry"] = r.dumpSpecifiedResultRetry(r.policy.FailurePolicy)
	}
	if r.errMsg != "" {
		dm["err_msg"] = r.errMsg
	}
	return dm
}
