package tidbcluster

import (
	"bufio"
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"path/filepath"
	"strings"
	"sync"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"

	apiresouce "github.com/pingcap/test-infra/common/model/resource"
	"github.com/pingcap/test-infra/sdk/e2etest/framework"
	"github.com/pingcap/test-infra/sdk/resource"
)

var _ = Describe("WorkloadNode Test", func() {
	Context("WorkloadNode.Create", func() {
		f := framework.NewDefaultFramework("testbed-workloadnode-create")

		It("WorkloadNode's testbed should equal to testbed's name", func() {
			node := f.CreateWorkloadNode(f.CreateTestbed())
			Expect(node.Testbed).Should(Equal(f.Testbed()))
		})

		It("WorkloadNode should be ready after a while", func() {
			node := f.CreateWorkloadNode(f.CreateTestbed())
			err := f.WaitForWorkloadNodeReady(node, framework.TestContext.ResourceWaitReadyTimeout)
			framework.ExpectNoError(err)
		})

		It("WorkloadNode pod should be allocated in specific k8s machine node", func() {
			k8sNode := f.AllNodes().Items[0]
			node := f.CreateWorkloadNodeWithNodeSelector(f.CreateTestbed(), k8sNode.Name)
			err := f.WaitForWorkloadNodeReady(node, framework.TestContext.ResourceWaitReadyTimeout)
			framework.ExpectNoError(err)

			pods := f.AllPods()
			Expect(pods.Items[0].Spec.NodeName).Should(Equal(k8sNode.Name))
		})

		It("WorkloadNode should be created with a volume mount", func() {
			mountPath := "/path-not-exist"
			node := f.CreateWorkloadNodeWithVolumeMount(f.CreateTestbed(), mountPath)
			err := f.WaitForWorkloadNodeReady(node, framework.TestContext.ResourceWaitReadyTimeout)
			framework.ExpectNoError(err)

			_, _, exitCode, err := node.Exec(resource.WorkloadNodeExecOptions{Command: fmt.Sprintf("ls %s", mountPath)})
			framework.ExpectNoError(err)
			Expect(exitCode).Should(Equal(0))
		})

		It("WorkloadNode should be created with the specified storageClass", func() {
			storageClass := f.AllStorageClassPreferNotDefault()[0].Name
			mountPath := "/path-not-exist"
			node := f.CreateWorkloadNodeWithStorageClass(f.CreateTestbed(), storageClass, mountPath)
			err := f.WaitForWorkloadNodeReady(node, framework.TestContext.ResourceWaitReadyTimeout)
			framework.ExpectNoError(err)

			_, _, exitCode, err := node.Exec(resource.WorkloadNodeExecOptions{Command: fmt.Sprintf("ls %s", mountPath)})
			framework.ExpectNoError(err)
			Expect(exitCode).Should(Equal(0))
		})

		It("WorkloadNode should be created with additionalVolume", func() {
			testbed := f.CreateTestbed()
			volumeMountPath := "/data/hostpath/"
			node := f.CreateWorkloadNodeWithAdditionalVolume(testbed, volumeMountPath)
			framework.ExpectNoError(f.WaitForWorkloadNodeReady(node, framework.TestContext.ResourceWaitReadyTimeout))

			_, _, exitCode, err := node.Exec(resource.WorkloadNodeExecOptions{Command: fmt.Sprintf("ls %s", volumeMountPath)})
			framework.ExpectNoError(err)
			Expect(exitCode).Should(Equal(0))
		})

		It("WorkloadNode TiUP node should be ready after a while", func() {
			node := f.CreateTiUPWorkloadNode(f.CreateTestbed())
			err := f.WaitForWorkloadNodeReady(node, framework.TestContext.ResourceWaitReadyTimeout)
			framework.ExpectNoError(err)
		})
	})

	Context("WorkloadNode.Get", func() {
		f := framework.NewDefaultFramework("testbed-workloadnode-get")

		It("WorkloadNode's testbed should equal to testbed's name", func() {
			node := f.CreateWorkloadNode(f.CreateTestbed())
			node = f.GetWorkloadNode(node.Handle())

			Expect(node.Testbed).Should(Equal(f.Testbed()))
		})

		It("node should be ready after a while", func() {
			node := f.CreateWorkloadNode(f.CreateTestbed())
			node = f.GetWorkloadNode(node.Handle())

			err := f.WaitForWorkloadNodeReady(node, framework.TestContext.ResourceWaitReadyTimeout)
			framework.ExpectNoError(err)
		})
	})

	Context("WorkloadNode.GetEndpoints", func() {
		f := framework.NewDefaultFramework("testbed-workloadnode-getendpoints")

		It("WorkloadNode.GetEndpoints should return with Others after ready", func() {
			node := f.CreateWorkloadNodeWithNodePort(f.CreateTestbed(), "80")
			node = f.GetWorkloadNode(node.Handle())

			err := f.WaitForWorkloadNodeReady(node, framework.TestContext.ResourceWaitReadyTimeout)
			framework.ExpectNoError(err)

			endpoints := f.GetWorkloadNodeEndpoints(node.Handle())

			Expect(len(endpoints)).Should(Equal(1))
			Expect(endpoints[0].Type).Should(Equal(apiresouce.EndpointTypeOther))
			Expect(endpoints[0].ExternalSvcAddress).ShouldNot(Equal(""))

			_, err = http.Get("http://" + endpoints[0].ExternalSvcAddress)
			framework.ExpectNoError(err)
		})
	})

	Context("WorkloadNode.Archive", func() {
		f := framework.NewDefaultFramework("testbed-workloadnode-archive")

		It("WorkloadNode.Archive should work well", func() {
			node := f.CreateWorkloadNode(f.CreateTestbed())
			node = f.GetWorkloadNode(node.Handle())

			err := f.WaitForWorkloadNodeReady(node, framework.TestContext.ResourceWaitReadyTimeout)
			framework.ExpectNoError(err)

			_, _, _, err = node.Exec(resource.WorkloadNodeExecOptions{Command: "echo test-infra > /tmp/node.log"})
			framework.ExpectNoError(err)
			var data []string

			f.ArchiveWorkloadNode(node.Handle(), func(objectName string, reader io.Reader) (*url.URL, error) {
				body, err := ioutil.ReadAll(reader)
				framework.ExpectNoError(err)
				data = append(data, string(body))
				return nil, nil
			})
			Expect(len(data)).Should(Equal(1))
			Expect(data[0]).Should(ContainSubstring("test-infra"))
		})
	})

	Context("WorkloadNode.Delete", func() {
		f := framework.NewDefaultFramework("testbed-workloadnode-delete")

		It("WorkloadNode's delete should clean all resources", func() {
			node := f.CreateWorkloadNode(f.CreateTestbed())

			f.DeleteWorkloadNodeWaited(node)

			pods := f.AllPods()
			Expect(len(pods.Items)).Should(Equal(0))
			Expect(len(f.AllPVCs().Items)).Should(Equal(0))
			Expect(len(f.AllServices().Items)).Should(Equal(0))
			Expect(len(f.AllConfigmaps().Items)).Should(Equal(1))
		})
	})

	Context("WorkloadNode.Exec", func() {
		f := framework.NewDefaultFramework("testbed-workloadnode-exec")

		It("WorkloadNode exec an normal command should work well", func() {
			node := f.CreateWorkloadNode(f.CreateTestbed())
			f.WaitForWorkloadNodeReady(node, framework.TestContext.ResourceWaitReadyTimeout)

			_, _, exitCode, _ := node.Exec(resource.WorkloadNodeExecOptions{Command: "echo hello > tmpfile"})
			Expect(exitCode).Should(Equal(0))

			stdout, _, exitCode, _ := node.Exec(resource.WorkloadNodeExecOptions{Command: "cat tmpfile"})
			Expect(exitCode).Should(Equal(0))
			Expect(stdout).Should(Equal("hello"))
		})

		It("WorkloadNode exec an error command should return an exit code", func() {
			node := f.CreateWorkloadNode(f.CreateTestbed())
			f.WaitForWorkloadNodeReady(node, framework.TestContext.ResourceWaitReadyTimeout)

			stdout, stderr, exitCode, _ := node.Exec(resource.WorkloadNodeExecOptions{Command: "cat not_exist_file"})
			Expect(exitCode).Should(Equal(1))
			Expect(stdout).Should(Equal(""))
			Expect(stderr).Should(Equal("cat: can't open 'not_exist_file': No such file or directory"))
		})
	})

	Context("WorkloadNode.ExecAsync", func() {
		f := framework.NewDefaultFramework("testbed-workloadnode-exec-async")

		It("WorkloadNode ExecAsyncInterruptible an normal command should work well", func() {
			node := f.CreateWorkloadNode(f.CreateTestbed())
			f.WaitForWorkloadNodeReady(node, framework.TestContext.ResourceWaitReadyTimeout)

			reader, writer := io.Pipe()
			remote, _ := node.ExecAsyncInterruptible(writer, resource.WorkloadNodeExecOptions{Command: "for i in $(seq 1 4); do echo $i; sleep 5; done"})
			b := bufio.NewReader(reader)

			for i := 1; i <= 3; i++ {
				n, err := b.ReadString('\n')
				framework.ExpectNoError(err)
				Expect(n).Should(Equal(fmt.Sprintf("%d\r\n", i)))
			}
			remain, err := io.ReadAll(b)
			framework.ExpectNoError(err)
			Expect(remain).Should(Equal([]byte("4\r\n")))
			commandErr := remote.Wait()
			framework.ExpectNoError(commandErr.Err)
		})

		It("WorkloadNode ExecAsyncInterruptible close can interrupt the command", func() {
			node := f.CreateWorkloadNode(f.CreateTestbed())
			f.WaitForWorkloadNodeReady(node, framework.TestContext.ResourceWaitReadyTimeout)

			b := bytes.Buffer{}
			remote, _ := node.ExecAsyncInterruptible(resource.NopCloser(&b), resource.WorkloadNodeExecOptions{Command: "sleep 10; touch not_exist_file"})
			remote.Interrupt()
			commandErr := remote.Wait()
			Expect(commandErr.ExitCode).Should(Equal(130))

			stdout, stderr, exitCode, _ := node.Exec(resource.WorkloadNodeExecOptions{Command: "cat not_exist_file"})
			Expect(exitCode).Should(Equal(1))
			Expect(stdout).Should(Equal(""))
			Expect(stderr).Should(Equal("cat: can't open 'not_exist_file': No such file or directory"))
		})

		It("WorkloadNode ExecAsyncInterruptible an error command should return an exit code", func() {
			node := f.CreateWorkloadNode(f.CreateTestbed())
			f.WaitForWorkloadNodeReady(node, framework.TestContext.ResourceWaitReadyTimeout)

			reader, writer := io.Pipe()
			remote, _ := node.ExecAsyncInterruptible(writer, resource.WorkloadNodeExecOptions{Command: "cat not_exist_file"})

			b := bufio.NewReader(reader)
			n, _ := b.ReadString('\n')
			Expect(n).Should(Equal("cat: can't open 'not_exist_file': No such file or directory\r\n"))
			commandErr := remote.Wait()
			Expect(commandErr.ExitCode).Should(Equal(1))
		})

		It("WorkloadNode ExecAsync an error command should return an exit code", func() {
			node := f.CreateWorkloadNode(f.CreateTestbed())
			f.WaitForWorkloadNodeReady(node, framework.TestContext.ResourceWaitReadyTimeout)

			stdout := &bytes.Buffer{}
			stderr := &bytes.Buffer{}
			remote, _ := node.ExecAsync(resource.NopCloser(stdout), resource.NopCloser(stderr), resource.WorkloadNodeExecOptions{Command: "echo 1; cat not_exist_file"})
			commandErr := remote.Wait()

			Expect(stdout.String()).Should(Equal("1\n"))
			Expect(stderr.String()).Should(Equal("cat: can't open 'not_exist_file': No such file or directory\n"))

			Expect(commandErr.ExitCode).Should(Equal(1))
		})

		It("WorkloadNode ExecAsync should work well", func() {
			node := f.CreateWorkloadNode(f.CreateTestbed())
			f.WaitForWorkloadNodeReady(node, framework.TestContext.ResourceWaitReadyTimeout)

			reader, writer := io.Pipe()
			errReader, errWriter := io.Pipe()
			remote, _ := node.ExecAsync(writer, errWriter, resource.WorkloadNodeExecOptions{Command: "for i in $(seq 1 3); do echo $i; >&2 echo $i ; sleep 2; done"})
			var stderr []byte
			var stderrErr error
			var wg sync.WaitGroup
			wg.Add(1)
			go func() {
				stderr, stderrErr = io.ReadAll(errReader) // read 2 reader at single thread will block
				wg.Done()
			}()
			stdout, err := io.ReadAll(reader)
			framework.ExpectNoError(err)
			Expect(string(stdout)).Should(Equal("1\n2\n3\n"))
			commandErr := remote.Wait()

			wg.Wait()
			framework.ExpectNoError(stderrErr)
			Expect(string(stderr)).Should(Equal("1\n2\n3\n"))

			Expect(commandErr.ExitCode).Should(Equal(0))
		})
	})

	Context("WorkloadNode.WriteFile", func() {
		f := framework.NewDefaultFramework("testbed-workloadnode-writefile")

		It("WorkloadNode writeFile works well", func() {
			node := f.CreateWorkloadNode(f.CreateTestbed())
			f.WaitForWorkloadNodeReady(node, framework.TestContext.ResourceWaitReadyTimeout)

			err := node.WriteFile(resource.WorkloadNodeWriteFileOptions{
				Content:  strings.NewReader("hello, test-infra"),
				FilePath: "/tmp/welcome",
			})
			stdout, _, _, err := node.Exec(resource.WorkloadNodeExecOptions{Command: "cat /tmp/welcome"})
			Expect(stdout).Should(Equal("hello, test-infra"))

			By("WriteFile a same file again")
			framework.ExpectNoError(err)
			err = node.WriteFile(resource.WorkloadNodeWriteFileOptions{
				Content:  strings.NewReader("hello, test-infra!"),
				FilePath: "/tmp/welcome",
			})
			framework.ExpectNoError(err)
			stdout, _, _, err = node.Exec(resource.WorkloadNodeExecOptions{Command: "cat /tmp/welcome"})
			Expect(stdout).Should(Equal("hello, test-infra!"))

			output, err := io.ReadAll(node.TarStreamFrom(node.WorkloadNode.Name+"-0", "alpineblock", "/tmp/welcome"))
			framework.ExpectNoError(err)
			Expect(string(output)).Should(ContainSubstring("hello, test-infra!"))
		})
	})

	Context("WorkloadNode.Artifact", func() {
		f := framework.NewDefaultFramework("testbed-workloadnode-artifact")

		It("WorkloadNode should archive content successfully", func() {
			artifactPath := "/test-infra"
			node := f.CreateWorkloadNodeWithArtifactPath(f.CreateTestbed(), artifactPath)
			f.WaitForWorkloadNodeReady(node, framework.TestContext.ResourceWaitReadyTimeout)

			filePath := filepath.Join(artifactPath, "welcome")
			content := "hello, test-infra"
			err := node.WriteFile(resource.WorkloadNodeWriteFileOptions{
				Content:  strings.NewReader(content),
				FilePath: filePath,
			})
			framework.ExpectNoError(err)

			var data []string
			f.ArtifactWorkloadNode(node, f.Archiver(&data))
			Expect(len(data)).Should(Equal(1))
			Expect(data[0]).Should(ContainSubstring(content))
		})

		It("WorkloadNode should archive files with wildcards successfully", func() {
			node := f.CreateWorkloadNodeWithArtifactPath(f.CreateTestbed(), "/test-infra/wel*")
			f.WaitForWorkloadNodeReady(node, framework.TestContext.ResourceWaitReadyTimeout)

			content := "hello, test-infra"
			err := node.WriteFile(resource.WorkloadNodeWriteFileOptions{
				Content:  strings.NewReader(content),
				FilePath: "/test-infra/welcome",
			})
			framework.ExpectNoError(err)

			var data []string
			f.ArtifactWorkloadNode(node, f.Archiver(&data))
			Expect(len(data)).Should(Equal(1))
			Expect(data[0]).Should(ContainSubstring(content))
		})
	})
})