package utils

import (
	"context"
	"sync"

	"github.com/tidbcloud/tiinf/monitor/tilog"
	"google.golang.org/grpc"
)

type ApiCallRecorder struct {
	methodCalledTimes map[string]int
	lock              sync.Mutex
}

func NewApiCallRecorder() *ApiCallRecorder {
	return &ApiCallRecorder{
		methodCalledTimes: map[string]int{},
	}
}

func (recorder *ApiCallRecorder) GetMethodCalledTimesAndClear(method string) int {
	recorder.lock.Lock()
	defer recorder.lock.Unlock()
	times, ok := recorder.methodCalledTimes[method]
	if !ok {
		return 0
	}
	recorder.methodCalledTimes[method] = 0
	return times
}

func (recorder *ApiCallRecorder) GetAllCalledTimesAndClear() int {
	recorder.lock.Lock()
	defer recorder.lock.Unlock()
	all := 0
	for method, times := range recorder.methodCalledTimes {
		all += times
		recorder.methodCalledTimes[method] = 0
	}
	return all
}

func (recorder *ApiCallRecorder) UnaryServerInterceptor() grpc.UnaryServerInterceptor {
	return func(ctx context.Context, req any, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (any, error) {
		method := info.FullMethod
		resp, err := handler(ctx, req)
		recorder.lock.Lock()
		defer recorder.lock.Unlock()
		times, ok := recorder.methodCalledTimes[method]
		if !ok {
			recorder.methodCalledTimes[method] = 1
		} else {
			recorder.methodCalledTimes[method] = times + 1
		}
		tilog.Info(ctx, method)
		return resp, err
	}
}