package grpc

import (
	"context"
	"database/sql"
	"errors"
	"fmt"
	"net/url"
	"strings"
	"testing"
	"time"

	"github.com/pingcap/tidb/pkg/lightning/mydump"
	"github.com/stretchr/testify/require"
	"github.com/tidbcloud/dataflow-service/dataflow-agent-ng/agent/tidbcloud"
	agent "github.com/tidbcloud/dataflow-service/dataflow-agent-ng/idl/pbgen/proto/ng_agent"
	"github.com/tidbcloud/dataflow-service/dataflow-agent-ng/pkg/loader"
	"github.com/tidbcloud/pkgv2/bizerr"
	"go.uber.org/zap"
)

// stubScanDataFiles returns a mock scanDataFiles function that looks up results from a
// map keyed by URL. If the URL is not in the map, it returns loader.ErrFileNotFound.
func stubScanDataFiles(m map[string][]*mydump.SourceFileMeta) func(ctx context.Context, url string, maxScanFiles int) ([]*mydump.SourceFileMeta, int64, error) {
	return func(ctx context.Context, url string, maxScanFiles int) ([]*mydump.SourceFileMeta, int64, error) {
		files, ok := m[url]
		if !ok {
			return nil, 0, fmt.Errorf("%w: %s", loader.ErrFileNotFound, url)
		}
		var totalSize int64
		for _, f := range files {
			totalSize += f.FileSize
		}
		return files, totalSize, nil
	}
}

func TestGetImportFiles_WithTableMeta_SchemaExists(t *testing.T) {
	s := &AgentServer{
		scanDataFiles: stubScanDataFiles(map[string][]*mydump.SourceFileMeta{
			"s3://bucket/path/data*.csv": {
				{Path: "path/data001.csv", FileSize: 1024, Type: mydump.SourceTypeCSV},
				{Path: "path/data002.csv", FileSize: 2048, Type: mydump.SourceTypeCSV},
			},
			"s3://bucket/path/test.sbtest1-schema.sql*": {
				{Path: "path/test.sbtest1-schema.sql", FileSize: 256, Type: mydump.SourceTypeSQL},
			},
		}),
	}
	req := &agent.GetImportFilesReq{
		ClusterId: 1,
		SourceUri: "s3://bucket/path/",
		TableMeta: []*agent.GetImportFilesReq_TableMeta{
			{
				Database:     "test",
				Table:        "sbtest1",
				WildcardPath: "s3://bucket/path/data*.csv",
			},
		},
	}

	resp, err := s.GetImportFiles(context.Background(), req)
	require.NoError(t, err)
	require.Len(t, resp.TableMeta, 1)

	meta := resp.TableMeta[0]
	require.Equal(t, "test", meta.Database)
	require.Equal(t, "sbtest1", meta.Table)
	require.Len(t, meta.DataFiles, 2)
	require.Equal(t, uint64(1024+2048), meta.TotalSize)
	require.Equal(t, "s3://bucket/path/test.sbtest1-schema.sql", meta.SchemaFile)
}

func TestGetImportFiles_WithTableMeta_SchemaNotFound(t *testing.T) {
	// Only data files exist, no schema file.
	s := &AgentServer{
		scanDataFiles: stubScanDataFiles(map[string][]*mydump.SourceFileMeta{
			"s3://bucket/path/data*.csv": {
				{Path: "path/data001.csv", FileSize: 512, Type: mydump.SourceTypeCSV},
			},
		}),
	}
	req := &agent.GetImportFilesReq{
		ClusterId: 1,
		SourceUri: "s3://bucket/path/",
		TableMeta: []*agent.GetImportFilesReq_TableMeta{
			{
				Database:     "test",
				Table:        "sbtest1",
				WildcardPath: "s3://bucket/path/data*.csv",
			},
		},
	}

	resp, err := s.GetImportFiles(context.Background(), req)
	require.NoError(t, err)
	require.Len(t, resp.TableMeta, 1)

	meta := resp.TableMeta[0]
	require.Equal(t, "test", meta.Database)
	require.Equal(t, "sbtest1", meta.Table)
	require.Len(t, meta.DataFiles, 1)
	// SchemaFile should be empty when schema file does not exist
	require.Empty(t, meta.SchemaFile)
}

func TestGetImportFiles_WithTableMeta_ScanSchemaRealError(t *testing.T) {
	networkErr := errors.New("network timeout")

	callCount := 0
	s := &AgentServer{
		scanDataFiles: func(ctx context.Context, url string, maxScanFiles int) ([]*mydump.SourceFileMeta, int64, error) {
			callCount++
			if callCount == 1 {
				// First call: data files scan succeeds
				return []*mydump.SourceFileMeta{
					{Path: "path/data001.csv", FileSize: 100, Type: mydump.SourceTypeCSV},
				}, 100, nil
			}
			// Second call: schema file scan hits a real error (not file-not-found)
			return nil, 0, networkErr
		},
	}
	req := &agent.GetImportFilesReq{
		ClusterId: 1,
		SourceUri: "s3://bucket/path/",
		TableMeta: []*agent.GetImportFilesReq_TableMeta{
			{
				Database:     "test",
				Table:        "sbtest1",
				WildcardPath: "s3://bucket/path/data*.csv",
			},
		},
	}

	resp, err := s.GetImportFiles(context.Background(), req)
	// Real errors (non file-not-found) should be propagated
	require.Error(t, err)
	require.ErrorIs(t, err, networkErr)
	require.Nil(t, resp)
}

func TestGetImportFiles_WithTableMeta_MultipleTables(t *testing.T) {
	s := &AgentServer{
		scanDataFiles: stubScanDataFiles(map[string][]*mydump.SourceFileMeta{
			"s3://bucket/path/users*.csv": {
				{Path: "path/users001.csv", FileSize: 1000, Type: mydump.SourceTypeCSV},
			},
			"s3://bucket/path/mydb.users-schema.sql*": {
				{Path: "path/mydb.users-schema.sql", FileSize: 128, Type: mydump.SourceTypeSQL},
			},
			"s3://bucket/path/orders*.csv": {
				{Path: "path/orders001.csv", FileSize: 2000, Type: mydump.SourceTypeCSV},
				{Path: "path/orders002.csv", FileSize: 3000, Type: mydump.SourceTypeCSV},
			},
			// No schema file for orders
		}),
	}
	req := &agent.GetImportFilesReq{
		ClusterId: 1,
		SourceUri: "s3://bucket/path/",
		TableMeta: []*agent.GetImportFilesReq_TableMeta{
			{
				Database:     "mydb",
				Table:        "users",
				WildcardPath: "s3://bucket/path/users*.csv",
			},
			{
				Database:     "mydb",
				Table:        "orders",
				WildcardPath: "s3://bucket/path/orders*.csv",
			},
		},
	}

	resp, err := s.GetImportFiles(context.Background(), req)
	require.NoError(t, err)
	require.Len(t, resp.TableMeta, 2)

	// Table 1: users — has schema file
	usersMeta := resp.TableMeta[0]
	require.Equal(t, "mydb", usersMeta.Database)
	require.Equal(t, "users", usersMeta.Table)
	require.Len(t, usersMeta.DataFiles, 1)
	require.Equal(t, "s3://bucket/path/mydb.users-schema.sql", usersMeta.SchemaFile)

	// Table 2: orders — no schema file
	ordersMeta := resp.TableMeta[1]
	require.Equal(t, "mydb", ordersMeta.Database)
	require.Equal(t, "orders", ordersMeta.Table)
	require.Len(t, ordersMeta.DataFiles, 2)
	require.Equal(t, uint64(2000+3000), ordersMeta.TotalSize)
	require.Empty(t, ordersMeta.SchemaFile)
}

func TestGetImportFiles_WithTableMeta_OssSchemeConvertedToS3(t *testing.T) {
	var capturedDataURL string
	s := &AgentServer{
		scanDataFiles: func(ctx context.Context, url string, maxScanFiles int) ([]*mydump.SourceFileMeta, int64, error) {
			if capturedDataURL == "" {
				capturedDataURL = url
				return []*mydump.SourceFileMeta{
					{Path: "path/data001.csv", FileSize: 100, Type: mydump.SourceTypeCSV},
				}, 100, nil
			}
			// Schema scan — return not found
			return nil, 0, fmt.Errorf("%w: %s", loader.ErrFileNotFound, url)
		},
	}
	req := &agent.GetImportFilesReq{
		ClusterId: 1,
		SourceUri: "s3://bucket/path/?access-key=abc",
		TableMeta: []*agent.GetImportFilesReq_TableMeta{
			{
				Database:     "test",
				Table:        "t1",
				WildcardPath: "oss://bucket/path/data*.csv",
			},
		},
	}

	resp, err := s.GetImportFiles(context.Background(), req)
	require.NoError(t, err)
	require.Len(t, resp.TableMeta, 1)

	// The oss:// scheme should have been converted to s3:// and query params preserved
	require.Contains(t, capturedDataURL, "s3://")
	require.Contains(t, capturedDataURL, "access-key=abc")
}

func TestGetImportFiles_WithTableMeta_CompressedSchemaExists(t *testing.T) {
	s := &AgentServer{
		scanDataFiles: stubScanDataFiles(map[string][]*mydump.SourceFileMeta{
			"s3://bucket/path/data*.csv": {
				{Path: "path/data001.csv", FileSize: 1024, Type: mydump.SourceTypeCSV},
			},
			"s3://bucket/path/test.sbtest1-schema.sql*": {
				{Path: "path/test.sbtest1-schema.sql.gz", FileSize: 256, Type: mydump.SourceTypeSQL, Compression: mydump.CompressionGZ},
			},
		}),
	}
	req := &agent.GetImportFilesReq{
		ClusterId: 1,
		SourceUri: "s3://bucket/path/",
		TableMeta: []*agent.GetImportFilesReq_TableMeta{
			{
				Database:     "test",
				Table:        "sbtest1",
				WildcardPath: "s3://bucket/path/data*.csv",
			},
		},
	}

	resp, err := s.GetImportFiles(context.Background(), req)
	require.NoError(t, err)
	require.Len(t, resp.TableMeta, 1)
	require.Equal(t, "s3://bucket/path/test.sbtest1-schema.sql.gz", resp.TableMeta[0].SchemaFile)
}

func TestGetImportFiles_WithTableMeta_CompressedSchemaWithQueryCredentials(t *testing.T) {
	const sourceURI = "s3://bucket/path/?access-key=AKIA123&external-id=10233074202638838702&secret-access-key=abc%2Bdef"

	callCount := 0
	s := &AgentServer{
		scanDataFiles: func(ctx context.Context, rawURL string, maxScanFiles int) ([]*mydump.SourceFileMeta, int64, error) {
			callCount++
			switch callCount {
			case 1:
				return []*mydump.SourceFileMeta{
					{Path: "path/data001.csv", FileSize: 1024, Type: mydump.SourceTypeCSV},
				}, 1024, nil
			case 2:
				parsedURL, err := url.Parse(rawURL)
				require.NoError(t, err)
				require.Equal(t, "/path/test.sbtest1-schema.sql*", parsedURL.Path)
				require.Equal(t, "AKIA123", parsedURL.Query().Get("access-key"))
				require.Equal(t, "10233074202638838702", parsedURL.Query().Get("external-id"))
				require.Equal(t, "abc+def", parsedURL.Query().Get("secret-access-key"))
				return []*mydump.SourceFileMeta{
					{Path: "path/test.sbtest1-schema.sql.gz", FileSize: 256, Type: mydump.SourceTypeSQL, Compression: mydump.CompressionGZ},
				}, 256, nil
			default:
				return nil, 0, fmt.Errorf("unexpected scanDataFiles call %d", callCount)
			}
		},
	}
	req := &agent.GetImportFilesReq{
		ClusterId: 1,
		SourceUri: sourceURI,
		TableMeta: []*agent.GetImportFilesReq_TableMeta{
			{
				Database:     "test",
				Table:        "sbtest1",
				WildcardPath: "s3://bucket/path/data*.csv",
			},
		},
	}

	resp, err := s.GetImportFiles(context.Background(), req)
	require.NoError(t, err)
	require.Len(t, resp.TableMeta, 1)
	require.Equal(t, "s3://bucket/path/test.sbtest1-schema.sql.gz?access-key=AKIA123&external-id=10233074202638838702&secret-access-key=abc%2Bdef", resp.TableMeta[0].SchemaFile)
}

// When the wildcard points into a subdirectory, the schema lookup must target
// that same subdirectory, not the source URI root. Previously the lookup was
// always anchored at the source URI root, so schemas next to the data in any
// subfolder were silently missed and downstream import failed with
// "Table '<db>.<tbl>' doesn't exist".
func TestGetImportFiles_WithTableMeta_SchemaInWildcardSubdir(t *testing.T) {
	s := &AgentServer{
		scanDataFiles: stubScanDataFiles(map[string][]*mydump.SourceFileMeta{
			"s3://bucket/path/sub/data*.csv": {
				{Path: "path/sub/data001.csv", FileSize: 1024, Type: mydump.SourceTypeCSV},
			},
			"s3://bucket/path/sub/test.sbtest1-schema.sql*": {
				{Path: "path/sub/test.sbtest1-schema.sql", FileSize: 256, Type: mydump.SourceTypeSQL},
			},
		}),
	}
	req := &agent.GetImportFilesReq{
		ClusterId: 1,
		SourceUri: "s3://bucket/path/",
		TableMeta: []*agent.GetImportFilesReq_TableMeta{
			{
				Database:     "test",
				Table:        "sbtest1",
				WildcardPath: "s3://bucket/path/sub/data*.csv",
			},
		},
	}

	resp, err := s.GetImportFiles(context.Background(), req)
	require.NoError(t, err)
	require.Len(t, resp.TableMeta, 1)

	meta := resp.TableMeta[0]
	require.Len(t, meta.DataFiles, 1)
	require.Equal(t, "s3://bucket/path/sub/test.sbtest1-schema.sql", meta.SchemaFile)
}

// Regression test for the original failure report: source URI at the root and
// data files in a subdirectory with a URL-escaped space. The schema file sits
// next to the data and must still be picked up.
func TestGetImportFiles_WithTableMeta_SchemaInSubdirWithEscapedSpace(t *testing.T) {
	const (
		wildcard       = "s3://bucket/e2e/case-ux-3/path%20with%20space/db1.t_ux_3.csv"
		schemaPattern  = "s3://bucket/e2e/case-ux-3/path%20with%20space/db1.t_ux_3-schema.sql*"
		expectedSchema = "s3://bucket/e2e/case-ux-3/path%20with%20space/db1.t_ux_3-schema.sql"
	)
	s := &AgentServer{
		scanDataFiles: stubScanDataFiles(map[string][]*mydump.SourceFileMeta{
			wildcard: {
				{Path: "e2e/case-ux-3/path with space/db1.t_ux_3.csv", FileSize: 124, Type: mydump.SourceTypeCSV},
			},
			schemaPattern: {
				{Path: "e2e/case-ux-3/path with space/db1.t_ux_3-schema.sql", FileSize: 573, Type: mydump.SourceTypeSQL},
			},
		}),
	}
	req := &agent.GetImportFilesReq{
		ClusterId: 1,
		SourceUri: "s3://bucket/e2e/case-ux-3/",
		TableMeta: []*agent.GetImportFilesReq_TableMeta{
			{
				Database:     "db1",
				Table:        "t_ux_3",
				WildcardPath: wildcard,
			},
		},
	}

	resp, err := s.GetImportFiles(context.Background(), req)
	require.NoError(t, err)
	require.Len(t, resp.TableMeta, 1)
	require.Equal(t, expectedSchema, resp.TableMeta[0].SchemaFile)
}

func TestGetImportFiles_WithTableMeta_InvalidCompressionSchemaIgnored(t *testing.T) {
	s := &AgentServer{
		scanDataFiles: stubScanDataFiles(map[string][]*mydump.SourceFileMeta{
			"s3://bucket/path/data*.csv": {
				{Path: "path/data001.csv", FileSize: 1024, Type: mydump.SourceTypeCSV},
			},
			"s3://bucket/path/test.sbtest1-schema.sql*": {
				{Path: "path/test.sbtest1-schema.sql.zip", FileSize: 256, Type: mydump.SourceTypeSQL},
			},
		}),
	}
	req := &agent.GetImportFilesReq{
		ClusterId: 1,
		SourceUri: "s3://bucket/path/",
		TableMeta: []*agent.GetImportFilesReq_TableMeta{
			{
				Database:     "test",
				Table:        "sbtest1",
				WildcardPath: "s3://bucket/path/data*.csv",
			},
		},
	}

	resp, err := s.GetImportFiles(context.Background(), req)
	require.NoError(t, err)
	require.Len(t, resp.TableMeta, 1)
	require.Empty(t, resp.TableMeta[0].SchemaFile)
}

func TestCreateImportFilesTask(t *testing.T) {
	blockCh := make(chan struct{})
	s := &AgentServer{
		importTaskManager: newImportAsyncTaskManagerWithRetention(time.Hour),
		scanDataFiles: func(ctx context.Context, rawURL string, maxScanFiles int) ([]*mydump.SourceFileMeta, int64, error) {
			<-blockCh
			return []*mydump.SourceFileMeta{
				{Path: "path/data001.csv", FileSize: 1024, Type: mydump.SourceTypeCSV},
			}, 1024, nil
		},
	}

	createResp, err := s.CreateImportFilesTask(context.Background(), &agent.CreateImportFilesTaskReq{
		TaskId:    "files-task-1",
		ClusterId: 1,
		SourceUri: "s3://bucket/path/",
		TableMeta: []*agent.GetImportFilesReq_TableMeta{
			{
				Database:     "test",
				Table:        "t1",
				WildcardPath: "s3://bucket/path/data*.csv",
			},
		},
	})
	require.NoError(t, err)
	require.Equal(t, "files-task-1", createResp.GetTaskId())

	runningResp, err := s.GetImportFilesTask(context.Background(), &agent.GetImportFilesTaskReq{TaskId: "files-task-1"})
	require.NoError(t, err)
	require.Equal(t, agent.ImportFilesTaskStatus_IMPORT_FILES_TASK_STATUS_RUNNING, runningResp.GetStatus())
	require.Nil(t, runningResp.Result)

	close(blockCh)

	require.Eventually(t, func() bool {
		resp, getErr := s.GetImportFilesTask(context.Background(), &agent.GetImportFilesTaskReq{TaskId: "files-task-1"})
		require.NoError(t, getErr)
		if resp.GetStatus() != agent.ImportFilesTaskStatus_IMPORT_FILES_TASK_STATUS_SUCCEEDED {
			return false
		}
		require.NotNil(t, resp.GetResult())
		require.Len(t, resp.GetResult().GetTableMeta(), 1)
		require.Equal(t, uint64(1024), resp.GetResult().GetTableMeta()[0].GetTotalSize())
		return true
	}, 3*time.Second, 10*time.Millisecond)
}

func TestGetImportFilesTaskFailed(t *testing.T) {
	s := &AgentServer{
		importTaskManager: newImportAsyncTaskManagerWithRetention(time.Hour),
		scanDataFiles: func(ctx context.Context, rawURL string, maxScanFiles int) ([]*mydump.SourceFileMeta, int64, error) {
			return nil, 0, bizerr.New(bizerr.ErrCodeGeneralBusiness, "boom")
		},
	}

	_, err := s.CreateImportFilesTask(context.Background(), &agent.CreateImportFilesTaskReq{
		TaskId:    "files-task-failed",
		ClusterId: 1,
		SourceUri: "s3://bucket/path/",
		TableMeta: []*agent.GetImportFilesReq_TableMeta{
			{
				Database:     "test",
				Table:        "t1",
				WildcardPath: "s3://bucket/path/data*.csv",
			},
		},
	})
	require.NoError(t, err)

	require.Eventually(t, func() bool {
		resp, getErr := s.GetImportFilesTask(context.Background(), &agent.GetImportFilesTaskReq{TaskId: "files-task-failed"})
		require.NoError(t, getErr)
		if resp.GetStatus() != agent.ImportFilesTaskStatus_IMPORT_FILES_TASK_STATUS_FAILED {
			return false
		}
		require.Equal(t, "boom", resp.GetErrorMsg())
		require.Nil(t, resp.Result)
		return true
	}, 3*time.Second, 10*time.Millisecond)
}

func TestGetImportFilesTaskReturnsInternalError(t *testing.T) {
	s := &AgentServer{
		importTaskManager: newImportAsyncTaskManagerWithRetention(time.Hour),
		scanDataFiles: func(ctx context.Context, rawURL string, maxScanFiles int) ([]*mydump.SourceFileMeta, int64, error) {
			return nil, 0, errors.New("boom")
		},
	}

	_, err := s.CreateImportFilesTask(context.Background(), &agent.CreateImportFilesTaskReq{
		TaskId:    "files-task-internal-failed",
		ClusterId: 1,
		SourceUri: "s3://bucket/path/",
		TableMeta: []*agent.GetImportFilesReq_TableMeta{
			{
				Database:     "test",
				Table:        "t1",
				WildcardPath: "s3://bucket/path/data*.csv",
			},
		},
	})
	require.NoError(t, err)

	require.Eventually(t, func() bool {
		resp, getErr := s.GetImportFilesTask(context.Background(), &agent.GetImportFilesTaskReq{TaskId: "files-task-internal-failed"})
		if getErr == nil {
			require.NotNil(t, resp)
			require.Equal(t, agent.ImportFilesTaskStatus_IMPORT_FILES_TASK_STATUS_RUNNING, resp.GetStatus())
			return false
		}
		require.Nil(t, resp)
		require.ErrorContains(t, getErr, "boom")
		return true
	}, time.Second, 10*time.Millisecond)
}

func TestCreateSchemasTask(t *testing.T) {
	blockCh := make(chan struct{})
	s := &AgentServer{
		importTaskManager: newImportAsyncTaskManagerWithRetention(time.Hour),
		createSchemasFn: func(ctx context.Context, req *agent.CreateSchemasReq) (*agent.CreateSchemasResp, error) {
			<-blockCh
			require.Equal(t, uint64(1), req.GetClusterId())
			require.NotNil(t, req.GetCluster())
			require.Equal(t, "root", req.GetCluster().GetUser())
			require.True(t, req.GetCluster().GetPasswordEncrypted())
			require.Len(t, req.GetSchemas(), 1)
			return &agent.CreateSchemasResp{}, nil
		},
	}

	createResp, err := s.CreateSchemasTask(context.Background(), &agent.CreateSchemasTaskReq{
		TaskId:    "schema-task-1",
		ClusterId: 1,
		SourceUri: "s3://bucket/path/",
		Schemas: []*agent.CreateSchemasReq_Schema{
			{
				Database: "test",
				Table:    "t1",
			},
		},
		Cluster: &agent.Cluster{
			ClusterId:         1,
			User:              "root",
			PasswordEncrypted: true,
			EncryptedPassword: &agent.EncryptedData{Pem: []byte("pem"), Data: []byte("data")},
		},
	})
	require.NoError(t, err)
	require.Equal(t, "schema-task-1", createResp.GetTaskId())

	runningResp, err := s.GetSchemasTask(context.Background(), &agent.GetSchemasTaskReq{TaskId: "schema-task-1"})
	require.NoError(t, err)
	require.Equal(t, agent.CreateSchemasTaskStatus_CREATE_SCHEMAS_TASK_STATUS_RUNNING, runningResp.GetStatus())
	require.Nil(t, runningResp.Result)

	close(blockCh)

	require.Eventually(t, func() bool {
		resp, getErr := s.GetSchemasTask(context.Background(), &agent.GetSchemasTaskReq{TaskId: "schema-task-1"})
		require.NoError(t, getErr)
		if resp.GetStatus() != agent.CreateSchemasTaskStatus_CREATE_SCHEMAS_TASK_STATUS_SUCCEEDED {
			return false
		}
		require.NotNil(t, resp.GetResult())
		require.False(t, resp.GetResult().GetPemChanged())
		return true
	}, time.Second, 10*time.Millisecond)
}

type captureClusterProvider struct {
	opts *tidbcloud.DBOptions
	err  error
}

func (p *captureClusterProvider) GetDB(
	_ context.Context,
	_ string,
	_ string,
	_ string,
	opts *tidbcloud.DBOptions,
) (*sql.DB, error) {
	p.opts = opts
	return nil, p.err
}

func TestRunCreateSchemasDisablesForeignKeyChecksOnSchemaConnection(t *testing.T) {
	provider := &captureClusterProvider{err: errors.New("stop after capturing opts")}
	s := &AgentServer{clusterProvider: provider}

	_, err := s.runCreateSchemas(context.Background(), &agent.CreateSchemasReq{
		ClusterId: 1,
		User:      "root",
		Password:  "password",
		SourceUri: "s3://bucket/path/",
		Schemas: []*agent.CreateSchemasReq_Schema{
			{Database: "test", Table: "table_a"},
		},
	})

	require.ErrorContains(t, err, "stop after capturing opts")
	require.NotNil(t, provider.opts)
	require.Equal(t, "0", provider.opts.Params["foreign_key_checks"])
}

type fakeSchemaCreateSDK struct {
	allCalled     bool
	createdTables []string
	failTable     string
	failAll       bool
}

func (f *fakeSchemaCreateSDK) CreateSchemasAndTables(context.Context) error {
	f.allCalled = true
	if f.failAll {
		return errors.New("boom")
	}
	return nil
}

func (f *fakeSchemaCreateSDK) CreateSchemaAndTableByName(_ context.Context, database, table string) error {
	f.createdTables = append(f.createdTables, database+"."+table)
	if table == f.failTable {
		return errors.New("boom")
	}
	return nil
}

func TestRunCreateSchemasContinuesAfterSchemaError(t *testing.T) {
	sdk := &fakeSchemaCreateSDK{failTable: "t2"}

	resp := executeCreateSchemas(context.Background(), zap.NewNop(), sdk, []*agent.CreateSchemasReq_Schema{
		{Database: "db1", Table: "t1"},
		{Database: "db1", Table: "t2"},
		{Database: "db1", Table: "t3"},
	})

	require.Len(t, resp.GetTableResults(), 3)
	require.Equal(t, agent.CreateSchemasResp_TableResult_STATUS_SUCCEEDED, resp.GetTableResults()[0].GetStatus())
	require.Equal(t, agent.CreateSchemasResp_TableResult_STATUS_FAILED, resp.GetTableResults()[1].GetStatus())
	require.Equal(t, "failed to create schemas: schema=db1, table=t2, err=boom", resp.GetTableResults()[1].GetErrorMsg())
	require.Equal(t, agent.CreateSchemasResp_TableResult_STATUS_SUCCEEDED, resp.GetTableResults()[2].GetStatus())
	require.Equal(t, []string{"db1.t1", "db1.t2", "db1.t3"}, sdk.createdTables)
	require.False(t, sdk.allCalled)
}

func TestRunCreateSchemasWithoutExplicitSchemasSurfacesBatchError(t *testing.T) {
	sdk := &fakeSchemaCreateSDK{failAll: true}

	resp := executeCreateSchemas(context.Background(), zap.NewNop(), sdk, nil)

	require.True(t, sdk.allCalled)
	require.Equal(t, "failed to create schemas: err=boom", resp.GetErrorMsg())
	require.Empty(t, resp.GetTableResults())
}

func TestGetSchemasTaskFailed(t *testing.T) {
	s := &AgentServer{
		importTaskManager: newImportAsyncTaskManagerWithRetention(time.Hour),
		createSchemasFn: func(ctx context.Context, req *agent.CreateSchemasReq) (*agent.CreateSchemasResp, error) {
			return nil, bizerr.New(bizerr.ErrCodeGeneralBusiness, "boom")
		},
	}

	_, err := s.CreateSchemasTask(context.Background(), &agent.CreateSchemasTaskReq{
		TaskId:    "schema-task-failed",
		ClusterId: 1,
		SourceUri: "s3://bucket/path/",
		Schemas: []*agent.CreateSchemasReq_Schema{
			{
				Database: "test",
				Table:    "t1",
			},
		},
		Cluster: &agent.Cluster{
			ClusterId:         1,
			User:              "root",
			PasswordEncrypted: true,
			EncryptedPassword: &agent.EncryptedData{Pem: []byte("pem"), Data: []byte("data")},
		},
	})
	require.NoError(t, err)

	require.Eventually(t, func() bool {
		resp, getErr := s.GetSchemasTask(context.Background(), &agent.GetSchemasTaskReq{TaskId: "schema-task-failed"})
		require.NoError(t, getErr)
		if resp.GetStatus() != agent.CreateSchemasTaskStatus_CREATE_SCHEMAS_TASK_STATUS_FAILED {
			return false
		}
		require.Equal(t, "boom", resp.GetErrorMsg())
		require.Nil(t, resp.Result)
		return true
	}, time.Second, 10*time.Millisecond)
}

func TestGetSchemasTaskLiftsBusinessErrorMessage(t *testing.T) {
	msg := "ddl failed: table already exists"
	s := &AgentServer{
		importTaskManager: newImportAsyncTaskManagerWithRetention(time.Hour),
		createSchemasFn: func(ctx context.Context, req *agent.CreateSchemasReq) (*agent.CreateSchemasResp, error) {
			return &agent.CreateSchemasResp{
				ErrorMsg: &msg,
			}, nil
		},
	}

	_, err := s.CreateSchemasTask(context.Background(), &agent.CreateSchemasTaskReq{
		TaskId:    "schema-task-business-failed",
		ClusterId: 1,
		SourceUri: "s3://bucket/path/",
		Schemas: []*agent.CreateSchemasReq_Schema{
			{
				Database: "test",
				Table:    "t1",
			},
		},
		Cluster: &agent.Cluster{
			ClusterId:         1,
			User:              "root",
			PasswordEncrypted: true,
			EncryptedPassword: &agent.EncryptedData{Pem: []byte("pem"), Data: []byte("data")},
		},
	})
	require.NoError(t, err)

	require.Eventually(t, func() bool {
		resp, getErr := s.GetSchemasTask(context.Background(), &agent.GetSchemasTaskReq{TaskId: "schema-task-business-failed"})
		require.NoError(t, getErr)
		if resp.GetStatus() != agent.CreateSchemasTaskStatus_CREATE_SCHEMAS_TASK_STATUS_SUCCEEDED {
			return false
		}
		require.Equal(t, msg, resp.GetErrorMsg())
		require.NotNil(t, resp.Result)
		require.Equal(t, msg, resp.GetResult().GetErrorMsg())
		return true
	}, time.Second, 10*time.Millisecond)
}

func TestCreateImportIntoJobTask(t *testing.T) {
	blockCh := make(chan struct{})
	s := &AgentServer{
		importTaskManager: newImportAsyncTaskManagerWithRetention(time.Hour),
		execSQLOnClusterFn: func(ctx context.Context, req *agent.ExecSQLReq) (*agent.ExecSQLResp, error) {
			require.True(t, strings.HasPrefix(req.GetSql(), "IMPORT INTO "))
			<-blockCh
			return &agent.ExecSQLResp{
				ColumnNames: []string{"Job_ID"},
				Rows: []*agent.Row{
					{
						TypedColumnValues: []*agent.ColumnValue{
							{Value: &agent.ColumnValue_Int64Value{Int64Value: 1001}},
						},
					},
				},
			}, nil
		},
	}

	createResp, err := s.CreateImportIntoJobTask(context.Background(), &agent.CreateImportIntoJobTaskReq{
		TaskId:    "submit-task-1",
		ClusterId: 1,
		Sql:       "IMPORT INTO `db`.`t` FROM 's3://bucket/path/*.csv' WITH DETACHED",
		Cluster: &agent.Cluster{
			ClusterId:         1,
			User:              "root",
			PasswordEncrypted: true,
			EncryptedPassword: &agent.EncryptedData{Pem: []byte("pem"), Data: []byte("data")},
		},
	})
	require.NoError(t, err)
	require.Equal(t, "submit-task-1", createResp.GetTaskId())

	runningResp, err := s.GetImportIntoJobTask(context.Background(), &agent.GetImportIntoJobTaskReq{TaskId: "submit-task-1"})
	require.NoError(t, err)
	require.Equal(t, agent.ImportIntoJobTaskStatus_IMPORT_INTO_JOB_TASK_STATUS_RUNNING, runningResp.GetStatus())
	require.Nil(t, runningResp.Result)

	close(blockCh)

	require.Eventually(t, func() bool {
		resp, getErr := s.GetImportIntoJobTask(context.Background(), &agent.GetImportIntoJobTaskReq{TaskId: "submit-task-1"})
		require.NoError(t, getErr)
		if resp.GetStatus() != agent.ImportIntoJobTaskStatus_IMPORT_INTO_JOB_TASK_STATUS_SUCCEEDED {
			return false
		}
		require.NotNil(t, resp.GetResult())
		require.Equal(t, int64(1001), resp.GetResult().GetJobId())
		return true
	}, 3*time.Second, 10*time.Millisecond)
}

func TestGetImportIntoJobTaskFailed(t *testing.T) {
	s := &AgentServer{
		importTaskManager: newImportAsyncTaskManagerWithRetention(time.Hour),
		execSQLOnClusterFn: func(ctx context.Context, req *agent.ExecSQLReq) (*agent.ExecSQLResp, error) {
			return nil, bizerr.New(bizerr.ErrCodeGeneralBusiness, "boom")
		},
	}

	_, err := s.CreateImportIntoJobTask(context.Background(), &agent.CreateImportIntoJobTaskReq{
		TaskId:    "submit-task-failed",
		ClusterId: 1,
		Sql:       "IMPORT INTO `db`.`t` FROM 's3://bucket/path/*.csv' WITH DETACHED",
		Cluster: &agent.Cluster{
			ClusterId:         1,
			User:              "root",
			PasswordEncrypted: true,
			EncryptedPassword: &agent.EncryptedData{Pem: []byte("pem"), Data: []byte("data")},
		},
	})
	require.NoError(t, err)

	require.Eventually(t, func() bool {
		resp, getErr := s.GetImportIntoJobTask(context.Background(), &agent.GetImportIntoJobTaskReq{TaskId: "submit-task-failed"})
		require.NoError(t, getErr)
		if resp.GetStatus() != agent.ImportIntoJobTaskStatus_IMPORT_INTO_JOB_TASK_STATUS_FAILED {
			return false
		}
		require.Equal(t, "boom", resp.GetErrorMsg())
		require.Nil(t, resp.GetResult())
		return true
	}, time.Second, 10*time.Millisecond)
}

func TestCreateImportIntoJobTaskRejectsNonImportSQL(t *testing.T) {
	s := &AgentServer{importTaskManager: newImportAsyncTaskManagerWithRetention(time.Hour)}

	_, err := s.CreateImportIntoJobTask(context.Background(), &agent.CreateImportIntoJobTaskReq{
		TaskId:    "submit-task-invalid",
		ClusterId: 1,
		Sql:       "SELECT 1",
		Cluster: &agent.Cluster{
			ClusterId:         1,
			User:              "root",
			PasswordEncrypted: true,
			EncryptedPassword: &agent.EncryptedData{Pem: []byte("pem"), Data: []byte("data")},
		},
	})
	require.Error(t, err)
}

func TestBatchCreateImportIntoJobTasksAndBatchGetImportIntoJobTasks(t *testing.T) {
	s := &AgentServer{
		importTaskManager: newImportAsyncTaskManagerWithRetention(time.Hour),
		execSQLOnClusterFn: func(ctx context.Context, req *agent.ExecSQLReq) (*agent.ExecSQLResp, error) {
			require.True(t, strings.HasPrefix(req.GetSql(), "IMPORT INTO "))
			jobID := int64(1001)
			if strings.Contains(req.GetSql(), "`t2`") {
				jobID = 1002
			}
			return &agent.ExecSQLResp{
				ColumnNames: []string{"Job_ID"},
				Rows: []*agent.Row{
					{
						TypedColumnValues: []*agent.ColumnValue{
							{Value: &agent.ColumnValue_Int64Value{Int64Value: jobID}},
						},
					},
				},
			}, nil
		},
	}

	cluster := &agent.Cluster{
		ClusterId:         1,
		User:              "root",
		PasswordEncrypted: true,
		EncryptedPassword: &agent.EncryptedData{Pem: []byte("pem"), Data: []byte("data")},
	}
	createResp, err := s.BatchCreateImportIntoJobTasks(context.Background(), &agent.BatchCreateImportIntoJobTasksReq{
		ClusterId: 1,
		Cluster:   cluster,
		Tasks: []*agent.BatchCreateImportIntoJobTasksReq_Task{
			{TaskId: "submit-task-1", Sql: "IMPORT INTO `db`.`t1` FROM 's3://bucket/path/*.csv' WITH DETACHED"},
			{TaskId: "submit-task-2", Sql: "IMPORT INTO `db`.`t2` FROM 's3://bucket/path/*.csv' WITH DETACHED"},
		},
	})
	require.NoError(t, err)
	require.Len(t, createResp.GetTasks(), 2)

	require.Eventually(t, func() bool {
		getResp, getErr := s.BatchGetImportIntoJobTasks(context.Background(), &agent.BatchGetImportIntoJobTasksReq{
			TaskIds: []string{"submit-task-1", "missing-task", "submit-task-2"},
		})
		require.NoError(t, getErr)
		if len(getResp.GetTasks()) != 2 {
			return false
		}
		for _, task := range getResp.GetTasks() {
			if task.GetStatus() != agent.ImportIntoJobTaskStatus_IMPORT_INTO_JOB_TASK_STATUS_SUCCEEDED {
				return false
			}
		}
		require.Equal(t, "submit-task-1", getResp.GetTasks()[0].GetTaskId())
		require.Equal(t, int64(1001), getResp.GetTasks()[0].GetResult().GetJobId())
		require.Equal(t, "submit-task-2", getResp.GetTasks()[1].GetTaskId())
		require.Equal(t, int64(1002), getResp.GetTasks()[1].GetResult().GetJobId())
		return true
	}, time.Second, 10*time.Millisecond)
}

func TestBatchGetImportIntoJobTasksOmitsInternalErrors(t *testing.T) {
	s := &AgentServer{importTaskManager: newImportAsyncTaskManagerWithRetention(time.Hour)}

	_, created := s.importTaskManager.CreateIfAbsent(context.Background(), "submit-task-ok", func(context.Context) (any, error) {
		return &agent.GetImportIntoJobTaskResp_Result{JobId: 1001}, nil
	})
	require.True(t, created)

	_, created = s.importTaskManager.CreateIfAbsent(context.Background(), "submit-task-error", func(context.Context) (any, error) {
		return nil, errors.New("boom")
	})
	require.True(t, created)

	require.Eventually(t, func() bool {
		resp, err := s.BatchGetImportIntoJobTasks(context.Background(), &agent.BatchGetImportIntoJobTasksReq{
			TaskIds: []string{"submit-task-error", "submit-task-ok"},
		})
		require.NoError(t, err)
		if len(resp.GetTasks()) != 1 {
			return false
		}
		okTask := resp.GetTasks()[0]
		return okTask.GetTaskId() == "submit-task-ok" &&
			okTask.GetStatus() == agent.ImportIntoJobTaskStatus_IMPORT_INTO_JOB_TASK_STATUS_SUCCEEDED &&
			okTask.GetResult().GetJobId() == int64(1001)
	}, time.Second, 10*time.Millisecond)
}

func TestBatchCreateImportIntoJobTasksRejectsInvalidRequestBeforeCreating(t *testing.T) {
	calls := 0
	s := &AgentServer{
		importTaskManager: newImportAsyncTaskManagerWithRetention(time.Hour),
		execSQLOnClusterFn: func(ctx context.Context, req *agent.ExecSQLReq) (*agent.ExecSQLResp, error) {
			calls++
			return &agent.ExecSQLResp{}, nil
		},
	}

	_, err := s.BatchCreateImportIntoJobTasks(context.Background(), &agent.BatchCreateImportIntoJobTasksReq{
		Cluster: &agent.Cluster{ClusterId: 1},
		Tasks: []*agent.BatchCreateImportIntoJobTasksReq_Task{
			{TaskId: "submit-task-1", Sql: "IMPORT INTO `db`.`t1` FROM 's3://bucket/path/*.csv' WITH DETACHED"},
			{TaskId: "submit-task-2", Sql: "SELECT 1"},
		},
	})
	require.Error(t, err)
	require.Equal(t, 0, calls)
}

func TestBatchCancelImportIntoJobs(t *testing.T) {
	s := &AgentServer{
		importTaskManager: newImportAsyncTaskManagerWithRetention(time.Hour),
		execSQLOnClusterFn: func(ctx context.Context, req *agent.ExecSQLReq) (*agent.ExecSQLResp, error) {
			switch req.GetSql() {
			case "CANCEL IMPORT JOB 1001":
				return &agent.ExecSQLResp{}, nil
			case "CANCEL IMPORT JOB 1002":
				return &agent.ExecSQLResp{ErrorMsg: "job already finished"}, nil
			default:
				require.FailNow(t, "unexpected SQL", req.GetSql())
				return nil, nil
			}
		},
	}

	resp, err := s.BatchCancelImportIntoJobs(context.Background(), &agent.BatchCancelImportIntoJobsReq{
		ClusterId: 1,
		Cluster:   &agent.Cluster{ClusterId: 1},
		JobIds:    []uint64{1001, 1002},
	})
	require.NoError(t, err)
	require.Len(t, resp.GetResults(), 2)
	require.Equal(t, uint64(1001), resp.GetResults()[0].GetJobId())
	require.Empty(t, resp.GetResults()[0].GetErrorMsg())
	require.Equal(t, uint64(1002), resp.GetResults()[1].GetJobId())
	require.Equal(t, "job already finished", resp.GetResults()[1].GetErrorMsg())
}