package precheck

import (
	"context"
	"fmt"
	"testing"
	"time"

	"github.com/stretchr/testify/require"
	"github.com/stretchr/testify/suite"
	"github.com/tidbcloud/pkgv2/bizerr"
	"github.com/tidbcloud/pkgv2/import/cloud_storage"
	"github.com/tidbcloud/pkgv2/log"
	"go.uber.org/mock/gomock"
	"go.uber.org/zap"
)

type CheckSchemasTestSuite struct {
	suite.Suite
	ctrl          *gomock.Controller
	store         *cloud_storage.MockCloudStorage
	fileRouter    *MockFileRouter
	schemaChecker *MockSchemaChecker
	logger        *zap.Logger
	checker       *Checker
}

func (suite *CheckSchemasTestSuite) SetupTest() {
	suite.ctrl = gomock.NewController(suite.T())
	suite.store = cloud_storage.NewMockCloudStorage(suite.ctrl)
	suite.fileRouter = NewMockFileRouter(suite.ctrl)
	suite.schemaChecker = NewMockSchemaChecker(suite.ctrl)
	suite.logger = log.WithContextL(context.Background())
	suite.checker = &Checker{
		store:         suite.store,
		fileRouter:    suite.fileRouter,
		schemaChecker: suite.schemaChecker,
	}
}

func (suite *CheckSchemasTestSuite) TearDownTest() {
	suite.ctrl.Finish()
}

func (suite *CheckSchemasTestSuite) TestCheckSchemas_Success() {
	assert := require.New(suite.T())
	ctx := context.Background()

	// Create mock data
	dbMap := map[string]*DbMeta{
		"test": {
			Name: "test",
			TableMap: map[string]*TableMeta{
				"table1": {
					Name: "table1",
					SchemaFiles: []RouteResult{{
						Table:    Table{Schema: "test", Name: "table1"},
						Type:     TableSchema,
						FilePath: "test_db/test.table1-schema.sql",
					}},
					DataFiles: []RouteResult{{
						Table:    Table{Schema: "test", Name: "table1"},
						Type:     TypeSQL,
						FilePath: "test_db/test.table1.sql",
					}},
				},
				"table2": {
					Name: "table2",
					SchemaFiles: []RouteResult{{
						Table:    Table{Schema: "test", Name: "table2"},
						Type:     TableSchema,
						FilePath: "test_db/test.table2-schema.sql",
					}},
					DataFiles: []RouteResult{{
						Table:    Table{Schema: "test", Name: "table2"},
						Type:     TypeSQL,
						FilePath: "test_db/test.table2.sql",
					}},
				},
			},
		},
	}

	// Setup expectations - all queries will succeed
	suite.schemaChecker.EXPECT().
		CheckSchema(gomock.Any(), "test", "table1").
		Return(&SchemaCheckResult{Exists: false, Empty: false}, nil)

	suite.schemaChecker.EXPECT().
		CheckSchema(gomock.Any(), "test", "table2").
		Return(&SchemaCheckResult{Exists: false, Empty: false}, nil)

	// Call the method
	results, bizErr := suite.checker.checkSchemas(ctx, dbMap, false)

	// Verify results
	assert.NoError(bizErr)
	assert.ElementsMatch([]TableResult{
		{
			Table:         Table{Schema: "test", Name: "table1"},
			DataFiles:     []string{"test_db/test.table1.sql"},
			ErrorMsg:      "",
			UseSchemaFile: true,
			SchemaFile:    "test_db/test.table1-schema.sql",
		},
		{
			Table:         Table{Schema: "test", Name: "table2"},
			DataFiles:     []string{"test_db/test.table2.sql"},
			ErrorMsg:      "",
			UseSchemaFile: true,
			SchemaFile:    "test_db/test.table2-schema.sql",
		},
	}, results.Results)
	assert.False(results.IsCancelled)
}

func (suite *CheckSchemasTestSuite) TestCheckSchemas_ContextCanceled() {
	assert := require.New(suite.T())

	// Create a context that will be canceled
	ctx, cancel := context.WithCancel(context.Background())

	// Create mock data
	dbMap := map[string]*DbMeta{
		"test": {
			Name: "test",
			TableMap: map[string]*TableMeta{
				"table1": {
					Name: "table1",
					SchemaFiles: []RouteResult{{
						Table:    Table{Schema: "test", Name: "table1"},
						Type:     TableSchema,
						FilePath: "test_db/test.table1-schema.sql",
					}},
				},
				"table2": {
					Name: "table2",
					SchemaFiles: []RouteResult{{
						Table:    Table{Schema: "test", Name: "table2"},
						Type:     TableSchema,
						FilePath: "test_db/test.table2-schema.sql",
					}},
				},
			},
		},
	}

	// Setup expectation for table1 - it will be called before context cancels
	suite.schemaChecker.EXPECT().
		CheckSchema(gomock.Any(), "test", "table1").
		DoAndReturn(func(ctx context.Context, schema, table string) (*SchemaCheckResult, bizerr.BizError) {
			return &SchemaCheckResult{Exists: true, Empty: true}, nil
		})

	// Setup expectation for table2 - it might or might not be called depending on timing
	suite.schemaChecker.EXPECT().
		CheckSchema(gomock.Any(), "test", "table2").
		DoAndReturn(func(ctx context.Context, schema, table string) (*SchemaCheckResult, bizerr.BizError) {
			// Sleep to allow context to be canceled
			time.Sleep(1 * time.Second)
			cancel()
			return nil, bizerr.New(bizerr.ErrCodeTimeout, "")
		}).MaxTimes(1)

	// Call the method
	results, bizErr := suite.checker.checkSchemas(ctx, dbMap, false)

	// Verify results - should have results from table1 that completed before context was canceled
	assert.NoError(bizErr)
	assert.ElementsMatch([]TableResult{
		{
			Table:         Table{Schema: "test", Name: "table1"},
			DataFiles:     []string{},
			ErrorMsg:      "",
			UseSchemaFile: false,
			SchemaFile:    "test_db/test.table1-schema.sql",
		},
	}, results.Results)
	assert.True(results.IsCancelled)
}

func (suite *CheckSchemasTestSuite) TestCheckSchemas_Error() {
	assert := require.New(suite.T())
	ctx := context.Background()

	// Create mock data
	dbMap := map[string]*DbMeta{
		"test": {
			Name: "test",
			TableMap: map[string]*TableMeta{
				"table1": {
					Name: "table1",
					SchemaFiles: []RouteResult{{
						Table:    Table{Schema: "test", Name: "table1"},
						Type:     TableSchema,
						FilePath: "test_db/test.table1-schema.sql",
					}},
				},
				"table2": {
					Name: "table2",
					SchemaFiles: []RouteResult{{
						Table:    Table{Schema: "test", Name: "table2"},
						Type:     TableSchema,
						FilePath: "test_db/test.table2-schema.sql",
					}},
				},
			},
		},
	}

	// Setup error expectation for table1
	suite.schemaChecker.EXPECT().
		CheckSchema(gomock.Any(), "test", "table1").
		Return(nil, bizerr.New(bizerr.ErrCodeUnknown, ""))

	// Set up expectation for table2, but it should not be called due to error handling
	suite.schemaChecker.EXPECT().
		CheckSchema(gomock.Any(), "test", "table2").
		Return(&SchemaCheckResult{Exists: true, Empty: true}, nil).
		MaxTimes(1)

	// Call the method - should return error
	results, bizErr := suite.checker.checkSchemas(ctx, dbMap, false)

	// Verify error is returned
	assert.Error(bizErr)
	assert.Equal(bizerr.ErrCodeUnknown, bizErr.RawErrCode())
	assert.Nil(results)
}

// Test for case where concurrent goroutines complete successfully
func TestCheckSchemas_Concurrency(t *testing.T) {
	assert := require.New(t)
	// Create a simple mock checker
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	// For this test we'll use a real errgroup that will demonstrate the concurrency
	store := cloud_storage.NewMockCloudStorage(ctrl)
	fileRouter := NewMockFileRouter(ctrl)
	schemaChecker := NewMockSchemaChecker(ctrl)

	checker := &Checker{
		store:         store,
		fileRouter:    fileRouter,
		schemaChecker: schemaChecker,
	}

	// Create a map with 20 tables to test concurrency (limit is set to 5)
	dbMap := map[string]*DbMeta{
		"test": {
			Name:     "test",
			TableMap: make(map[string]*TableMeta),
		},
	}

	// Add 20 tables to the map
	for i := 0; i < 20; i++ {
		tableName := fmt.Sprintf("table%d", i)
		dbMap["test"].TableMap[tableName] = &TableMeta{
			Name: tableName,
			DataFiles: []RouteResult{{
				Table:    Table{Schema: "test", Name: tableName},
				Type:     TypeSQL,
				FilePath: fmt.Sprintf("test_db/test.%s.sql", tableName),
			}},
		}

		// Set up expectations for each table
		schemaChecker.EXPECT().
			CheckSchema(gomock.Any(), "test", tableName).
			Return(&SchemaCheckResult{Exists: true, Empty: true}, nil)
	}

	// Call the method
	results, bizErr := checker.checkSchemas(context.Background(), dbMap, false)

	// Verify results
	assert.NoError(bizErr)
	assert.Len(results.Results, 20)
	assert.False(results.IsCancelled)
}

func TestCheckSchemasTestSuite(t *testing.T) {
	suite.Run(t, new(CheckSchemasTestSuite))
}