package grpc

import (
	"context"
	"database/sql"
	"encoding/csv"
	"fmt"
	"io/ioutil"
	"os"
	"testing"

	"github.com/DATA-DOG/go-sqlmock"
	"github.com/go-sql-driver/mysql"
	"github.com/pingcap/tidb/pkg/errno"
	"github.com/pingcap/tidb/pkg/meta/model"
	"github.com/pingcap/tidb/pkg/parser/ast"
	pmysql "github.com/pingcap/tidb/pkg/parser/mysql"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"github.com/tidbcloud/dataflow-service/idl/pbgen/proto/agent"
	"github.com/tidbcloud/dataflow-service/internal/db/entity"
	fileparser "github.com/tidbcloud/dataflow-service/pkg/agent/file-parser"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
)

const (
	schema    = "test_db"
	table     = "test_table"
	tableName = "`test_db`.`test_table`"
)

func TestCheckTableEmpty(t *testing.T) {
	db, dbMock, err := sqlmock.New()
	require.NoError(t, err)
	ctx := context.Background()

	query := fmt.Sprintf("SELECT 1 FROM %s USE INDEX\\(\\) LIMIT 1", tableName)

	dbMock.ExpectQuery(query).WillReturnRows(sqlmock.NewRows([]string{"1"}).AddRow(1))
	err = checkTableEmpty(ctx, db, schema, table)
	require.Equal(t, err, status.Error(codes.PermissionDenied, "table(`test_db`.`test_table`) is not empty, please truncate table and retry it."))

	dbMock.ExpectQuery(query).WillReturnError(sql.ErrNoRows)
	err = checkTableEmpty(ctx, db, schema, table)
	require.NoError(t, err)

	dbMock.ExpectQuery(query).WillReturnError(&mysql.MySQLError{Number: errno.ErrNoSuchTable})
	err = checkTableEmpty(ctx, db, schema, table)
	require.NoError(t, err)
}

func TestCheckColumnCompatitable(t *testing.T) {
	db, dbMock, err := sqlmock.New()
	require.NoError(t, err)
	ctx := context.Background()

	mockDb := func() {
		showCreateTableSql := fmt.Sprintf("SHOW CREATE TABLE %s", tableName)
		createTableSql := "CREATE TABLE `test_table` (\n  `id` int(11) NOT NULL AUTO_INCREMENT,\n  `name` varchar(45) NOT NULL DEFAULT 'a',\n  `age` int(11) NOT NULL,\n `date_created` timestamp NOT NULL DEFAULT CURRENT_TIMESTAMP,\n address_info JSON,\n city VARCHAR(64) AS (JSON_UNQUOTE(JSON_EXTRACT(address_info, '$.city'))),\n PRIMARY KEY (`id`)\n) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_bin"
		dbMock.ExpectQuery(showCreateTableSql).WillReturnRows(sqlmock.NewRows([]string{"Table", "Create Table"}).AddRow(table, createTableSql))
		dbMock.ExpectQuery("SHOW VARIABLES LIKE 'sql_mode'").WillReturnRows(sqlmock.NewRows([]string{"Variable_name", "Value"}).AddRow("sql_mode", ""))
	}

	// age is not null, and it's not in csv file
	fileName := writeCsvFile(t, [][]string{
		{"id", "name", "date_created", "address_info", "city"},
		{"1", "a", "2020-01-01 00:00:00", `{"city": "beijing"}`, "beijing"},
		{"2", "b", "2020-01-01 00:00:00", `{"city": "shanghai"}`, "shanghai"},
	})
	//nolint: errcheck
	defer os.Remove(fileName)

	// *********** test no header, it will check in order ***********
	builder := fileparser.NewParserBuilder(ctx, entity.ImportDataFormatCsv, &agent.CsvConfig{
		Separator:       ",",
		Delimiter:       "\"",
		BackslashEscape: true,
		Header:          false,
	}, nil)

	mockDb()
	warnTable, err := checkColumnCompatitable(ctx, db, schema, table, fileName, builder)
	require.NoError(t, err)
	assert.Nil(t, warnTable)

	// *********** test header, it will check by column name ***********
	builder = fileparser.NewParserBuilder(ctx, entity.ImportDataFormatCsv, &agent.CsvConfig{
		Separator:       ",",
		Delimiter:       "\"",
		BackslashEscape: true,
		Header:          true,
	}, nil)

	mockDb()
	warnTable, err = checkColumnCompatitable(ctx, db, schema, table, fileName, builder)
	require.NoError(t, err)
	assert.NotNil(t, warnTable)
	assert.Equal(t, &warningTable{
		schema: schema,
		table:  table,
		cols:   []string{"age"},
	}, warnTable)
}

func TestColumnCanBeEmpty(t *testing.T) {
	db, dbMock, err := sqlmock.New()
	require.NoError(t, err)
	ctx := context.Background()

	showCreateTableSql := fmt.Sprintf("SHOW CREATE TABLE %s", tableName)
	createTableSql := "CREATE TABLE `test_table` (\n  `id` int(11) NOT NULL AUTO_INCREMENT,\n  `name` varchar(45) NOT NULL DEFAULT 'a',\n  `age` int(11) NOT NULL,\n `date_created` timestamp NOT NULL DEFAULT CURRENT_TIMESTAMP,\n address_info JSON,\n city VARCHAR(64) AS (JSON_UNQUOTE(JSON_EXTRACT(address_info, '$.city'))),\n PRIMARY KEY (`id`)\n) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_bin"
	dbMock.ExpectQuery(showCreateTableSql).WillReturnRows(sqlmock.NewRows([]string{"Table", "Create Table"}).AddRow(table, createTableSql))
	dbMock.ExpectQuery("SHOW VARIABLES LIKE 'sql_mode'").WillReturnRows(sqlmock.NewRows([]string{"Variable_name", "Value"}).AddRow("sql_mode", ""))

	tblInfo, err := getTableInfo(ctx, db, schema, table)
	require.NoError(t, err)

	// column0: `id` int(11) NOT NULL AUTO_INCREMENT
	assert.True(t, pmysql.HasAutoIncrementFlag(tblInfo.Columns[0].GetFlag()))
	// column1: `name` varchar(45) NOT NULL DEFAULT 'a'
	assert.NotNil(t, tblInfo.Columns[1].GetDefaultValue())
	// column2: `age` int(11) NOT NULL
	assert.True(t, pmysql.HasNotNullFlag(tblInfo.Columns[2].GetFlag()))
	// column3: `date_created` timestamp NOT NULL DEFAULT CURRENT_TIMESTAMP
	assert.NotNil(t, tblInfo.Columns[3].GetDefaultValue())
	// column4: address_info JSON
	// column5: city VARCHAR(64) AS (JSON_UNQUOTE(JSON_EXTRACT(address_info, '$.city')))
	assert.True(t, tblInfo.Columns[5].IsGenerated())

	expected := []bool{true, true, false, true, true, true}

	for i, exp := range expected {
		require.Equal(t, exp, columnCanBeEmpty(tblInfo.Columns[i]))
	}
}

func writeCsvFile(t *testing.T, data [][]string) string {
	tempFile, err := ioutil.TempFile("", "temp*.csv")
	require.NoError(t, err)

	fileWriter := csv.NewWriter(tempFile)
	for _, record := range data {
		require.NoError(t, fileWriter.Write(record))
	}
	fileWriter.Flush()
	require.NoError(t, fileWriter.Error())

	return tempFile.Name()
}

func TestFindNoExistCols(t *testing.T) {
	fileName := writeCsvFile(t, [][]string{
		{"Name", "Age"},
		{"John Doe", "25"},
		{"Jane Smith", "30"},
	})
	//nolint: errcheck
	defer os.Remove(fileName)

	cols := []*model.ColumnInfo{
		{
			Name: ast.NewCIStr("Name"),
		},
		{
			Name: ast.NewCIStr("Gender"),
		},
		{
			Name: ast.NewCIStr("Age"),
		},
	}

	ctx := context.Background()

	// ********** test no header, will find in order **********
	builder := fileparser.NewParserBuilder(ctx, entity.ImportDataFormatCsv, &agent.CsvConfig{
		Separator:       ",",
		Delimiter:       "\"",
		BackslashEscape: true,
		Header:          false,
	}, nil)

	fileParser, err := builder.Build(fmt.Sprintf("file://%s", fileName))
	require.NoError(t, err)
	require.NoError(t, fileParser.ReadRow())
	noExistCol, err := findNoExistCols(cols, fileParser)
	require.NoError(t, err)
	require.Equal(t, map[string]struct{}{
		"age": {},
	}, noExistCol)

	// ********** test header, will find by header **********
	builder = fileparser.NewParserBuilder(ctx, entity.ImportDataFormatCsv, &agent.CsvConfig{
		Separator:       ",",
		Delimiter:       "\"",
		BackslashEscape: true,
		Header:          true,
	}, nil)

	fileParser, err = builder.Build(fmt.Sprintf("file://%s", fileName))
	require.NoError(t, err)
	require.NoError(t, fileParser.ReadRow())
	noExistCol, err = findNoExistCols(cols, fileParser)
	require.NoError(t, err)
	require.Equal(t, map[string]struct{}{
		"gender": {},
	}, noExistCol)
}

func TestCheckFileColsInDb(t *testing.T) {
	fileName := writeCsvFile(t, [][]string{
		{"Name", "Age"},
		{"John Doe", "25"},
		{"Jane Smith", "30"},
	})
	//nolint: errcheck
	defer os.Remove(fileName)

	sourceUri := fmt.Sprintf("file://%s", fileName)

	testCases := []struct {
		name        string
		tableInfo   *model.TableInfo
		withHeader  bool
		expectedErr string
	}{
		{
			name: "happy path: column exist in db",
			tableInfo: &model.TableInfo{
				Name: ast.NewCIStr("test_table"),
				Columns: []*model.ColumnInfo{
					{
						Name:   ast.NewCIStr("Name"),
						State:  model.StatePublic,
						Offset: 0,
					},
					{
						Name:   ast.NewCIStr("Gender"),
						State:  model.StatePublic,
						Offset: 1,
					},
					{
						Name:   ast.NewCIStr("Age"),
						State:  model.StatePublic,
						Offset: 2,
					},
				},
			},
			withHeader:  true,
			expectedErr: "",
		},
		{
			name: "bad path(header): column not exist in db",
			tableInfo: &model.TableInfo{
				Name: ast.NewCIStr("test_table"),
				Columns: []*model.ColumnInfo{
					{
						Name:   ast.NewCIStr("Name"),
						State:  model.StatePublic,
						Offset: 0,
					},
					{
						Name:   ast.NewCIStr("Gender"),
						State:  model.StatePublic,
						Offset: 1,
					},
				},
			},
			withHeader:  true,
			expectedErr: fmt.Sprintf("rpc error: code = PermissionDenied desc = The column(s) age in file %s does not exist in table test_table. Please check CSV file header line and retry it. If the CSV file has a header line, ensure column names of the header line match the target table. If the file doesn't have a header line, disable \"With Header\" in \"CSV configuration\".", sourceUri),
		},
		{
			name: "bad path(no header): column not exist in db",
			tableInfo: &model.TableInfo{
				Name: ast.NewCIStr("test_table"),
				Columns: []*model.ColumnInfo{
					{
						Name:   ast.NewCIStr("Name"),
						State:  model.StatePublic,
						Offset: 0,
					},
				},
			},
			withHeader:  false,
			expectedErr: fmt.Sprintf("rpc error: code = PermissionDenied desc = The column count in file %s(2) exceeds the table test_table(1). Please provide valid files and retry again.", sourceUri),
		},
	}

	ctx := context.Background()
	for _, tc := range testCases {
		builder := fileparser.NewParserBuilder(ctx, entity.ImportDataFormatCsv, &agent.CsvConfig{
			Separator:       ",",
			Delimiter:       "\"",
			BackslashEscape: true,
			Header:          tc.withHeader,
		}, nil)
		fileParser, err := builder.Build(sourceUri)
		require.NoError(t, err)
		require.NoError(t, fileParser.ReadRow())
		err = checkFileColsInDb(tc.tableInfo, fileParser, sourceUri)
		if tc.expectedErr == "" {
			require.NoError(t, err, tc.name)
		} else {
			require.EqualError(t, err, tc.expectedErr, tc.name)
		}
	}
}