package grpc

import (
	"context"
	"net"
	"strconv"
	"testing"

	"github.com/stretchr/testify/require"
	agent "github.com/tidbcloud/dataflow-service/dataflow-agent-ng/idl/pbgen/proto/ng_agent"
	"github.com/tidbcloud/dataflow-service/dataflow-agent-ng/pkg/crypto"
)

func TestExecSQLOnDMSourceReturnsSourceNetworkErrorInPayload(t *testing.T) {
	encryptor, pemBytes := newTestRSAEncryptor(t)
	port := closedLocalTCPPort(t)

	resp, err := (&AgentServer{encryptor: encryptor}).ExecSQLOnDMSource(context.Background(), &agent.ExecSQLOnDMSourceRequest{
		Sql:         "SELECT 1",
		ConnProfile: newEncryptedDMSourceConnProfile(t, pemBytes, "127.0.0.1", int32(port)),
	})

	require.NoError(t, err)
	require.NotNil(t, resp)
	require.Contains(t, resp.GetErrorMsg(), "dial tcp")
	require.Contains(t, resp.GetErrorMsg(), "connection refused")
}

func TestExecSQLOnDMSourceReturnsPEMChangedInPayload(t *testing.T) {
	serverEncryptor, _ := newTestRSAEncryptor(t)
	_, stalePEMBytes := newTestRSAEncryptor(t)

	resp, err := (&AgentServer{encryptor: serverEncryptor}).ExecSQLOnDMSource(context.Background(), &agent.ExecSQLOnDMSourceRequest{
		Sql:         "SELECT 1",
		ConnProfile: newEncryptedDMSourceConnProfile(t, stalePEMBytes, "127.0.0.1", 3306),
	})

	require.NoError(t, err)
	require.NotNil(t, resp)
	require.True(t, resp.GetPemChanged())
	require.Empty(t, resp.GetErrorMsg())
}

func newTestRSAEncryptor(t *testing.T) (*crypto.RSAEncryptor, []byte) {
	t.Helper()

	encryptor, err := crypto.NewRSAEncryptor(2048)
	require.NoError(t, err)
	pemBytes, err := encryptor.ExportPublicKeyAsPEM()
	require.NoError(t, err)
	return encryptor, pemBytes
}

func newEncryptedDMSourceConnProfile(t *testing.T, pemBytes []byte, host string, port int32) *agent.ConnProfile {
	t.Helper()

	encryptedPassword, err := crypto.EncryptWithPEM(pemBytes, []byte("source-password"))
	require.NoError(t, err)

	return &agent.ConnProfile{
		User: "source-user",
		Port: port,
		EncryptedPassword: &agent.EncryptedData{
			Pem:  pemBytes,
			Data: encryptedPassword,
		},
		Endpoint: &agent.ConnProfile_Host{Host: host},
	}
}

func closedLocalTCPPort(t *testing.T) int {
	t.Helper()

	listener, err := net.Listen("tcp", "127.0.0.1:0")
	require.NoError(t, err)
	port := listener.Addr().(*net.TCPAddr).Port
	require.NoError(t, listener.Close())
	_, err = net.Dial("tcp", net.JoinHostPort("127.0.0.1", strconv.Itoa(port)))
	require.Error(t, err)
	return port
}