package pimport

import (
	"fmt"
	"os"
	"sort"
	"testing"
	"time"

	"github.com/docker/go-units"
	"github.com/pingcap/tidb/pkg/parser"
	"github.com/pingcap/tidb/pkg/parser/ast"
	"github.com/pingcap/tidb/pkg/parser/mysql"
	"github.com/pingcap/tidb/pkg/types"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	agentpb "github.com/tidbcloud/dataflow-service/idl/pbgen/proto/agent"
)

func TestInferDataType(t *testing.T) {
	testcases := []struct {
		input    types.Datum
		expected byte
	}{
		{
			input:    types.NewStringDatum("2542"),
			expected: mysql.TypeLonglong,
		},
		{
			input:    types.NewStringDatum("2542.645"),
			expected: mysql.TypeDouble,
		},
		{
			input:    types.NewStringDatum("2011-11-10 11:11:11.111111"),
			expected: mysql.TypeDatetime,
		},
		{
			input:    types.NewStringDatum("2011-11-10"),
			expected: mysql.TypeDate,
		},
		{
			input:    types.NewStringDatum("11:11:11"),
			expected: mysql.TypeDuration,
		},
		{
			input:    types.NewStringDatum(`{"a": "9223372036854775809"}`),
			expected: mysql.TypeJSON,
		},
		{
			input:    types.NewStringDatum("hello, world"),
			expected: mysql.TypeVarchar,
		},
	}

	for _, tc := range testcases {
		s, err := tc.input.ToString()
		assert.NoError(t, err)
		ft := inferDataType(tc.input, s)
		assert.Equal(t, tc.expected, ft.GetType())
		if tc.expected == mysql.TypeVarchar {
			assert.Equal(t, 512, ft.GetFlen())
		}
	}
}

func TestDefaultDataType(t *testing.T) {
	expected := []string{
		"BIT(1)",
		"TEXT",
		"DATE",
		"DATETIME",
		"FLOAT",
		"DOUBLE",
		"MEDIUMINT(9)",
		"JSON",
		"INT(11)",
		"BIGINT(20)",
		"LONGTEXT",
		"MEDIUMTEXT",
		"SMALLINT(6)",
		"TIME",
		"TIMESTAMP",
		"TINYINT(4)",
		"TINYTEXT",
		"YEAR(4)",
		"TINYINT(1)",
		"DECIMAL(32,4)",
		"CHAR(32)",
		"VARCHAR(64)",
		"VARCHAR(256)",
		"VARCHAR(512)",
		"VECTOR",
		"BINARY(32)",
		"VARBINARY(64)",
		"VARBINARY(256)",
		"VARBINARY(512)",
	}
	sort.StringSlice(expected).Sort()
	for i, s := range GetDefaultDataType() {
		assert.Equal(t, expected[i], s)
	}
}

func TestCheckColumn(t *testing.T) {
	testCases := []struct {
		createSql       string
		expectedPkuk    bool
		expectedAutoGen bool
	}{
		// pk + auto_increment
		{
			createSql: `CREATE TABLE example_table (
				id INT NOT NULL AUTO_INCREMENT,
				PRIMARY KEY (id)
			  );`,
			expectedPkuk:    true,
			expectedAutoGen: true,
		},
		{
			createSql: `CREATE TABLE example_table (
				id INT NOT NULL PRIMARY KEY AUTO_INCREMENT
			  );`,
			expectedPkuk:    true,
			expectedAutoGen: true,
		},
		{
			createSql: `CREATE TABLE example_table (
				id INT NOT NULL AUTO_INCREMENT
			  );`,
			expectedPkuk:    false,
			expectedAutoGen: true,
		},
		{
			createSql: `CREATE TABLE example_table (
				id INT NOT NULL PRIMARY KEY
			  );`,
			expectedPkuk:    true,
			expectedAutoGen: false,
		},
		{
			createSql: `CREATE TABLE example_table (
				id INT NOT NULL,
				PRIMARY KEY (id)
			  );`,
			expectedPkuk:    true,
			expectedAutoGen: false,
		},
		// pk + auto_random
		{
			createSql: `CREATE TABLE example_table (
				id INT NOT NULL AUTO_RANDOM,
				PRIMARY KEY (id)
			  );`,
			expectedPkuk:    true,
			expectedAutoGen: true,
		},
		{
			createSql: `CREATE TABLE example_table (
				id INT NOT NULL PRIMARY KEY AUTO_RANDOM
			  );`,
			expectedPkuk:    true,
			expectedAutoGen: true,
		},
		{
			createSql: `CREATE TABLE example_table (
				id INT NOT NULL AUTO_RANDOM
			  );`,
			expectedPkuk:    false,
			expectedAutoGen: true,
		},
		{
			createSql: `CREATE TABLE example_table (
				id INT NOT NULL PRIMARY KEY
			  );`,
			expectedPkuk:    true,
			expectedAutoGen: false,
		},
		{
			createSql: `CREATE TABLE example_table (
				id INT NOT NULL,
				PRIMARY KEY (id)
			  );`,
			expectedPkuk:    true,
			expectedAutoGen: false,
		},
		// uk + auto_increment
		{
			createSql: `CREATE TABLE example_table (
				id INT NOT NULL AUTO_INCREMENT,
				UNIQUE KEY (id)
			  );`,
			expectedPkuk:    true,
			expectedAutoGen: true,
		},
		{
			createSql: `CREATE TABLE example_table (
				id INT NOT NULL UNIQUE KEY AUTO_INCREMENT
			  );`,
			expectedPkuk:    true,
			expectedAutoGen: true,
		},
		{
			createSql: `CREATE TABLE example_table (
				id INT NOT NULL AUTO_INCREMENT
			  );`,
			expectedPkuk:    false,
			expectedAutoGen: true,
		},
		{
			createSql: `CREATE TABLE example_table (
				id INT NOT NULL UNIQUE KEY
			  );`,
			expectedPkuk:    true,
			expectedAutoGen: false,
		},
		{
			createSql: `CREATE TABLE example_table (
				id INT NOT NULL,
				UNIQUE KEY (id)
			  );`,
			expectedPkuk:    true,
			expectedAutoGen: false,
		},
		// uk + auto_random
		{
			createSql: `CREATE TABLE example_table (
				id INT NOT NULL AUTO_RANDOM,
				UNIQUE KEY (id)
			  );`,
			expectedPkuk:    true,
			expectedAutoGen: true,
		},
		{
			createSql: `CREATE TABLE example_table (
				id INT NOT NULL UNIQUE KEY AUTO_RANDOM
			  );`,
			expectedPkuk:    true,
			expectedAutoGen: true,
		},
		{
			createSql: `CREATE TABLE example_table (
				id INT NOT NULL AUTO_RANDOM
			  );`,
			expectedPkuk:    false,
			expectedAutoGen: true,
		},
		{
			createSql: `CREATE TABLE example_table (
				id INT NOT NULL UNIQUE KEY
			  );`,
			expectedPkuk:    true,
			expectedAutoGen: false,
		},
		{
			createSql: `CREATE TABLE example_table (
				id INT NOT NULL,
				UNIQUE KEY (id)
			  );`,
			expectedPkuk:    true,
			expectedAutoGen: false,
		},

		// VECTOR
		{
			createSql: `CREATE TABLE example_table (
				id INT,
				embedding VECTOR(3)
			  );`,
			expectedPkuk:    false,
			expectedAutoGen: false,
		},
		{
			createSql: `CREATE TABLE example_table (
				id INT,
				embedding VECTOR
			  );`,
			expectedPkuk:    false,
			expectedAutoGen: false,
		},

		// noop
		{
			createSql: `CREATE TABLE example_table (
				id INT
			  );`,
			expectedPkuk:    false,
			expectedAutoGen: false,
		},
	}

	for i, tc := range testCases {
		stmt, err := parser.New().ParseOneStmt(tc.createSql, "", "")
		require.NoError(t, err)

		createTableStmt, ok := stmt.(*ast.CreateTableStmt)
		require.True(t, ok)

		col := createTableStmt.Cols[0]
		require.Equal(t, tc.expectedPkuk, columnIsPkuk(col, createTableStmt.Constraints), fmt.Sprintf("%d case column: %s", i+1, col.Name))
		require.Equal(t, tc.expectedAutoGen, columnIsAutoGen(col), fmt.Sprintf("%d case column: %s", i+1, col.Name))
	}
}

func TestEstimateProgress(t *testing.T) {
	testCases := []struct {
		timeBeforeNow   time.Duration
		fileSize        uint32
		expectedPercent uint32
	}{
		{
			timeBeforeNow:   5 * time.Second,
			fileSize:        uint32(10 * units.MiB),
			expectedPercent: 50,
		},
		{
			timeBeforeNow:   5 * time.Second,
			fileSize:        uint32(50 * units.MiB),
			expectedPercent: 10,
		},
		{
			timeBeforeNow:   10 * time.Second,
			fileSize:        uint32(10 * units.MiB),
			expectedPercent: 99,
		},
	}

	for _, tc := range testCases {
		require.Equal(t, estimateProgress(time.Now().Add(-tc.timeBeforeNow), tc.fileSize), tc.expectedPercent)
	}
}

func TestGetLineTerminator(t *testing.T) {
	testCases := []struct {
		name     string
		filePath string
		expected string
	}{
		{
			name:     "Windows line endings",
			filePath: "testdata/windows_line_endings.csv",
			expected: "\\r\\n",
		},
		{
			name:     "Unix line endings",
			filePath: "testdata/unix_line_endings.csv",
			expected: "\\n",
		},
		{
			name:     "One line",
			filePath: "testdata/one_line.csv",
			expected: "\\n",
		},
	}

	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			file, err := os.Open(tc.filePath)
			require.NoError(t, err)
			result, err := getLineTerminator(file)
			require.Equal(t, tc.expected, result)
			require.NoError(t, err)
			require.NoError(t, file.Close())
		})
	}
}

func TestReorderDataColumns(t *testing.T) {
	testCases := []struct {
		rowData        []*agentpb.Row
		sourceColNames []string
		targetColNames []string
		expected       []*agentpb.Row
	}{
		// normal case: ordered column or the target table does not exist
		{
			rowData: []*agentpb.Row{
				{ColumnValues: []string{"0", "1", "2"}},
			},
			sourceColNames: []string{"column1", "column2", "column3"},
			targetColNames: []string{"column1", "column2", "column3"},
			expected: []*agentpb.Row{
				{ColumnValues: []string{"0", "1", "2"}},
			},
		},
		// upper case column names
		{
			rowData: []*agentpb.Row{
				{ColumnValues: []string{"0", "1", "2"}},
			},
			sourceColNames: []string{"column1", "column2", "column3"},
			targetColNames: []string{"COLUMN2", "COLUMN1", "COLUMN3"},
			expected: []*agentpb.Row{
				{ColumnValues: []string{"1", "0", "2"}},
			},
		},
		// disordered column
		{
			rowData: []*agentpb.Row{
				{ColumnValues: []string{"0", "1", "2"}},
				{ColumnValues: []string{"3", "4", "5"}},
				{ColumnValues: []string{"6", "7", "8"}},
			},
			sourceColNames: []string{"column1", "column2", "column3"},
			targetColNames: []string{"column1", "column3", "column2"},
			expected: []*agentpb.Row{
				{ColumnValues: []string{"0", "2", "1"}},
				{ColumnValues: []string{"3", "5", "4"}},
				{ColumnValues: []string{"6", "8", "7"}},
			},
		},
		// the column number of target table is greater than source data
		{
			rowData: []*agentpb.Row{
				{ColumnValues: []string{"0", "1", "2"}},
			},
			sourceColNames: []string{"column1", "column4", "column3"},
			targetColNames: []string{"column1", "column2", "column3", "column4"},
			expected: []*agentpb.Row{
				{ColumnValues: []string{"0", "", "2", "1"}},
			},
		},
	}

	for i, tc := range testCases {
		require.Equal(t, tc.expected, reorderDataColumns(tc.rowData, tc.sourceColNames, tc.targetColNames), fmt.Sprintf("case #%d", i))
	}
}