/*
 *     Copyright 2023 The Dragonfly Authors
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package storage

import (
	"fmt"
	"io/fs"
	"io/ioutil"
	"os"
	"path/filepath"
	"reflect"
	"regexp"
	"testing"

	"github.com/gocarina/gocsv"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"

	schedulerstorage "d7y.io/dragonfly/v2/scheduler/storage"
)

var mockModelKey = "bar"

func TestStorage_New(t *testing.T) {
	tests := []struct {
		name    string
		baseDir string
		expect  func(t *testing.T, s Storage)
	}{
		{
			name:    "new storage",
			baseDir: os.TempDir(),
			expect: func(t *testing.T, s Storage) {
				assert := assert.New(t)
				assert.Equal(reflect.TypeOf(s).Elem().Name(), "storage")
			},
		},
	}

	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			tc.expect(t, New(tc.baseDir))
		})
	}
}

func TestStorage_ListDownload(t *testing.T) {
	require := require.New(t)
	testData, err := os.ReadFile("./testdata/download.csv")
	require.Nil(err, "load test file")

	tests := []struct {
		name    string
		baseDir string
		mock    func(t *testing.T, s Storage, baseDir, modelKey string, download []byte)
		expect  func(t *testing.T, s Storage, baseDir, modelKey string, download []byte)
	}{
		{
			name:    "empty csv file given",
			baseDir: os.TempDir(),
			mock: func(t *testing.T, s Storage, baseDir, modelKey string, download []byte) {
				file, err := os.OpenFile(filepath.Join(baseDir, "download-bar.csv"), os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0600)
				if err != nil {
					t.Fatal(err)
				}
				defer file.Close()
			},
			expect: func(t *testing.T, s Storage, baseDir, modelKey string, download []byte) {
				assert := assert.New(t)
				_, err := s.ListDownload(modelKey)
				assert.EqualError(err, "empty csv file given")
			},
		},
		{
			name:    "get file failed",
			baseDir: os.TempDir(),
			mock: func(t *testing.T, s Storage, baseDir, modelKey string, download []byte) {
				file, err := os.OpenFile(filepath.Join(baseDir, "download-bar.csv"), os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0600)
				if err != nil {
					t.Fatal(err)
				}
				defer file.Close()

				if _, err = file.Write(download); err != nil {
					t.Fatal(err)
				}
				s.(*storage).baseDir = "bas"
			},
			expect: func(t *testing.T, s Storage, baseDir, modelKey string, download []byte) {
				assert := assert.New(t)
				_, err := s.ListDownload(modelKey)
				assert.EqualError(err, "open bas/download-bar.csv: no such file or directory")
				s.(*storage).baseDir = baseDir
			},
		},
		{
			name:    "list downloads of a file",
			baseDir: os.TempDir(),
			mock: func(t *testing.T, s Storage, baseDir, modelKey string, download []byte) {
				file, err := os.OpenFile(filepath.Join(baseDir, "download-bar.csv"), os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0600)
				if err != nil {
					t.Fatal(err)
				}
				defer file.Close()

				if _, err = file.Write(download); err != nil {
					t.Fatal(err)
				}
			},
			expect: func(t *testing.T, s Storage, baseDir, modelKey string, download []byte) {
				assert := assert.New(t)
				list, err := s.ListDownload(modelKey)
				assert.NoError(err)
				assert.Equal(len(list), 1)
			},
		},
	}

	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			s := New(tc.baseDir)
			tc.mock(t, s, tc.baseDir, mockModelKey, testData)
			tc.expect(t, s, tc.baseDir, mockModelKey, testData)
			if err := s.ClearDownload(mockModelKey); err != nil {
				t.Fatal(err)
			}
		})
	}
}

func TestStorage_ListNetworkTopology(t *testing.T) {
	require := require.New(t)
	testData, err := os.ReadFile("./testdata/networktopology.csv")
	require.Nil(err, "load test file")

	tests := []struct {
		name    string
		baseDir string
		mock    func(t *testing.T, s Storage, baseDir, modelKey string, networkTopology []byte)
		expect  func(t *testing.T, s Storage, baseDir, modelKey string, networkTopology []byte)
	}{
		{
			name:    "empty csv file given",
			baseDir: os.TempDir(),
			mock: func(t *testing.T, s Storage, baseDir, modelKey string, networkTopology []byte) {
				file, err := os.OpenFile(filepath.Join(baseDir, "networktopology-bar.csv"), os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0600)
				if err != nil {
					t.Fatal(err)
				}
				defer file.Close()
			},
			expect: func(t *testing.T, s Storage, baseDir, modelKey string, networkTopology []byte) {
				assert := assert.New(t)
				_, err := s.ListNetworkTopology(modelKey)
				assert.EqualError(err, "empty csv file given")
			},
		},
		{
			name:    "get file failed",
			baseDir: os.TempDir(),
			mock: func(t *testing.T, s Storage, baseDir, modelKey string, networkTopology []byte) {
				file, err := os.OpenFile(filepath.Join(baseDir, "networktopology-bar.csv"), os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0600)
				if err != nil {
					t.Fatal(err)
				}
				defer file.Close()

				if _, err = file.Write(networkTopology); err != nil {
					t.Fatal(err)
				}
				s.(*storage).baseDir = "foo"
			},
			expect: func(t *testing.T, s Storage, baseDir, modelKey string, networkTopology []byte) {
				assert := assert.New(t)
				_, err := s.ListNetworkTopology(modelKey)
				assert.EqualError(err, "open foo/networktopology-bar.csv: no such file or directory")
				s.(*storage).baseDir = baseDir
			},
		},
		{
			name:    "list network topologies of a file",
			baseDir: os.TempDir(),
			mock: func(t *testing.T, s Storage, baseDir, modelKey string, networkTopology []byte) {
				file, err := os.OpenFile(filepath.Join(baseDir, "networktopology-bar.csv"), os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0600)
				if err != nil {
					t.Fatal(err)
				}
				defer file.Close()

				if _, err = file.Write(networkTopology); err != nil {
					t.Fatal(err)
				}
			},
			expect: func(t *testing.T, s Storage, baseDir, modelKey string, networkTopology []byte) {
				assert := assert.New(t)
				list, err := s.ListNetworkTopology(modelKey)
				assert.NoError(err)
				assert.Equal(len(list), 1)
			},
		},
	}

	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			s := New(tc.baseDir)
			tc.mock(t, s, tc.baseDir, mockModelKey, testData)
			tc.expect(t, s, tc.baseDir, mockModelKey, testData)
			if err := s.ClearNetworkTopology(mockModelKey); err != nil {
				t.Fatal(err)
			}
		})
	}
}

func TestStorage_OpenDownload(t *testing.T) {
	require := require.New(t)
	testData, err := os.ReadFile("./testdata/download.csv")
	require.Nil(err, "load test file")

	tests := []struct {
		name    string
		baseDir string
		mock    func(t *testing.T, s Storage, baseDir, modelKey string, download []byte)
		expect  func(t *testing.T, s Storage, baseDir, modelKey string, download []byte)
	}{
		{
			name:    "open file failed",
			baseDir: os.TempDir(),
			mock: func(t *testing.T, s Storage, baseDir, modelKey string, download []byte) {
				file, err := os.OpenFile(filepath.Join(baseDir, "download-bar.csv"), os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0600)
				if err != nil {
					t.Fatal(err)
				}
				defer file.Close()
				s.(*storage).baseDir = "baw"
			},
			expect: func(t *testing.T, s Storage, baseDir, modelKey string, download []byte) {
				assert := assert.New(t)
				_, err := s.OpenDownload(modelKey)
				assert.EqualError(err, "open baw/download-bar.csv: no such file or directory")
				s.(*storage).baseDir = baseDir
			},
		},
		{
			name:    "open storage with downloads of a file",
			baseDir: os.TempDir(),
			mock: func(t *testing.T, s Storage, baseDir, modelKey string, download []byte) {
				file, err := os.OpenFile(filepath.Join(baseDir, "download-bar.csv"), os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0600)
				if err != nil {
					t.Fatal(err)
				}
				defer file.Close()

				if _, err = file.Write(download); err != nil {
					t.Fatal(err)
				}
			},
			expect: func(t *testing.T, s Storage, baseDir, modelKey string, download []byte) {
				assert := assert.New(t)
				readCloser, err := s.OpenDownload(modelKey)
				assert.NoError(err)

				var downloads []schedulerstorage.Download
				assert.NoError(gocsv.UnmarshalWithoutHeaders(readCloser, &downloads))
				assert.Equal(len(downloads), 1)
			},
		},
	}

	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			s := New(tc.baseDir)
			tc.mock(t, s, tc.baseDir, mockModelKey, testData)
			tc.expect(t, s, tc.baseDir, mockModelKey, testData)
			if err := s.ClearDownload(mockModelKey); err != nil {
				t.Fatal(err)
			}
		})
	}
}

func TestStorage_OpenNetworkTopology(t *testing.T) {
	require := require.New(t)
	testData, err := os.ReadFile("./testdata/networktopology.csv")
	require.Nil(err, "load test file")

	tests := []struct {
		name    string
		baseDir string
		mock    func(t *testing.T, s Storage, baseDir, modelKey string, networkTopology []byte)
		expect  func(t *testing.T, s Storage, baseDir, modelKey string, networkTopology []byte)
	}{
		{
			name:    "open file failed",
			baseDir: os.TempDir(),
			mock: func(t *testing.T, s Storage, baseDir, modelKey string, networkTopology []byte) {
				file, err := os.OpenFile(filepath.Join(baseDir, "networktopology-bar.csv"), os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0600)
				if err != nil {
					t.Fatal(err)
				}
				defer file.Close()
				s.(*storage).baseDir = "bas"
			},
			expect: func(t *testing.T, s Storage, baseDir, modelKey string, networkTopology []byte) {
				assert := assert.New(t)
				_, err := s.OpenNetworkTopology(modelKey)
				assert.EqualError(err, "open bas/networktopology-bar.csv: no such file or directory")
				s.(*storage).baseDir = baseDir
			},
		},
		{
			name:    "open storage with network topologies of a file",
			baseDir: os.TempDir(),
			mock: func(t *testing.T, s Storage, baseDir, modelKey string, networkTopology []byte) {
				file, err := os.OpenFile(filepath.Join(baseDir, "networktopology-bar.csv"), os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0600)
				if err != nil {
					t.Fatal(err)
				}
				defer file.Close()

				if _, err = file.Write(networkTopology); err != nil {
					t.Fatal(err)
				}
			},
			expect: func(t *testing.T, s Storage, baseDir, modelKey string, networkTopology []byte) {
				assert := assert.New(t)
				readCloser, err := s.OpenNetworkTopology(modelKey)
				assert.NoError(err)

				var networkTopologies []schedulerstorage.NetworkTopology
				assert.NoError(gocsv.UnmarshalWithoutHeaders(readCloser, &networkTopologies))
				assert.Equal(len(networkTopologies), 1)
			},
		},
	}

	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			s := New(tc.baseDir)
			tc.mock(t, s, tc.baseDir, mockModelKey, testData)
			tc.expect(t, s, tc.baseDir, mockModelKey, testData)
			if err := s.ClearNetworkTopology(mockModelKey); err != nil {
				t.Fatal(err)
			}
		})
	}
}

func TestStorage_ClearDownload(t *testing.T) {
	tests := []struct {
		name    string
		baseDir string
		mock    func(t *testing.T, s Storage, baseDir, modelKey string)
		expect  func(t *testing.T, s Storage, baseDir, modelKey string)
	}{
		{
			name:    "clear file",
			baseDir: os.TempDir(),
			mock: func(t *testing.T, s Storage, baseDir, modelKey string) {
				file, err := os.OpenFile(filepath.Join(baseDir, "download-bar.csv"), os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0600)
				if err != nil {
					t.Fatal(err)
				}
				defer file.Close()
			},
			expect: func(t *testing.T, s Storage, baseDir, modelKey string) {
				assert := assert.New(t)
				assert.NoError(s.ClearDownload(modelKey))
				fileInfos, err := ioutil.ReadDir(filepath.Join(baseDir))
				assert.NoError(err)

				var backups []fs.FileInfo
				re := regexp.MustCompile(fmt.Sprintf("%s-%s", DownloadFilePrefix, modelKey))

				for _, fileInfo := range fileInfos {
					if !fileInfo.IsDir() && re.MatchString(fileInfo.Name()) {
						backups = append(backups, fileInfo)
					}
				}
				assert.Equal(len(backups), 0)
			},
		},
		{
			name:    "open file failed",
			baseDir: os.TempDir(),
			mock: func(t *testing.T, s Storage, baseDir, modelKey string) {
				file, err := os.OpenFile(filepath.Join(baseDir, "download-bar.csv"), os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0600)
				if err != nil {
					t.Fatal(err)
				}
				defer file.Close()
				s.(*storage).baseDir = "baz"
			},
			expect: func(t *testing.T, s Storage, baseDir, modelKey string) {
				assert := assert.New(t)
				assert.EqualError(s.ClearDownload(modelKey), "remove baz/download-bar.csv: no such file or directory")

				s.(*storage).baseDir = baseDir
				assert.NoError(s.ClearDownload(modelKey))
			},
		},
	}

	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			s := New(tc.baseDir)
			tc.mock(t, s, tc.baseDir, mockModelKey)
			tc.expect(t, s, tc.baseDir, mockModelKey)
		})
	}
}

func TestStorage_ClearNetworkTopology(t *testing.T) {
	tests := []struct {
		name    string
		baseDir string
		mock    func(t *testing.T, s Storage, baseDir, modelKey string)
		expect  func(t *testing.T, s Storage, baseDir, modelKey string)
	}{
		{
			name:    "clear file",
			baseDir: os.TempDir(),
			mock: func(t *testing.T, s Storage, baseDir, modelKey string) {
				file, err := os.OpenFile(filepath.Join(baseDir, "networktopology-bar.csv"), os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0600)
				if err != nil {
					t.Fatal(err)
				}
				defer file.Close()
			},
			expect: func(t *testing.T, s Storage, baseDir, modelKey string) {
				assert := assert.New(t)
				assert.NoError(s.ClearNetworkTopology(modelKey))
				fileInfos, err := ioutil.ReadDir(filepath.Join(baseDir))
				assert.NoError(err)

				var backups []fs.FileInfo
				re := regexp.MustCompile(fmt.Sprintf("%s-%s", NetworkTopologyFilePrefix, modelKey))
				for _, fileInfo := range fileInfos {
					if !fileInfo.IsDir() && re.MatchString(fileInfo.Name()) {
						backups = append(backups, fileInfo)
					}
				}
				assert.Equal(len(backups), 0)
			},
		},
		{
			name:    "open file failed",
			baseDir: os.TempDir(),
			mock: func(t *testing.T, s Storage, baseDir, modelKey string) {
				file, err := os.OpenFile(filepath.Join(baseDir, "networktopology-bar.csv"), os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0600)
				if err != nil {
					t.Fatal(err)
				}
				defer file.Close()
				s.(*storage).baseDir = "baz"
			},
			expect: func(t *testing.T, s Storage, baseDir, modelKey string) {
				assert := assert.New(t)
				assert.EqualError(s.ClearNetworkTopology(modelKey), "remove baz/networktopology-bar.csv: no such file or directory")
				s.(*storage).baseDir = baseDir
				assert.NoError(s.ClearNetworkTopology(modelKey))
			},
		},
	}

	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			s := New(tc.baseDir)
			tc.mock(t, s, tc.baseDir, mockModelKey)
			tc.expect(t, s, tc.baseDir, mockModelKey)
		})
	}
}

func TestStorage_Clear(t *testing.T) {
	tests := []struct {
		name    string
		baseDir string
		mock    func(t *testing.T, s Storage, baseDir string)
		expect  func(t *testing.T, s Storage, baseDir string)
	}{
		{
			name:    "clear file",
			baseDir: os.TempDir(),
			mock: func(t *testing.T, s Storage, baseDir string) {
				s.(*storage).baseDir = filepath.Join(baseDir, "bae")
				if err := os.MkdirAll(s.(*storage).baseDir, os.ModePerm); err != nil {
					t.Fatal(err)
				}

				downloadFile, err := os.OpenFile(filepath.Join(s.(*storage).baseDir, "download-bar.csv"), os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0600)
				if err != nil {
					t.Fatal(err)
				}
				defer downloadFile.Close()

				networkTopologyFile, err := os.OpenFile(filepath.Join(s.(*storage).baseDir, "networktopology-bar.csv"), os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0600)
				if err != nil {
					t.Fatal(err)
				}
				defer networkTopologyFile.Close()
			},
			expect: func(t *testing.T, s Storage, baseDir string) {
				assert := assert.New(t)
				assert.NoError(s.Clear())
				_, err := os.Stat(filepath.Join(baseDir, "bae"))
				assert.EqualError(err, fmt.Sprintf("stat %s: no such file or directory", filepath.Join(baseDir, "bae")))
			},
		},
	}

	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			s := New(tc.baseDir)
			tc.mock(t, s, tc.baseDir)
			tc.expect(t, s, tc.baseDir)
		})
	}
}

func TestStorage_downloadFilename(t *testing.T) {
	baseDir := os.TempDir()
	s := New(baseDir)

	filename := s.(*storage).downloadFilename(mockModelKey)
	re := regexp.MustCompile(fmt.Sprintf("%s-%s.%s$", DownloadFilePrefix, mockModelKey, CSVFileExt))
	assert := assert.New(t)
	assert.True(re.MatchString(filename))
}

func TestStorage_networkTopologyFilename(t *testing.T) {
	baseDir := os.TempDir()
	s := New(baseDir)

	filename := s.(*storage).networkTopologyFilename(mockModelKey)
	re := regexp.MustCompile(fmt.Sprintf("%s-%s.%s$", NetworkTopologyFilePrefix, mockModelKey, CSVFileExt))
	assert := assert.New(t)
	assert.True(re.MatchString(filename))
}