package service

import (
	"context"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"crypto/x509"
	"encoding/pem"
	"errors"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	apiv1 "github.com/tidbcloud/cluster-service-ng/api/tidbcloud/cluster-agent/v1"

	agentcrypto "github.com/tidbcloud/cluster-service-ng/internal/cluster-agent/crypto"
	"github.com/tidbcloud/cluster-service-ng/internal/cluster-agent/internal/sqlexec"
	"github.com/tidbcloud/cluster-service-ng/internal/cluster-agent/internal/utils"
)

type queryMultiCall struct {
	clusterID string
	cmdNames  []string
}

type stubCDCTablesEligibilityExecutor struct {
	out   []sqlexec.QueryResult
	err   error
	calls []queryMultiCall
}

func (e *stubCDCTablesEligibilityExecutor) QueryMulti(
	ctx context.Context,
	clusterID string,
	cmds []sqlexec.QueryCommand,
) ([]sqlexec.QueryResult, error) {
	_ = ctx

	call := queryMultiCall{
		clusterID: clusterID,
		cmdNames:  make([]string, 0, len(cmds)),
	}
	for _, cmd := range cmds {
		if cmd == nil {
			call.cmdNames = append(call.cmdNames, "<nil>")
			continue
		}
		call.cmdNames = append(call.cmdNames, cmd.Name())
	}
	e.calls = append(e.calls, call)

	if e.err != nil {
		return nil, e.err
	}
	return e.out, nil
}

func TestAgentService_QueryTiCDCTablesEligibility(t *testing.T) {
	t.Setenv("NAMESPACE", "tidb123")

	encryptor, err := agentcrypto.NewEncryptor()
	require.NoError(t, err)

	stubExec := &stubCDCTablesEligibilityExecutor{
		out: []sqlexec.QueryResult{
			{
				Rows: [][]any{
					{"db1", "t1"},
				},
			},
			{
				Rows: [][]any{
					{"db2", "t2"},
					{"db3", "t3"},
				},
			},
		},
	}

	oldFactory := newCDCTablesEligibilityExecutor
	newCDCTablesEligibilityExecutor = func(
		resolver utils.AddressResolver,
		auth sqlexec.DBAuth,
		tlsProvider sqlexec.TLSConfigProvider,
	) cdcTablesEligibilityExecutor {
		require.NotNil(t, resolver)
		require.Equal(t, "cloud_admin", auth.Username)
		require.Equal(t, "super_secret", auth.Password.PlainText())
		require.NotNil(t, tlsProvider)
		return stubExec
	}
	t.Cleanup(func() {
		newCDCTablesEligibilityExecutor = oldFactory
	})

	svc := &AgentService{
		encryptor: encryptor,
	}

	req := &apiv1.QueryTiCDCTablesEligibilityReq{
		ClusterId:         "cluster-1",
		EncryptedUsername: encryptForTest(t, encryptor, "cloud_admin"),
		EncryptedPassword: encryptForTest(t, encryptor, "super_secret"),
	}

	resp, err := svc.QueryTiCDCTablesEligibility(context.Background(), req)
	require.NoError(t, err)
	require.NotNil(t, resp)
	assert.Equal(t, []string{"db1"}, resp.GetEligibleTables().GetTableSchemas())
	assert.Equal(t, []string{"t1"}, resp.GetEligibleTables().GetTableNames())
	assert.Equal(t, []string{"db2", "db3"}, resp.GetIneligibleTables().GetTableSchemas())
	assert.Equal(t, []string{"t2", "t3"}, resp.GetIneligibleTables().GetTableNames())

	require.Len(t, stubExec.calls, 1)
	assert.Equal(t, "cluster-1", stubExec.calls[0].clusterID)
	assert.Equal(t, []string{"QueryCDCEligibleTables", "QueryCDCIneligibleTables"}, stubExec.calls[0].cmdNames)
}

func TestAgentService_QueryTiCDCTablesEligibility_Failures(t *testing.T) {
	t.Setenv("NAMESPACE", "tidb123")

	encryptor, err := agentcrypto.NewEncryptor()
	require.NoError(t, err)

	tests := []struct {
		name string
		req  *apiv1.QueryTiCDCTablesEligibilityReq
		out  []sqlexec.QueryResult
		err  error
		want bool
	}{
		{
			name: "executor query failed",
			req: &apiv1.QueryTiCDCTablesEligibilityReq{
				ClusterId:         "cluster-1",
				EncryptedUsername: encryptForTest(t, encryptor, "cloud_admin"),
				EncryptedPassword: encryptForTest(t, encryptor, "super_secret"),
			},
			err:  errors.New("query failed"),
			want: true,
		},
		{
			name: "build ineligible response failed",
			req: &apiv1.QueryTiCDCTablesEligibilityReq{
				ClusterId:         "cluster-1",
				EncryptedUsername: encryptForTest(t, encryptor, "cloud_admin"),
				EncryptedPassword: encryptForTest(t, encryptor, "super_secret"),
			},
			out: []sqlexec.QueryResult{
				{
					Rows: [][]any{
						{"db1", "t1"},
					},
				},
				{
					Rows: [][]any{
						{"db2"},
					},
				},
			},
			want: true,
		},
		{
			name: "nil request",
			req:  nil,
			want: true,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			stubExec := &stubCDCTablesEligibilityExecutor{
				out: tt.out,
				err: tt.err,
			}

			oldFactory := newCDCTablesEligibilityExecutor
			newCDCTablesEligibilityExecutor = func(
				resolver utils.AddressResolver,
				auth sqlexec.DBAuth,
				tlsProvider sqlexec.TLSConfigProvider,
			) cdcTablesEligibilityExecutor {
				require.NotNil(t, resolver)
				require.NotEmpty(t, auth.Username)
				require.False(t, auth.Password.IsEmpty())
				require.NotNil(t, tlsProvider)
				return stubExec
			}
			t.Cleanup(func() {
				newCDCTablesEligibilityExecutor = oldFactory
			})

			svc := &AgentService{
				encryptor: encryptor,
			}

			_, err := svc.QueryTiCDCTablesEligibility(context.Background(), tt.req)
			if tt.want {
				require.Error(t, err)
				return
			}
			require.NoError(t, err)
		})
	}
}

func Test_buildTableSchemaInfoFromQueryResult(t *testing.T) {
	tests := []struct {
		name    string
		out     sqlexec.QueryResult
		want    *apiv1.TableSchemaInfo
		wantErr bool
	}{
		{
			name: "bytes",
			out: sqlexec.QueryResult{
				Rows: [][]any{
					{[]byte("db1"), []byte("t1")},
					{[]byte("db2"), []byte("t2")},
				},
			},
			want: &apiv1.TableSchemaInfo{
				TableSchemas: []string{"db1", "db2"},
				TableNames:   []string{"t1", "t2"},
			},
			wantErr: false,
		},
		{
			name: "string",
			out: sqlexec.QueryResult{
				Rows: [][]any{
					{"db1", "t1"},
				},
			},
			want: &apiv1.TableSchemaInfo{
				TableSchemas: []string{"db1"},
				TableNames:   []string{"t1"},
			},
			wantErr: false,
		},
		{
			name: "empty",
			out: sqlexec.QueryResult{
				Rows: nil,
			},
			want:    &apiv1.TableSchemaInfo{},
			wantErr: false,
		},
		{
			name: "missing columns",
			out: sqlexec.QueryResult{
				Rows: [][]any{
					{[]byte("db1")},
				},
			},
			want:    nil,
			wantErr: true,
		},
		{
			name: "nil value",
			out: sqlexec.QueryResult{
				Rows: [][]any{
					{nil, []byte("t1")},
				},
			},
			want:    nil,
			wantErr: true,
		},
		{
			name: "unsupported type",
			out: sqlexec.QueryResult{
				Rows: [][]any{
					{int64(1), []byte("t1")},
				},
			},
			want:    nil,
			wantErr: true,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got, err := buildTableSchemaInfoFromQueryResult(tt.out)
			if tt.wantErr {
				assert.Error(t, err)
				return
			}

			require.NoError(t, err)
			assert.Equal(t, tt.want.GetTableSchemas(), got.GetTableSchemas())
			assert.Equal(t, tt.want.GetTableNames(), got.GetTableNames())
		})
	}
}

func encryptForTest(t *testing.T, encryptor *agentcrypto.Encryptor, plaintext string) []byte {
	t.Helper()

	pubPEM, err := encryptor.ExportPublicKeyAsPEM()
	require.NoError(t, err)

	block, _ := pem.Decode(pubPEM)
	require.NotNil(t, block)
	require.Equal(t, "RSA PUBLIC KEY", block.Type)

	pubInterface, err := x509.ParsePKIXPublicKey(block.Bytes)
	require.NoError(t, err)

	pubKey, ok := pubInterface.(*rsa.PublicKey)
	require.True(t, ok)

	encrypted, err := rsa.EncryptOAEP(sha256.New(), rand.Reader, pubKey, []byte(plaintext), nil)
	require.NoError(t, err)
	return encrypted
}