package agent

import (
	"context"
	"encoding/json"
	"net/http"
	"testing"

	"github.com/stretchr/testify/require"
	"github.com/tidbcloud/dataflow-service/internal/constants"
	infraapi "github.com/tidbcloud/dataflow-service/pkg/infra-api"
)

func TestLightningClient_RequestLightning(t *testing.T) {
	mockClient := &mockHttpRequester{}
	l := &LightningClient{
		AgentClient: &AgentClientImpl{
			BaseURL:      "test",
			ProviderType: infraapi.ProviderTypeEKS,
			Client:       mockClient,
		},
		lightningAddr: "fake.lightning.addr",
	}
	ctx := context.Background()

	mockClient.On("RequestWithHeaders", ctx, "GET", "http://test/proxy/lightning/happy_path", nil, map[string]string{
		constants.LightningAddrHeader: "fake.lightning.addr",
	}).Return(http.StatusOK, []byte("test"), nil)
	respBytes, err := l.requestLightning(ctx, "GET", "happy_path", nil)
	require.NoError(t, err)
	require.Equal(t, []byte("test"), respBytes)
}

func TestLightningClient_GetTaskProgress(t *testing.T) {
	mockClient := &mockHttpRequester{}
	l := &LightningClient{
		AgentClient: &AgentClientImpl{
			BaseURL:      "test",
			ProviderType: infraapi.ProviderTypeEKS,
			Client:       mockClient,
		},
		lightningAddr: "fake.lightning.addr",
	}
	ctx := context.Background()

	mockClient.On("RequestWithHeaders", ctx, "GET", "http://test/proxy/lightning/progress/task", nil, map[string]string{
		constants.LightningAddrHeader: "fake.lightning.addr",
	}).Return(http.StatusOK, []byte("test"), nil)
	_, err := l.GetTaskProgress(ctx)
	require.Error(t, err)

	expected := &TaskProgress{
		Tables: map[string]*TableProgress{
			"test": {
				TotalWritten: 1,
				TotalSize:    2,
				Status:       taskStatusRunning,
				Message:      "test",
			},
		},
		Status:  taskStatusRunning,
		Message: "test",
	}
	mockResp, err := json.Marshal(expected)
	require.NoError(t, err)

	l.lightningAddr = "happy.lightning.addr"
	mockClient.On("RequestWithHeaders", ctx, "GET", "http://test/proxy/lightning/progress/task", nil, map[string]string{
		constants.LightningAddrHeader: "happy.lightning.addr",
	}).Return(http.StatusOK, mockResp, nil)
	taskProgress, err := l.GetTaskProgress(ctx)
	require.NoError(t, err)
	require.Equal(t, expected, taskProgress)
}