package slprober

import (
	"context"
	"errors"
	"testing"
	"time"

	"github.com/pingcap/tidb-operator/pkg/apis/pingcap/v1alpha1"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/mock"
	v1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
)

func TestProbeTiDBCluster(t *testing.T) {
	tests := []struct {
		name            string
		tc              *v1alpha1.TidbCluster
		tidbSvc         *v1.Service
		tiproxySvc      *v1.Service
		handshakeErrors map[string]error
		expectedSuccess bool
	}{
		{
			name: "no TiProxy, TiDB alive",
			tc: &v1alpha1.TidbCluster{
				ObjectMeta: metav1.ObjectMeta{
					Namespace: "tidb123",
					Name:      "db",
				},
			},
			tidbSvc: &v1.Service{
				ObjectMeta: metav1.ObjectMeta{
					Namespace: "tidb123",
					Name:      "db-tidb",
				},
				Spec: v1.ServiceSpec{
					ClusterIP: "1.2.3.4",
					Ports: []v1.ServicePort{
						{Name: "mysql-client", Port: 4000},
					},
				},
			},
			tiproxySvc: nil,
			handshakeErrors: map[string]error{
				"1.2.3.4:4000": nil,
			},
			expectedSuccess: true,
		},
		{
			name: "both services alive",
			tc: &v1alpha1.TidbCluster{
				ObjectMeta: metav1.ObjectMeta{
					Namespace: "tidb123",
					Name:      "db",
				},
				Spec: v1alpha1.TidbClusterSpec{
					TiProxy: &v1alpha1.TiProxySpec{Replicas: 2},
				},
			},
			tidbSvc: &v1.Service{
				ObjectMeta: metav1.ObjectMeta{
					Namespace: "tidb123",
					Name:      "db-tidb",
				},
				Spec: v1.ServiceSpec{
					ClusterIP: "1.2.3.4",
					Ports: []v1.ServicePort{
						{Name: "mysql-client", Port: 4000},
					},
				},
			},
			tiproxySvc: &v1.Service{
				ObjectMeta: metav1.ObjectMeta{
					Namespace: "tidb123",
					Name:      "db-tiproxy",
				},
				Spec: v1.ServiceSpec{
					ClusterIP: "5.6.7.8",
					Ports: []v1.ServicePort{
						{Name: "tiproxy-sql", Port: 6000},
					},
				},
			},
			handshakeErrors: map[string]error{
				"1.2.3.4:4000": nil,
				"5.6.7.8:6000": nil,
			},
			expectedSuccess: true,
		},
		{
			name: "tiproxy service dead, tidb service alive",
			tc: &v1alpha1.TidbCluster{
				ObjectMeta: metav1.ObjectMeta{
					Namespace: "tidb123",
					Name:      "db",
				},
				Spec: v1alpha1.TidbClusterSpec{
					TiProxy: &v1alpha1.TiProxySpec{Replicas: 2},
				},
			},
			tidbSvc: &v1.Service{
				ObjectMeta: metav1.ObjectMeta{
					Namespace: "tidb123",
					Name:      "db-tidb",
				},
				Spec: v1.ServiceSpec{
					ClusterIP: "1.2.3.4",
					Ports: []v1.ServicePort{
						{Name: "mysql-client", Port: 4000},
					},
				},
			},
			tiproxySvc: &v1.Service{
				ObjectMeta: metav1.ObjectMeta{
					Namespace: "tidb123",
					Name:      "db-tiproxy",
				},
				Spec: v1.ServiceSpec{
					ClusterIP: "5.6.7.8",
					Ports: []v1.ServicePort{
						{Name: "tiproxy-sql", Port: 6000},
					},
				},
			},
			handshakeErrors: map[string]error{
				"1.2.3.4:4000": nil,
				"5.6.7.8:6000": errors.New("connection failed"),
			},
			expectedSuccess: false,
		},
		{
			name: "tidb service dead, tiproxy service alive",
			tc: &v1alpha1.TidbCluster{
				ObjectMeta: metav1.ObjectMeta{
					Namespace: "tidb123",
					Name:      "db",
				},
				Spec: v1alpha1.TidbClusterSpec{
					TiProxy: &v1alpha1.TiProxySpec{Replicas: 2},
				},
			},
			tidbSvc: &v1.Service{
				ObjectMeta: metav1.ObjectMeta{
					Namespace: "tidb123",
					Name:      "db-tidb",
				},
				Spec: v1.ServiceSpec{
					ClusterIP: "1.2.3.4",
					Ports: []v1.ServicePort{
						{Name: "mysql-client", Port: 4000},
					},
				},
			},
			tiproxySvc: &v1.Service{
				ObjectMeta: metav1.ObjectMeta{
					Namespace: "tidb123",
					Name:      "db-tiproxy",
				},
				Spec: v1.ServiceSpec{
					ClusterIP: "5.6.7.8",
					Ports: []v1.ServicePort{
						{Name: "tiproxy-sql", Port: 6000},
					},
				},
			},
			handshakeErrors: map[string]error{
				"1.2.3.4:4000": errors.New("connection failed"),
				"5.6.7.8:6000": nil,
			},
			expectedSuccess: false,
		},
		{
			name: "tidb service timeout, tiproxy service alive",
			tc: &v1alpha1.TidbCluster{
				ObjectMeta: metav1.ObjectMeta{
					Namespace: "tidb123",
					Name:      "db",
				},
				Spec: v1alpha1.TidbClusterSpec{
					TiProxy: &v1alpha1.TiProxySpec{Replicas: 2},
				},
			},
			tidbSvc: &v1.Service{
				ObjectMeta: metav1.ObjectMeta{
					Namespace: "tidb123",
					Name:      "db-tidb",
				},
				Spec: v1.ServiceSpec{
					ClusterIP: "1.2.3.4",
					Ports: []v1.ServicePort{
						{Name: "mysql-client", Port: 4000},
					},
				},
			},
			tiproxySvc: &v1.Service{
				ObjectMeta: metav1.ObjectMeta{
					Namespace: "tidb123",
					Name:      "db-tiproxy",
				},
				Spec: v1.ServiceSpec{
					ClusterIP: "5.6.7.8",
					Ports: []v1.ServicePort{
						{Name: "tiproxy-sql", Port: 6000},
					},
				},
			},
			handshakeErrors: map[string]error{
				"1.2.3.4:4000": context.DeadlineExceeded,
				"5.6.7.8:6000": nil,
			},
			expectedSuccess: false,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			mockHandshakeCli := new(MockHandshakeClient)
			shouldTimeout := false
			for addr, err := range tt.handshakeErrors {
				if errors.Is(err, context.DeadlineExceeded) {
					shouldTimeout = true
					mockHandshakeCli.On("TryConnect", mock.Anything, addr).After(1*time.Second).Return(nil, err)
				} else {
					mockHandshakeCli.On("TryConnect", mock.Anything, addr).Return(nil, err)
				}
			}

			resultCh := make(chan tcResult, 1)
			ctx, cancel := context.WithTimeout(context.Background(), 500*time.Millisecond)
			defer cancel()

			go probeTiDBCluster(ctx, tt.tc, tt.tidbSvc, tt.tiproxySvc, resultCh, mockHandshakeCli)

			select {
			case result := <-resultCh:
				if shouldTimeout {
					t.Fatal("Expected no result due to context cancellation")
				} else {
					assert.Equal(t, tt.expectedSuccess, result.success)
				}
			case <-ctx.Done():
				if !shouldTimeout {
					t.Fatal("test timed out")
				}
			}

			time.Sleep(200 * time.Millisecond) // Wait for goroutines to finish
			mockHandshakeCli.AssertExpectations(t)
		})
	}
}