package controller

import (
	"context"
	"dataservices/commons/logger"
	"dataservices/commons/metrics"
	"dataservices/commons/utils"
	"fmt"
	"os"
	"sync"
	"sync/atomic"
	"time"

	"github.com/tidbcloud/tiinf/monitor/tilog"
	tigrpc "github.com/tidbcloud/tiinf/transport/tigrpc/client"
	"go.uber.org/zap"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/connectivity"
	"google.golang.org/grpc/keepalive"
	"google.golang.org/grpc/status"
)

const (
	defaultServiceConfig = `{
		"loadBalancingConfig": [{"round_robin":{}}],
		"methodConfig": [{
		  "name": [{"service": "dataapi.DataAPIExecutorService"}],
		  "retryPolicy": {
			  "maxAttempts": 4,
			  "initialBackoff": "0.1s",
			  "maxBackoff": "1s",
			  "backoffMultiplier": 2,
			  "retryableStatusCodes": [
				"UNAVAILABLE",
				"INTERNAL"
			  ]
		  }
		}]}`

	endpointServiceConfig = `{
		"methodConfig": [{
		  "name": [{"service": "dataapi.DataAPIExecutorService"}],
		  "retryPolicy": {
			  "maxAttempts": 2,
			  "initialBackoff": "0.1s",
			  "maxBackoff": "0.1s",
			  "backoffMultiplier": 1,
			  "retryableStatusCodes": [
				"UNAVAILABLE",
				"INTERNAL"
			  ]
		  }
		}]}`

	endpointMaxFailureTimes = 3
	endpointMaxIdleSec      = 300

	grpcConnectTimeout  = 1500 * time.Millisecond
	keepaliveInterval   = 10 * time.Second
	keepaliveTimeout    = 1500 * time.Millisecond
	grpcClienMaxMsgSize = 8 * 1024 * 1024 // default value: 8M
)

type Conn struct {
	conn         *grpc.ClientConn `json:"-"`
	addr         string
	createdAt    int64
	lastActiveAt int64
	failTimes    int64
}

func (c *Conn) String() string {
	return fmt.Sprintf("{addr:%s, createdAt:%d, lastActiveAt:%d, failTimes:%d}", c.addr, c.createdAt, c.lastActiveAt, c.failTimes)
}

type ExecutorClient struct {
	addr         string
	defaultConn  *grpc.ClientConn
	endpointPool sync.Map
	mutex        sync.Mutex
}

func newExecutorClient(addr string) (*ExecutorClient, error) {
	client := new(ExecutorClient)
	client.addr = addr

	// The default connection always targets ExecutorSvcAddr. In AliCloud V2 this
	// is expected to be the only executor instance; AWS can still override it per
	// request with endpointAddr when Redis-backed routing is enabled.
	conn, err := newGrpcConn(addr)
	if err != nil {
		return nil, err
	}
	client.defaultConn = conn
	client.runMonitor()
	return client, nil
}

func newGrpcConn(addr string) (*grpc.ClientConn, error) {
	return tigrpc.NewClientConn(addr,
		tigrpc.WithDefault(),
		tigrpc.WithGrpcOption(grpc.WithDefaultCallOptions(grpc.MaxCallRecvMsgSize(grpcClienMaxMsgSize))),
		tigrpc.WithGrpcOption(grpc.WithDefaultServiceConfig(defaultServiceConfig)))
}

type Option func(o *clientOption)

type clientOption struct {
	endpointAddr      string
	enableLoadBalance bool
}

func WithEndpointAddr(endpointAddr string) Option {
	return func(o *clientOption) {
		o.endpointAddr = endpointAddr
	}
}

func WithLoadBalance(enable bool) Option {
	return func(o *clientOption) {
		o.enableLoadBalance = enable
	}
}

func (c *ExecutorClient) Get() *grpc.ClientConn {
	return c.defaultConn
}

func (c *ExecutorClient) Exec(ctx context.Context, f func(ctx context.Context, conn *grpc.ClientConn) error, opts ...Option) (err error) {
	start := time.Now()
	defer func() {
		cost := time.Since(start)
		metrics.ExecutionLinkDuration.WithLabelValues("endpoint_exec", "run_handler").Observe(cost.Seconds())
		logger.LogWithError(ctx, err, "ExecutorClient.Exec completed", zap.Duration("cost", cost))
	}()
	var o clientOption
	for _, opt := range opts {
		opt(&o)
	}
	if o.endpointAddr != "" && (!o.enableLoadBalance || IsExecutorLoadOK(ctx, o.endpointAddr)) {
		var conn *Conn
		conn, err = c.getEndpointConn(ctx, o.endpointAddr)
		if err != nil {
			logger.Warn(ctx, "get executor service endpoint connection failed, retry with default connection", zap.String("endpointAddr", o.endpointAddr), zap.Error(err))
		} else {
			logger.Info(ctx, "get executor service endpoint connection succeed, begin execute user handler", zap.String("endpointAddr", o.endpointAddr))
			err = f(ctx, conn.conn)
			if !needRetry(err) {
				return err
			}
			c.incrConnFailures(ctx, conn)
			logger.Warn(ctx, "executor client exec failed, retry with default connection", zap.String("endpointAddr", o.endpointAddr), zap.Uint32("code", uint32(status.Code(err))), zap.Error(err))
		}
	}
	return f(ctx, c.defaultConn)
}

func (c *ExecutorClient) getEndpointConn(ctx context.Context, endpoint string) (*Conn, error) {
	start := time.Now()
	defer func() {
		cost := time.Since(start)
		metrics.ExecutionLinkDuration.WithLabelValues("endpoint_exec", "get_executor_endpoint").Observe(cost.Seconds())
	}()
	val, _ := c.endpointPool.Load(endpoint)
	endpointConn, ok := val.(*Conn)
	if ok {
		atomic.StoreInt64(&endpointConn.lastActiveAt, time.Now().Unix())
		return endpointConn, nil
	}

	c.mutex.Lock()
	defer c.mutex.Unlock()

	val, _ = c.endpointPool.Load(endpoint)
	endpointConn, ok = val.(*Conn)
	if ok {
		atomic.StoreInt64(&endpointConn.lastActiveAt, time.Now().Unix())
		return endpointConn, nil
	}

	ctx, cancel := context.WithTimeout(ctx, grpcConnectTimeout)
	defer cancel()
	kaParams := keepalive.ClientParameters{
		Time:                keepaliveInterval,
		Timeout:             keepaliveTimeout,
		PermitWithoutStream: true,
	}
	grpcConn, err := tigrpc.NewClientContext(ctx, endpoint,
		tigrpc.WithGrpcOption(grpc.WithBlock()),
		tigrpc.WithGrpcOption(grpc.WithDefaultCallOptions(grpc.MaxCallRecvMsgSize(grpcClienMaxMsgSize))),
		tigrpc.WithGrpcOption(grpc.WithKeepaliveParams(kaParams)),
		tigrpc.WithDefault(),
		tigrpc.WithGrpcOption(grpc.WithDefaultServiceConfig(endpointServiceConfig)))
	if err != nil {
		return nil, err
	}

	conn := &Conn{
		conn:         grpcConn,
		addr:         endpoint,
		createdAt:    time.Now().Unix(),
		lastActiveAt: time.Now().Unix(),
	}
	c.endpointPool.Store(endpoint, conn)
	tilog.Info(ctx, "new endpoint connection", zap.String("connection", conn.String()))
	return conn, nil
}

func (c *ExecutorClient) incrConnFailures(ctx context.Context, conn *Conn) {
	if atomic.AddInt64(&conn.failTimes, 1) < endpointMaxFailureTimes {
		return
	}
	c.endpointPool.Delete(conn.addr)
	conn.conn.Close()
	tilog.Info(ctx, "endpoint connection was cleaned up due to exceeding maximum failure times", zap.String("connection", conn.String()), zap.Int("endpointMaxFailureTimes", endpointMaxFailureTimes))
}

func (c *ExecutorClient) runMonitor() {
	utils.GO(func() {
		ctx := context.Background()
		podIP := os.Getenv("POD_IP")
		logger := tilog.WithContextL(ctx).With(zap.String("pod_ip", podIP))
		ticker := time.NewTicker(time.Second * 10)
		var times int64
		for range ticker.C {
			times++
			var size int
			logger.Info("[rpc_client_pool_monitor] endpoint connection monitor triggered", zap.Int64("loop_times", times))
			c.endpointPool.Range(func(key, value any) bool {
				conn, ok := value.(*Conn)
				connState := conn.conn.GetState()
				if connState != connectivity.Ready {
					c.endpointPool.Delete(key)
					conn.conn.Close()
					logger.Info("[rpc_client_pool_monitor] endpoint connection was cleaned up due to invalid state", zap.String("connection", conn.String()), zap.String("state", connState.String()))
				} else if ok && time.Now().Unix()-atomic.LoadInt64(&conn.lastActiveAt) > endpointMaxIdleSec {
					c.endpointPool.Delete(key)
					conn.conn.Close()
					logger.Info("[rpc_client_pool_monitor] endpoint connection was cleaned up due to exceeding the maximum idle time", zap.String("connection", conn.String()), zap.String("state", connState.String()), zap.Int("endpointMaxIdleSec", endpointMaxIdleSec))
				} else {
					size++
					logger.Info("[rpc_client_pool_monitor] endpoint connection living", zap.String("connection", conn.String()), zap.String("state", connState.String()))
				}
				return true
			})
			logger.Info("[rpc_client_pool_monitor] endpoint connection monitor finished", zap.Int64("loop_times", times), zap.Int("pool_size", size))
		}
	})
}

func needRetry(err error) bool {
	switch status.Code(err) {
	case codes.Unavailable, codes.NotFound, codes.Unimplemented, codes.Internal:
		return true
	default:
		return false
	}
}