package fileparser

import (
	"compress/gzip"
	"context"
	"encoding/csv"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"testing"

	"github.com/pkg/errors"
	"github.com/stretchr/testify/require"
	"github.com/tidbcloud/dataflow-service/idl/pbgen/proto/agent"
	"github.com/tidbcloud/dataflow-service/internal/db/entity"
)

type (
	writeCsvFileFunc func(t *testing.T, data [][]string) string
	writeSqlFileFunc func(t *testing.T, data []string) string
)

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 writeCsvGzFile(t *testing.T, data [][]string) string {
	tempFile, err := ioutil.TempFile("", "temp*.csv.gz")
	require.NoError(t, err)

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

	require.NoError(t, gzipWriter.Close())

	return tempFile.Name()
}

func TestReadRow(t *testing.T) {
	runCsvReadRowTest(t, writeCsvFile)
	runCsvReadRowTest(t, writeCsvGzFile)

	runSqlReadRowTest(t, writeSqlFile)
	runSqlReadRowTest(t, writeSqlGzFile)
}

func TestCsv_BackslashEscape(t *testing.T) {
	ctx := context.Background()
	dir, err := os.Getwd()
	require.NoError(t, err)

	parser, err := NewCsvParser(ctx, fmt.Sprintf("file://%s/%s", dir, "backslash_escape.csv"), &agent.CsvConfig{
		Separator:       ",",
		Delimiter:       "\"",
		BackslashEscape: true,
		Header:          true,
	}, nil)
	require.NoError(t, err)

	// 1,"\"tom\"","2017-09-12"
	// 2,"\"bob\"","2022-11-23"
	writeRow, err := ReadRow(parser)
	require.NoError(t, err)
	require.Equal(t, &agent.Row{ColumnValues: []string{"1", "\"tom", "2017-09-12"}}, writeRow)
	writeRow, err = ReadRow(parser)
	require.NoError(t, err)
	require.Equal(t, &agent.Row{ColumnValues: []string{"2", "\"bob", "2022-11-23"}}, writeRow)
	_, err = ReadRow(parser)
	require.Equal(t, io.EOF, errors.Cause(err))
	require.NoError(t, parser.Close())
}

func runCsvReadRowTest(t *testing.T, write writeCsvFileFunc) {
	ctx := context.Background()

	fileName := write(t, [][]string{
		{"Name", "Age", "City"},
		{"John Doe", "25", "New York"},
		{"Jane Smith", "30", "San Francisco"},
	})
	// nolint: errcheck
	defer os.Remove(fileName)

	parser, err := NewCsvParser(ctx, fmt.Sprintf("file://%s", fileName), &agent.CsvConfig{
		Separator:       ",",
		Delimiter:       "\"",
		BackslashEscape: true,
		Header:          false,
	}, nil)
	require.NoError(t, err)

	writeRow, err := ReadRow(parser)
	require.NoError(t, err)
	require.Equal(t, &agent.Row{ColumnValues: []string{"Name", "Age", "City"}}, writeRow)
	writeRow, err = ReadRow(parser)
	require.NoError(t, err)
	require.Equal(t, &agent.Row{ColumnValues: []string{"John Doe", "25", "New York"}}, writeRow)
	writeRow, err = ReadRow(parser)
	require.NoError(t, err)
	require.Equal(t, &agent.Row{ColumnValues: []string{"Jane Smith", "30", "San Francisco"}}, writeRow)
	_, err = ReadRow(parser)
	require.Equal(t, io.EOF, errors.Cause(err))
	require.NoError(t, parser.Close())

	parser, err = NewCsvParser(ctx, fmt.Sprintf("file://%s", fileName), &agent.CsvConfig{
		Separator:       ",",
		Delimiter:       "\"",
		BackslashEscape: true,
		// test header
		Header: true,
	}, nil)
	require.NoError(t, err)

	writeRow, err = ReadRow(parser)
	require.NoError(t, err)
	require.Equal(t, &agent.Row{ColumnValues: []string{"John Doe", "25", "New York"}}, writeRow)
	writeRow, err = ReadRow(parser)
	require.NoError(t, err)
	require.Equal(t, &agent.Row{ColumnValues: []string{"Jane Smith", "30", "San Francisco"}}, writeRow)
	_, err = ReadRow(parser)
	require.Equal(t, io.EOF, errors.Cause(err))
	require.NoError(t, parser.Close())
}

func writeSqlFile(t *testing.T, queries []string) string {
	tempFile, err := ioutil.TempFile("", "temp*.sql")
	require.NoError(t, err)

	for _, query := range queries {
		_, err = tempFile.WriteString(query + "\n")
		require.NoError(t, err)
	}

	return tempFile.Name()
}

func writeSqlGzFile(t *testing.T, queries []string) string {
	tempFile, err := ioutil.TempFile("", "temp*.sql.gz")
	require.NoError(t, err)

	gzipWriter := gzip.NewWriter(tempFile)
	for _, query := range queries {
		_, err := gzipWriter.Write([]byte(query + "\n"))
		require.NoError(t, err)
	}
	require.NoError(t, gzipWriter.Close())

	return tempFile.Name()
}

func runSqlReadRowTest(t *testing.T, write writeSqlFileFunc) {
	ctx := context.Background()

	fileName := write(t, []string{
		"INSERT INTO table_name VALUES ('John Doe', '25', 'New York');",
		"INSERT INTO table_name VALUES ('Jane Smith', '30', 'San Francisco');",
	})
	// nolint: errcheck
	defer os.Remove(fileName)

	builder := NewParserBuilder(ctx, entity.ImportDataFormatSql, nil, nil)
	parser, err := builder.Build(fmt.Sprintf("file://%s", fileName))
	require.NoError(t, err)

	writeRow, err := ReadRow(parser)
	require.NoError(t, err)
	require.Equal(t, &agent.Row{ColumnValues: []string{"John Doe", "25", "New York"}}, writeRow)
	writeRow, err = ReadRow(parser)
	require.NoError(t, err)
	require.Equal(t, &agent.Row{ColumnValues: []string{"Jane Smith", "30", "San Francisco"}}, writeRow)
	_, err = ReadRow(parser)
	require.Equal(t, io.EOF, errors.Cause(err))
	require.NoError(t, parser.Close())
}
