package kubernetes

import (
	"context"
	"fmt"
	"os"
	"strings"
	"sync"
	"time"

	"github.com/openimsdk/tools/discovery"
	"github.com/openimsdk/tools/log"
	"github.com/openimsdk/tools/utils/datautil"
	"google.golang.org/grpc"
	"google.golang.org/grpc/connectivity"
	"google.golang.org/grpc/credentials/insecure"
	"google.golang.org/grpc/keepalive"
	v1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/client-go/informers"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/rest"
	"k8s.io/client-go/tools/cache"
)

// addrConn 存储连接和地址信息，用于连接复用
type addrConn struct {
	conn   *grpc.ClientConn
	addr   string
	reused bool // 标记是否被复用
}

type KubernetesConnManager struct {
	clientset   *kubernetes.Clientset
	namespace   string
	dialOptions []grpc.DialOption

	rpcTargets map[string]string
	selfTarget string

	// watchNames 只监听这些服务的 Endpoints 变化
	watchNames []string

	mu      sync.RWMutex
	connMap map[string][]*addrConn
}

// NewKubernetesConnManager creates a new connection manager that uses Kubernetes services for service discovery.
func NewKubernetesConnManager(namespace string, watchNames []string, options ...grpc.DialOption) (*KubernetesConnManager, error) {
	ctx := context.Background()
	log.ZInfo(ctx, "K8s Discovery: Initializing connection manager", "namespace", namespace, "watchNames", watchNames)

	// 获取集群内配置
	config, err := rest.InClusterConfig()
	if err != nil {
		log.ZError(ctx, "K8s Discovery: Failed to create in-cluster config", err)
		return nil, fmt.Errorf("failed to create in-cluster config: %v", err)
	}
	log.ZDebug(ctx, "K8s Discovery: Successfully created in-cluster config")

	// 创建 K8s API 客户端
	clientset, err := kubernetes.NewForConfig(config)
	if err != nil {
		log.ZError(ctx, "K8s Discovery: Failed to create clientset", err)
		return nil, fmt.Errorf("failed to create clientset: %v", err)
	}
	log.ZDebug(ctx, "K8s Discovery: Successfully created clientset")

	// 初始化连接管理器
	k := &KubernetesConnManager{
		clientset:   clientset,
		namespace:   namespace,
		dialOptions: options,
		connMap:     make(map[string][]*addrConn),
		rpcTargets:  make(map[string]string),
		watchNames:  watchNames,
	}

	// 启动后台 goroutine 监听 Endpoints 变化
	log.ZInfo(ctx, "K8s Discovery: Starting Endpoints watcher")
	go k.watchEndpoints()

	log.ZInfo(ctx, "K8s Discovery: Connection manager initialized successfully")
	return k, nil
}

// parseServiceName 解析服务名，去掉端口信息
// 例如：user-rpc-service:http-10320 -> user-rpc-service
func parseServiceName(serviceName string) string {
	if idx := strings.Index(serviceName, ":"); idx != -1 {
		return serviceName[:idx]
	}
	return serviceName
}

// initializeConns 初始化指定服务的所有 gRPC 连接（支持连接复用）
func (k *KubernetesConnManager) initializeConns(serviceName string) error {
	ctx := context.Background()
	log.ZInfo(ctx, "K8s Discovery: Starting to initialize connections", "serviceName", serviceName)

	// 步骤 1: 获取 Service 的端口
	port, err := k.getServicePort(serviceName)
	if err != nil {
		log.ZError(ctx, "K8s Discovery: Failed to get service port", err, "serviceName", serviceName)
		return fmt.Errorf("failed to get service port: %w", err)
	}
	log.ZDebug(ctx, "K8s Discovery: Got service port", "serviceName", serviceName, "port", port)

	// 步骤 2: 获取旧连接，建立地址到连接的映射（用于复用）
	k.mu.Lock()
	oldList := k.connMap[serviceName]
	addrMap := make(map[string]*addrConn, len(oldList))
	for _, ac := range oldList {
		addrMap[ac.addr] = ac
		ac.reused = false // 重置复用标记
	}
	k.mu.Unlock()

	log.ZDebug(ctx, "K8s Discovery: Old connections snapshot", "serviceName", serviceName, "count", len(oldList))

	// 步骤 3: 获取 Service 对应的 Endpoints
	endpoints, err := k.clientset.CoreV1().Endpoints(k.namespace).Get(
		ctx,
		serviceName,
		metav1.GetOptions{},
	)
	if err != nil {
		log.ZError(ctx, "K8s Discovery: Failed to get endpoints", err, "serviceName", serviceName)
		return fmt.Errorf("failed to get endpoints for service %s: %w", serviceName, err)
	}

	// 统计 Endpoints 数量
	var totalAddresses int
	for _, subset := range endpoints.Subsets {
		totalAddresses += len(subset.Addresses)
	}
	log.ZDebug(ctx, "K8s Discovery: Found endpoint addresses", "serviceName", serviceName, "count", totalAddresses)

	// 步骤 4: 为每个 Pod IP 创建或复用 gRPC 连接
	var newList []*addrConn
	var reusedCount, createdCount int

	for _, subset := range endpoints.Subsets {
		for _, address := range subset.Addresses {
			target := fmt.Sprintf("%s:%d", address.IP, port)

			// 检查是否可以复用旧连接
			if ac, ok := addrMap[target]; ok {
				// 复用旧连接
				ac.reused = true
				newList = append(newList, ac)
				reusedCount++
				log.ZDebug(ctx, "K8s Discovery: Reusing existing connection", "serviceName", serviceName, "target", target)
				continue
			}

			// 创建新连接
			log.ZDebug(ctx, "K8s Discovery: Creating new connection", "serviceName", serviceName, "target", target)
			conn, err := grpc.Dial(
				target,
				append(k.dialOptions,
					grpc.WithTransportCredentials(insecure.NewCredentials()),
					grpc.WithKeepaliveParams(keepalive.ClientParameters{
						Time:                10 * time.Second,
						Timeout:             3 * time.Second,
						PermitWithoutStream: true,
					}),
				)...,
			)
			if err != nil {
				log.ZWarn(ctx, "K8s Discovery: Failed to dial endpoint, skipping", err, "serviceName", serviceName, "target", target)
				// 跳过无法连接的端点，不终止整个初始化
				continue
			}

			state := conn.GetState()
			log.ZDebug(ctx, "K8s Discovery: New connection created", "serviceName", serviceName, "target", target, "state", state.String())

			newList = append(newList, &addrConn{conn: conn, addr: target, reused: false})
			createdCount++
		}
	}

	log.ZInfo(ctx, "K8s Discovery: Connection initialization summary", "serviceName", serviceName,
		"total", len(newList), "reused", reusedCount, "created", createdCount)

	// 步骤 5: 收集需要关闭的旧连接（未被复用的）
	var connsToClose []*addrConn
	for _, ac := range oldList {
		if !ac.reused {
			connsToClose = append(connsToClose, ac)
		}
	}

	// 步骤 6: 更新连接映射
	k.mu.Lock()
	k.connMap[serviceName] = newList
	k.mu.Unlock()

	log.ZDebug(ctx, "K8s Discovery: Connection map updated", "serviceName", serviceName,
		"oldCount", len(oldList), "newCount", len(newList), "toClose", len(connsToClose))

	// 步骤 7: 延迟关闭未复用的旧连接
	if len(connsToClose) > 0 {
		log.ZInfo(ctx, "K8s Discovery: Scheduling delayed close for unused connections", "serviceName", serviceName, "count", len(connsToClose), "delaySeconds", 5)
		go func() {
			time.Sleep(5 * time.Second)
			log.ZDebug(ctx, "K8s Discovery: Closing unused old connections", "serviceName", serviceName, "count", len(connsToClose))
			closedCount := 0
			for _, ac := range connsToClose {
				if err := ac.conn.Close(); err != nil {
					log.ZError(ctx, "K8s Discovery: Failed to close old connection", err, "serviceName", serviceName, "addr", ac.addr)
				} else {
					closedCount++
				}
			}
			log.ZInfo(ctx, "K8s Discovery: Closed unused connections", "serviceName", serviceName, "closed", closedCount, "total", len(connsToClose))
		}()
	}

	log.ZInfo(ctx, "K8s Discovery: Connection initialization completed", "serviceName", serviceName)
	return nil
}

// GetConns returns gRPC client connections for a given Kubernetes service name.
func (k *KubernetesConnManager) GetConns(ctx context.Context, serviceName string, opts ...grpc.DialOption) ([]grpc.ClientConnInterface, error) {
	// 解析服务名，去掉端口信息
	svcName := parseServiceName(serviceName)
	log.ZDebug(ctx, "K8s Discovery: GetConns called", "serviceName", serviceName, "parsedName", svcName)

	// 步骤 1: 第一次检查缓存（读锁）
	k.mu.RLock()
	conns, exists := k.connMap[svcName]
	k.mu.RUnlock()

	// 步骤 2: 如果缓存中有连接，检查健康状态
	if exists && len(conns) > 0 {
		log.ZDebug(ctx, "K8s Discovery: Found connections in cache, checking health", "serviceName", svcName, "count", len(conns))

		// 检查连接健康状态
		validConns := k.filterValidConns(ctx, svcName, conns)

		// 如果还有有效连接，更新缓存并返回
		if len(validConns) > 0 {
			if len(validConns) < len(conns) {
				log.ZWarn(ctx, "K8s Discovery: Removed invalid connections", nil, "serviceName", svcName, "removed", len(conns)-len(validConns), "remaining", len(validConns))
				k.mu.Lock()
				k.connMap[svcName] = validConns
				k.mu.Unlock()
			} else {
				log.ZDebug(ctx, "K8s Discovery: All connections are healthy", "serviceName", svcName, "count", len(validConns))
			}
			// 转换为接口类型
			result := make([]grpc.ClientConnInterface, len(validConns))
			for i, ac := range validConns {
				result[i] = ac.conn
			}
			return result, nil
		}

		// 如果所有连接都失效，清除缓存并重新初始化
		log.ZWarn(ctx, "K8s Discovery: All connections are invalid, reinitializing", nil, "serviceName", svcName)
		k.mu.Lock()
		delete(k.connMap, svcName)
		k.mu.Unlock()
	} else {
		log.ZDebug(ctx, "K8s Discovery: No connections in cache, initializing", "serviceName", svcName)
	}

	// 步骤 3: 缓存中没有连接或所有连接都失效，重新初始化
	k.mu.Lock()
	conns, exists = k.connMap[svcName]
	if exists && len(conns) > 0 {
		log.ZDebug(ctx, "K8s Discovery: Connections were initialized by another goroutine", "serviceName", svcName)
		k.mu.Unlock()
		result := make([]grpc.ClientConnInterface, len(conns))
		for i, ac := range conns {
			result[i] = ac.conn
		}
		return result, nil
	}
	k.mu.Unlock()

	// 初始化新连接
	log.ZDebug(ctx, "K8s Discovery: Initializing new connections", "serviceName", svcName)
	if err := k.initializeConns(svcName); err != nil {
		log.ZError(ctx, "K8s Discovery: Failed to initialize connections", err, "serviceName", svcName)
		return nil, fmt.Errorf("failed to initialize connections for service %s: %w", svcName, err)
	}

	// 返回新初始化的连接
	k.mu.RLock()
	conns = k.connMap[svcName]
	k.mu.RUnlock()

	log.ZDebug(ctx, "K8s Discovery: Returning connections", "serviceName", svcName, "count", len(conns))
	result := make([]grpc.ClientConnInterface, len(conns))
	for i, ac := range conns {
		result[i] = ac.conn
	}
	return result, nil
}

// filterValidConns 过滤出有效的连接
func (k *KubernetesConnManager) filterValidConns(ctx context.Context, serviceName string, conns []*addrConn) []*addrConn {
	validConns := make([]*addrConn, 0, len(conns))
	invalidStates := make(map[string]int)

	for _, ac := range conns {
		state := ac.conn.GetState()

		// 只保留 Ready 和 Idle 状态的连接
		if state == connectivity.Ready || state == connectivity.Idle {
			validConns = append(validConns, ac)
		} else {
			invalidStates[state.String()]++
			log.ZDebug(ctx, "K8s Discovery: Connection is invalid, closing", "serviceName", serviceName, "addr", ac.addr, "state", state.String())
			if err := ac.conn.Close(); err != nil {
				log.ZError(ctx, "K8s Discovery: Failed to close invalid connection", err, "serviceName", serviceName, "addr", ac.addr)
			}
		}
	}

	if len(invalidStates) > 0 {
		log.ZWarn(ctx, "K8s Discovery: Found invalid connections", nil, "serviceName", serviceName, "invalidStates", invalidStates)
	}

	return validConns
}

// GetConn returns a single gRPC client connection for a given Kubernetes service name.
// 重要：GetConn 使用 DNS，避免连接被强制关闭
func (k *KubernetesConnManager) GetConn(ctx context.Context, serviceName string, opts ...grpc.DialOption) (grpc.ClientConnInterface, error) {
	// 解析服务名，去掉端口信息
	svcName := parseServiceName(serviceName)
	log.ZDebug(ctx, "K8s Discovery: GetConn called (using DNS)", "serviceName", serviceName, "parsedName", svcName)

	var target string

	// 检查是否有自定义目标
	if k.rpcTargets[svcName] == "" {
		// 获取 Service 端口
		svcPort, err := k.getServicePort(svcName)
		if err != nil {
			log.ZError(ctx, "K8s Discovery: Failed to get service port", err, "serviceName", svcName)
			return nil, err
		}

		// 构建 K8s DNS 名称
		target = fmt.Sprintf("%s.%s.svc.cluster.local:%d", svcName, k.namespace, svcPort)
		log.ZDebug(ctx, "K8s Discovery: Using DNS target", "serviceName", svcName, "target", target)
	} else {
		target = k.rpcTargets[svcName]
		log.ZDebug(ctx, "K8s Discovery: Using custom target", "serviceName", svcName, "target", target)
	}

	// 创建 gRPC 连接
	log.ZDebug(ctx, "K8s Discovery: Dialing DNS target", "serviceName", svcName, "target", target)
	conn, err := grpc.DialContext(
		ctx,
		target,
		append([]grpc.DialOption{
			grpc.WithTransportCredentials(insecure.NewCredentials()),
			grpc.WithDefaultCallOptions(
				grpc.MaxCallRecvMsgSize(1024*1024*10),
				grpc.MaxCallSendMsgSize(1024*1024*20),
			),
			grpc.WithKeepaliveParams(keepalive.ClientParameters{
				Time:                10 * time.Second,
				Timeout:             3 * time.Second,
				PermitWithoutStream: true,
			}),
		}, k.dialOptions...)...,
	)
	if err != nil {
		log.ZError(ctx, "K8s Discovery: Failed to dial DNS target", err, "serviceName", svcName, "target", target)
		return nil, err
	}

	log.ZDebug(ctx, "K8s Discovery: GetConn completed successfully", "serviceName", svcName)
	return conn, nil
}

// IsSelfNode checks if the given connection is to the current node
func (k *KubernetesConnManager) IsSelfNode(cc grpc.ClientConnInterface) bool {
	return false
}

// watchEndpoints 监听 Endpoints 资源变化
func (k *KubernetesConnManager) watchEndpoints() {
	ctx := context.Background()
	log.ZInfo(ctx, "K8s Discovery: Starting Endpoints watcher")

	informerFactory := informers.NewSharedInformerFactoryWithOptions(k.clientset, time.Minute*10,
		informers.WithNamespace(k.namespace))
	informer := informerFactory.Core().V1().Endpoints().Informer()
	log.ZDebug(ctx, "K8s Discovery: Endpoints Informer created")

	informer.AddEventHandler(cache.ResourceEventHandlerFuncs{
		AddFunc: func(obj interface{}) {
			k.handleEndpointChange(obj)
		},
		UpdateFunc: func(oldObj, newObj interface{}) {
			oldEndpoint := oldObj.(*v1.Endpoints)
			newEndpoint := newObj.(*v1.Endpoints)

			if k.endpointsChanged(oldEndpoint, newEndpoint) {
				k.handleEndpointChange(newObj)
			}
		},
		DeleteFunc: func(obj interface{}) {
			k.handleEndpointChange(obj)
		},
	})

	log.ZDebug(ctx, "K8s Discovery: Starting Informer factory")
	informerFactory.Start(ctx.Done())

	log.ZDebug(ctx, "K8s Discovery: Waiting for Informer cache to sync")
	if !cache.WaitForCacheSync(ctx.Done(), informer.HasSynced) {
		log.ZError(ctx, "K8s Discovery: Failed to sync Informer cache", nil)
		return
	}
	log.ZInfo(ctx, "K8s Discovery: Informer cache synced successfully")

	log.ZInfo(ctx, "K8s Discovery: Endpoints watcher is running")
	<-ctx.Done()
	log.ZInfo(ctx, "K8s Discovery: Endpoints watcher stopped")
}

// endpointsChanged 检查 Endpoints 是否有实际变化
func (k *KubernetesConnManager) endpointsChanged(old, new *v1.Endpoints) bool {
	oldAddresses := make(map[string]bool)
	for _, subset := range old.Subsets {
		for _, address := range subset.Addresses {
			oldAddresses[address.IP] = true
		}
	}

	newAddresses := make(map[string]bool)
	for _, subset := range new.Subsets {
		for _, address := range subset.Addresses {
			newAddresses[address.IP] = true
		}
	}

	if len(oldAddresses) != len(newAddresses) {
		return true
	}

	for ip := range oldAddresses {
		if !newAddresses[ip] {
			return true
		}
	}

	return false
}

// handleEndpointChange 处理 Endpoints 资源变化
func (k *KubernetesConnManager) handleEndpointChange(obj interface{}) {
	ctx := context.Background()

	endpoint, ok := obj.(*v1.Endpoints)
	if !ok {
		log.ZError(ctx, "K8s Discovery: Expected *v1.Endpoints", nil, "actualType", fmt.Sprintf("%T", obj))
		return
	}

	serviceName := endpoint.Name

	// 只处理 watchNames 中的服务
	if len(k.watchNames) > 0 && !datautil.Contain(serviceName, k.watchNames...) {
		log.ZDebug(ctx, "K8s Discovery: Ignoring Endpoints change (not in watchNames)", "serviceName", serviceName)
		return
	}

	log.ZInfo(ctx, "K8s Discovery: Handling Endpoints change", "serviceName", serviceName)

	var totalAddresses int
	for _, subset := range endpoint.Subsets {
		totalAddresses += len(subset.Addresses)
	}
	log.ZDebug(ctx, "K8s Discovery: Endpoint addresses count", "serviceName", serviceName, "count", totalAddresses)

	if err := k.initializeConns(serviceName); err != nil {
		log.ZError(ctx, "K8s Discovery: Failed to initialize connections", err, "serviceName", serviceName)
	} else {
		log.ZInfo(ctx, "K8s Discovery: Successfully updated connections", "serviceName", serviceName)
	}
}

// getServicePort 获取 Service 的 RPC 端口
func (k *KubernetesConnManager) getServicePort(serviceName string) (int32, error) {
	ctx := context.Background()
	log.ZDebug(ctx, "K8s Discovery: Getting service port", "serviceName", serviceName)

	svc, err := k.clientset.CoreV1().Services(k.namespace).Get(
		ctx,
		serviceName,
		metav1.GetOptions{},
	)
	if err != nil {
		log.ZError(ctx, "K8s Discovery: Failed to get service", err, "serviceName", serviceName, "namespace", k.namespace)
		return 0, fmt.Errorf("failed to get service %s: %w", serviceName, err)
	}

	if len(svc.Spec.Ports) == 0 {
		log.ZError(ctx, "K8s Discovery: Service has no ports defined", nil, "serviceName", serviceName)
		return 0, fmt.Errorf("service %s has no ports defined", serviceName)
	}

	var svcPort int32
	for _, port := range svc.Spec.Ports {
		if port.Port != 10001 {
			svcPort = port.Port
			break
		}
	}

	if svcPort == 0 {
		log.ZError(ctx, "K8s Discovery: Service has no RPC port", nil, "serviceName", serviceName)
		return 0, fmt.Errorf("service %s has no RPC port (all ports are 10001)", serviceName)
	}

	log.ZDebug(ctx, "K8s Discovery: Got service port", "serviceName", serviceName, "port", svcPort)
	return svcPort, nil
}

// Close 关闭所有连接
func (k *KubernetesConnManager) Close() {
	ctx := context.Background()
	log.ZInfo(ctx, "K8s Discovery: Closing all connections")

	k.mu.Lock()
	defer k.mu.Unlock()

	totalConns := 0
	for serviceName, conns := range k.connMap {
		log.ZDebug(ctx, "K8s Discovery: Closing connections for service", "serviceName", serviceName, "count", len(conns))
		for _, ac := range conns {
			if err := ac.conn.Close(); err != nil {
				log.ZError(ctx, "K8s Discovery: Failed to close connection", err, "serviceName", serviceName, "addr", ac.addr)
			}
		}
		totalConns += len(conns)
	}

	log.ZInfo(ctx, "K8s Discovery: Closed all connections", "totalCount", totalConns)
	k.connMap = make(map[string][]*addrConn)
}

// GetSelfConnTarget returns the connection target for the current service.
func (k *KubernetesConnManager) GetSelfConnTarget() string {
	ctx := context.Background()

	if k.selfTarget == "" {
		hostName := os.Getenv("HOSTNAME")
		log.ZDebug(ctx, "K8s Discovery: Getting self connection target", "hostname", hostName)

		pod, err := k.clientset.CoreV1().Pods(k.namespace).Get(ctx, hostName, metav1.GetOptions{})
		if err != nil {
			log.ZError(ctx, "K8s Discovery: Failed to get pod", err, "hostname", hostName)
		}

		for pod.Status.PodIP == "" {
			log.ZDebug(ctx, "K8s Discovery: Waiting for pod IP to be assigned", "hostname", hostName)
			pod, err = k.clientset.CoreV1().Pods(k.namespace).Get(context.TODO(), hostName, metav1.GetOptions{})
			if err != nil {
				log.ZError(ctx, "K8s Discovery: Failed to get pod", err)
			}
			time.Sleep(3 * time.Second)
		}

		var selfPort int32
		for _, port := range pod.Spec.Containers[0].Ports {
			if port.ContainerPort != 10001 {
				selfPort = port.ContainerPort
				break
			}
		}

		k.selfTarget = fmt.Sprintf("%s:%d", pod.Status.PodIP, selfPort)
		log.ZInfo(ctx, "K8s Discovery: Self connection target", "target", k.selfTarget)
	}

	return k.selfTarget
}

// AddOption appends gRPC dial options to the existing options.
func (k *KubernetesConnManager) AddOption(opts ...grpc.DialOption) {
	k.mu.Lock()
	defer k.mu.Unlock()
	k.dialOptions = append(k.dialOptions, opts...)
	log.ZDebug(context.Background(), "K8s Discovery: Added dial options", "count", len(opts))
}

// CloseConn closes a given gRPC client connection.
func (k *KubernetesConnManager) CloseConn(conn *grpc.ClientConn) {
	log.ZDebug(context.Background(), "K8s Discovery: Closing single connection")
	conn.Close()
}

func (k *KubernetesConnManager) Register(ctx context.Context, serviceName, host string, port int, opts ...grpc.DialOption) error {
	// K8s 环境下不需要注册，返回 nil
	return nil
}

func (k *KubernetesConnManager) UnRegister() error {
	// K8s 环境下不需要注销，返回 nil
	return nil
}

func (k *KubernetesConnManager) GetUserIdHashGatewayHost(ctx context.Context, userId string) (string, error) {
	// K8s 环境下不支持，返回空
	return "", nil
}

// KeyValue interface methods - K8s环境下不支持，返回 discovery.ErrNotSupported
// 这样调用方可以通过 errors.Is(err, discovery.ErrNotSupported) 来判断并忽略

func (k *KubernetesConnManager) SetKey(ctx context.Context, key string, value []byte) error {
	return discovery.ErrNotSupported
}

func (k *KubernetesConnManager) SetWithLease(ctx context.Context, key string, val []byte, ttl int64) error {
	return discovery.ErrNotSupported
}

func (k *KubernetesConnManager) GetKey(ctx context.Context, key string) ([]byte, error) {
	return nil, discovery.ErrNotSupported
}

func (k *KubernetesConnManager) GetKeyWithPrefix(ctx context.Context, key string) ([][]byte, error) {
	return nil, discovery.ErrNotSupported
}

func (k *KubernetesConnManager) WatchKey(ctx context.Context, key string, fn discovery.WatchKeyHandler) error {
	return discovery.ErrNotSupported
}
