package server

import (
	"bytes"
	"context"
	"fmt"
	"os"
	"path/filepath"
	"testing"
	"time"

	gomock "github.com/golang/mock/gomock"
	"github.com/pingcap/errors"
	"github.com/spf13/viper"
	"github.com/stretchr/testify/require"
	"gopkg.in/natefinch/lumberjack.v2"
	"k8s.io/apimachinery/pkg/util/rand"
)

func TestConfigToYAML(t *testing.T) {
	cfg := Config{
		LogDir:                "dir",
		BucketURL:             "url",
		Region:                "region",
		RoleARN:               "arn",
		ExternalID:            "id",
		SlowLogMaxBackups:     1,
		CleanOldestLogPercent: 0.8,
	}

	data, err := cfg.ToYAML()
	t.Log("yaml: ", string(data))
	require.Nil(t, err)
	vp := viper.New()
	vp.SetConfigType("yaml")
	err = vp.ReadConfig(bytes.NewReader(data))
	require.Nil(t, err)
	cfgBack := new(Config)
	err = vp.Unmarshal(cfgBack)
	require.Nil(t, err)
	require.Equal(t, cfg, *cfgBack)
}

func TestListLog(t *testing.T) {
	dir := t.TempDir()
	name := "testlist.log"

	items, err := listLog(dir, name)
	require.Nil(t, err)
	require.Len(t, items, 0)

	lg := lumberjack.Logger{
		Filename: filepath.Join(dir, name),
	}
	_, err = lg.Write([]byte("a"))
	require.Nil(t, err)
	err = lg.Rotate()
	require.Nil(t, err)
	time.Sleep(time.Millisecond) // avoid may use same backup name
	_, err = lg.Write([]byte("b"))
	require.Nil(t, err)
	err = lg.Rotate()
	require.Nil(t, err)

	items, err = listLog(dir, name)
	require.Nil(t, err)
	require.Len(t, items, 2)
	data, err := os.ReadFile(filepath.Join(dir, items[0].name))
	require.Nil(t, err)
	require.Equal(t, []byte("a"), data)
	data, err = os.ReadFile(filepath.Join(dir, items[1].name))
	require.Nil(t, err)
	require.Equal(t, []byte("b"), data)
}

func TestListLogNoExt(t *testing.T) {
	dir := t.TempDir()
	name := "slowlog"
	t.Log("use TempDir: ", dir)

	items, err := listLog(dir, name)
	require.Nil(t, err)
	require.Len(t, items, 0)

	lg := lumberjack.Logger{
		Filename: filepath.Join(dir, name),
	}
	_, err = lg.Write([]byte("a"))
	require.Nil(t, err)
	err = lg.Rotate()
	require.Nil(t, err)
	time.Sleep(time.Millisecond) // avoid may use same backup name
	_, err = lg.Write([]byte("b"))
	require.Nil(t, err)
	err = lg.Rotate()
	require.Nil(t, err)

	items, err = listLog(dir, name)
	require.Nil(t, err)
	require.Len(t, items, 2)
	data, err := os.ReadFile(filepath.Join(dir, items[0].name))
	require.Nil(t, err)
	require.Equal(t, []byte("a"), data)
	data, err = os.ReadFile(filepath.Join(dir, items[1].name))
	require.Nil(t, err)
	require.Equal(t, []byte("b"), data)
}

func TestRotateIfNeed(t *testing.T) {
	dir := t.TempDir()
	name := "test.log"

	reset := func(t *testing.T) {
		t.Helper()

		// delete all files under dir.
		logs, err := listLog(dir, name)
		require.Nil(t, err)
		for _, item := range logs {
			err := os.Remove(filepath.Join(dir, item.name))
			require.Nil(t, err)
		}
		err = os.Remove(filepath.Join(dir, name))
		if !os.IsNotExist(err) {
			require.Nil(t, err)
		}
	}

	t.Run("test not exist", func(t *testing.T) {
		reset(t)

		err := rotateIfNeed(dir, name)
		require.Nil(t, err)
		logs, err := listLog(dir, name)
		require.Nil(t, err)
		require.Len(t, logs, 0)
	})

	t.Run("test empty file", func(t *testing.T) {
		reset(t)

		fname := filepath.Join(dir, name)
		err := os.WriteFile(fname, nil, 0600)
		require.Nil(t, err)
		logs, err := listLog(dir, name)
		require.Nil(t, err)
		require.Len(t, logs, 0)
	})

	t.Run("test will rotate", func(t *testing.T) {
		reset(t)

		fname := filepath.Join(dir, name)
		err := os.WriteFile(fname, []byte("data"), 0600)
		require.Nil(t, err)

		err = rotateIfNeed(dir, name)
		require.Nil(t, err)
		logs, err := listLog(dir, name)
		require.Nil(t, err)
		require.Len(t, logs, 1)

		fname = filepath.Join(dir, logs[0].name)
		readData, err := os.ReadFile(fname)
		require.Nil(t, err)
		require.Equal(t, []byte("data"), readData)
	})
}

func TestServer(t *testing.T) {
	var dir string
	var audit *Server
	var vp *viper.Viper
	var mockHandler *MockLogHandler
	var configFile *os.File
	defer func() {
		if configFile != nil {
			os.Remove(configFile.Name())
		}
	}()

	reset := func(t *testing.T, yamlConfig string) {
		t.Helper()
		mockCtrl := gomock.NewController(t)
		mockHandler = NewMockLogHandler(mockCtrl)
		defaultLogHandler = mockHandler

		dir = filepath.Join(t.TempDir(), rand.String(5))
		var err error
		err = os.Mkdir(dir, 0777)
		require.Nil(t, err)

		if configFile != nil {
			configFile.Close()
			os.Remove(configFile.Name())
		}
		configFile, err = os.CreateTemp(dir, "*.yaml")
		require.Nil(t, err)
		data := fmt.Sprintf(`
logDir: %s
`, dir)
		data += yamlConfig
		_, err = configFile.WriteString(data)
		require.Nil(t, err)

		vp = viper.New()
		vp.SetConfigFile(configFile.Name())
		audit, err = New(vp)
		require.Nil(t, err)
		audit.skipCheckConn = true
	}

	t.Run("load config", func(t *testing.T) {
		reset(t, "")

		require.Equal(t, dir, audit.LogDir)
		// defaults in New()
		require.Equal(t, 0.8, audit.CleanOldestLogPercent)
		require.Equal(t, 1, audit.SlowLogMaxBackups)

		// test setting all config items.
		data := `
logDir: /tmp
bucketURL: url
region: region
roleARN: roleARN
externalID: externalID
slowLogMaxBackups: 10
cleanOldestLogPercent: 0.5
`
		_, err := configFile.WriteString(data)
		require.Nil(t, err)
		err = vp.ReadInConfig()
		require.Nil(t, err)
		audit.loadConfig()
		require.Equal(t, audit.LogDir, "/tmp")
		require.Equal(t, audit.BucketURL, "url")
		require.Equal(t, audit.Region, "region")
		require.Equal(t, audit.RoleARN, "roleARN")
		require.Equal(t, audit.ExternalID, "externalID")
		require.Equal(t, audit.SlowLogMaxBackups, 10)
		require.Equal(t, audit.CleanOldestLogPercent, 0.5)
	})

	t.Run("precheck", func(t *testing.T) {
		reset(t, "")
		audit.BucketURL = "s3://name/a/b"

		// test success push
		mockHandler.EXPECT().Push(
			gomock.Any(),
			gomock.Any(),
			gomock.Eq("a/b/"+checkFileName),
		).Return(nil)
		pass, _, err := audit.Precheck(context.Background())
		require.Nil(t, err)
		require.True(t, pass)

		// test fail to push
		mockHandler.EXPECT().Push(
			gomock.Any(),                    // ctx
			gomock.Any(),                    // fname
			gomock.Eq("a/b/"+checkFileName), // targetKey
		).Return(fmt.Errorf("no permission"))
		pass, reason, err := audit.Precheck(context.Background())
		require.Nil(t, err)
		require.False(t, pass)
		require.Equal(t, "no permission", reason)
	})

	t.Run("fail_push", func(t *testing.T) {
		reset(t, "cleanOldestLogPercent: 1.0")

		// write two rotated file and one current.
		writeLog(t, filepath.Join(dir, auditLogFileName), []string{"a", "b", "c"})
		items, err := listLog(dir, auditLogFileName)
		require.Nil(t, err)
		require.Len(t, items, 2)

		// set up fail to push
		callCh := make(chan struct{})
		mockHandler.EXPECT().Push(
			gomock.Any(),
			gomock.Any(),
			gomock.All(),
		).MinTimes(1).DoAndReturn(func(ctx context.Context, fname string, targetKey string) error {
			select {
			case callCh <- struct{}{}:
			default:
			}
			return errors.New("some error")
		})

		// run the audit
		ctx, cancel := context.WithCancel(context.Background())
		defer cancel()
		var runErr error
		runDone := make(chan struct{})
		go func() {
			defaultCheckLogDuration = time.Millisecond * 100
			defaultCleanLogDuration = time.Millisecond * 100
			runErr = audit.Run(ctx)
			close(runDone)
		}()

		// check files will not be deleted since it fail to push
		<-callCh
		<-callCh // make sure already retry
		items, err = listLog(dir, auditLogFileName)
		require.Nil(t, err)
		require.Len(t, items, 2)

		// test clean backup logs and slow logs
		// setting cleanOldestLogPercent as 0 to make it delete all rotated files.
		str := `
logDir: %s
cleanOldestLogPercent: 0.0
`
		data := fmt.Sprintf(str, dir)
		err = os.WriteFile(configFile.Name(), []byte(data), 0600)
		require.Nil(t, err)

		require.Equal(t, 1, audit.SlowLogMaxBackups)
		writeLog(t, filepath.Join(dir, slowLogFileName), []string{"a", "b", "c"})

		// wait it cleanup
		for i := 0; i < 100; i++ {
			items, err = listLog(dir, auditLogFileName)
			itemsSlow, errSlow := listLog(dir, slowLogFileName)
			if err == nil && len(items) == 0 &&
				errSlow == nil && len(itemsSlow) == 1 {
				break
			}
			time.Sleep(time.Millisecond * 100)
		}
		items, err = listLog(dir, auditLogFileName)
		require.Nil(t, err)
		require.Len(t, items, 0)
		items, err = listLog(dir, slowLogFileName)
		require.Nil(t, err)
		require.Len(t, items, 1)

		cancel()
		<-runDone
		require.NotNil(t, runErr)
	})

	t.Run("normal_push", func(t *testing.T) {
		reset(t, "cleanOldestLogPercent: 1.0")

		// write two rotated file and on current.
		writeLog(t, filepath.Join(dir, auditLogFileName), []string{"a", "b", "c"})
		items, err := listLog(dir, auditLogFileName)
		require.Nil(t, err)
		require.Len(t, items, 2)

		// set up fail to push
		mockHandler.EXPECT().Push(
			gomock.Any(),
			gomock.Any(),
			gomock.All(),
		).Return(nil).MinTimes(3).MaxTimes(3)

		ctx, cancel := context.WithCancel(context.Background())
		defer cancel()
		var runErr error
		runDone := make(chan struct{})
		go func() {
			defaultCheckLogDuration = time.Millisecond * 100
			defaultCleanLogDuration = time.Millisecond * 100
			runErr = audit.Run(ctx)
			close(runDone)
		}()

		// wait it pushed the 2 rotated files
		for i := 0; i < 100; i++ {
			items, err = listLog(dir, auditLogFileName)
			if err == nil && len(items) == 0 {
				break
			}
			time.Sleep(time.Millisecond * 100)
		}
		items, err = listLog(dir, auditLogFileName)
		require.Nil(t, err)
		require.Len(t, items, 0)

		cancel()
		<-runDone
		require.Nil(t, runErr)
	})
}

func writeLog(t *testing.T, fname string, files []string) {
	t.Helper()

	lg := lumberjack.Logger{
		Filename: fname,
	}

	for idx, data := range files {
		if idx != 0 {
			time.Sleep(time.Millisecond) // avoid it using same name
			err := lg.Rotate()
			require.Nil(t, err)
		}
		_, err := lg.Write([]byte(data))
		require.Nil(t, err)
	}
	lg.Close()
}

func TestAppendOrdinals(t *testing.T) {
	get, err := appendOrdinals("s3://host/", "tidb-0")
	require.Nil(t, err)
	require.Equal(t, "s3://host/tidb-0", get)

	get, err = appendOrdinals("s3://host/pre", "tidb-0")
	require.Nil(t, err)
	require.Equal(t, "s3://host/pre/tidb-0", get)

	get, err = appendOrdinals("s3://host/", "tidb-notid")
	require.Nil(t, err)
	require.Equal(t, "s3://host/", get)

	get, err = appendOrdinals("s3://host/pre", "tidb-notid")
	require.Nil(t, err)
	require.Equal(t, "s3://host/pre", get)
}