package schedule

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"strconv"
	"strings"
	"testing"

	"github.com/pkg/errors"
	"github.com/stretchr/testify/require"

	. "github.com/pingcap/test-infra/common/model"
	"github.com/pingcap/test-infra/common/model/resource"

	wf "github.com/argoproj/argo-workflows/v3/pkg/apis/workflow/v1alpha1"
	"github.com/argoproj/argo-workflows/v3/workflow/validate"
)

func newTestContext() CompileContext {
	return NewMockContext(
		[]TestCase{
			{
				ID: 1,
				TestCaseMeta: TestCaseMeta{
					Category: CaseCategoryFunctionality,
					Tags:     []string{"stable"},
				},
				TestCaseBase: TestCaseBase{
					Name: "bank",
					Framework: FrameworkUnion{
						Generic: &GenericFrameworkSpec{
							Command:         []string{"bank", "-txn-mode={{parameters.txn-mode}}"},
							Env:             map[string]string{},
							Image:           "hub.pingcap.net/qa/bank:v1.0.0",
							ImagePullPolicy: "IfNotPresent",
						},
					},
					Parameters: []CaseParamDecl{
						{Type: CaseParamTypeResource, Name: "tidb"},
						{Type: CaseParamTypeString, Name: "txn-mode", Default: "pessimistic"},
					},
				},
			},
			{
				ID: 2,
				TestCaseMeta: TestCaseMeta{
					Category: CaseCategoryFunctionality,
					Tags:     []string{"stable"},
				},
				TestCaseBase: TestCaseBase{
					Name: "jepsen",
					Framework: FrameworkUnion{
						Generic: &GenericFrameworkSpec{
							Command:         []string{"run-jepsen.sh", "--workload={{parameters.workload}}"},
							Env:             map[string]string{"nodes": "n1,n2,n3,n4,n5"},
							Image:           "hub.pingcap.net/qa/jepsen:v1.0.0",
							ImagePullPolicy: "IfNotPresent",
						},
					},
					Parameters: []CaseParamDecl{
						{Type: CaseParamTypeString, Name: "workload", Default: "append"},
					},
				},
			},
			{
				ID: 3,
				TestCaseMeta: TestCaseMeta{
					Category: CaseCategoryFunctionality,
					Tags:     []string{"unstable"},
				},
				TestCaseBase: TestCaseBase{
					Name: "pddl",
					Framework: FrameworkUnion{
						Generic: &GenericFrameworkSpec{
							Command:         []string{"schrddl", "-mode=parallel"},
							Image:           "hub.pingcap.net/qa/schrddl:v0.1.0-debug",
							ImagePullPolicy: "IfNotPresent",
							DefaultTimeout:  7200,
						},
					},
					Parameters: []CaseParamDecl{
						{Type: CaseParamTypeResource, Name: "target"},
					},
				},
			},
			{
				ID: 4,
				TestCaseMeta: TestCaseMeta{
					Category: CaseCategoryFunctionality,
					Tags:     []string{"unstable"},
				},
				TestCaseBase: TestCaseBase{
					Name: "sddl",
					Framework: FrameworkUnion{
						Generic: &GenericFrameworkSpec{
							Command:         []string{"schrddl", "-mode=serial"},
							Image:           "hub.pingcap.net/qa/schrddl:v0.1.0-debug",
							ImagePullPolicy: "IfNotPresent",
							DefaultTimeout:  7200,
						},
					},
					Parameters: []CaseParamDecl{
						{Type: CaseParamTypeResource, Name: "target"},
					},
				},
			},
			{
				ID: 5,
				TestCaseMeta: TestCaseMeta{
					Category: CaseCategoryFunctionality,
					Tags:     []string{"stable"},
				},
				DesignedBy: DesignedBy{
					ID:    5,
					Email: "foo@pingcap.com",
				},
				TestCaseBase: TestCaseBase{
					Name: "bank",
					Framework: FrameworkUnion{
						Generic: &GenericFrameworkSpec{
							Command:         []string{"bank", "--dsn={{parameters.dsn}} --spec={{parameters.spec}}"},
							Env:             map[string]string{},
							Image:           "hub.pingcap.net/qa/bank:v1.0.0",
							ImagePullPolicy: "IfNotPresent",
						},
					},
					Parameters: []CaseParamDecl{
						{Type: CaseParamTypeResource, Name: "tidb"},
						{Type: CaseParamTypeString, Name: "dsn"},
						{Type: CaseParamTypeString, Name: "spec"},
					},
				},
			},
		},
		[]TestPlan{
			{
				ID: 1,
				TestPlanBase: TestPlanBase{
					Name: "Debug DDL",
					Variables: []PlanVarDecl{
						{Name: "version", Default: "master"},
					},
					TestBed: TestBed{
						Name: "testbed-",
						Items: []TestBedItem{
							{
								ResourceHandle: resource.ResourceHandle{Name: "tidb", Type: resource.ResourceTypeTiDBCluster},
								Spec:           json.RawMessage(`{"version": "{{variables.version}}"}`),
							},
						},
						ResourcePool: "default",
						RmsUrl:       "http://rms-apiserver-service.test-infra",
					},
					Pipeline: []PipelineStep{
						{
							Name:    "main",
							Timeout: 3600,
							CaseRef: CaseRef{
								CaseID: 3,
								Arguments: []CaseParamBind{
									{Name: "target", Value: "{{testbed.tidb}}"},
								},
							},
						},
					},
				},
			},
			{
				ID: 2,
				TestPlanBase: TestPlanBase{
					Name:     "Transaction Suite",
					FailFast: true,
					Variables: []PlanVarDecl{
						{Name: "version", Default: "master"},
						{Name: "txn-mode", Default: "pessimistic"},
					},
					TestBed: TestBed{
						Name: "testbed-",
						Items: []TestBedItem{
							{
								ResourceHandle: resource.ResourceHandle{Name: "tidb", Type: resource.ResourceTypeTiDBCluster},
								Spec:           json.RawMessage(`{"version": "{{variables.version}}"}`),
							},
						},
						ResourcePool: "default",
						RmsUrl:       "http://rms-apiserver-service.test-infra",
					},
					Pipeline: []PipelineStep{
						{
							Name: "bank",
							CaseRef: CaseRef{
								CaseID: 1,
								Arguments: []CaseParamBind{
									{Name: "tidb", Value: "{{testbed.tidb}}"},
									{Name: "txn-mode", Value: "{{variables.txn-mode}}"},
								},
							},
						},
						{
							Name:  "append",
							After: []string{"bank"},
							CaseRef: CaseRef{
								CaseID: 2,
								Arguments: []CaseParamBind{
									{Name: "workload", Value: "append"},
								},
							},
						},
						{
							Name:  "register",
							After: []string{"bank"},
							CaseRef: CaseRef{
								CaseID: 2,
								Arguments: []CaseParamBind{
									{Name: "workload", Value: "register"},
								},
							},
						},
					},
				},
			},
			{
				ID: 3,
				TestPlanBase: TestPlanBase{
					Name: "Release Test",
					Variables: []PlanVarDecl{
						{Name: "version"},
					},
					TestBed: TestBed{
						Name: "testbed-",
						Items: []TestBedItem{
							{
								ResourceHandle: resource.ResourceHandle{Name: "tidb", Type: resource.ResourceTypeTiDBCluster},
								Spec:           json.RawMessage(`{"version": "{{variables.version}}"}`),
							},
						},
						ResourcePool: "default",
						RmsUrl:       "http://rms-apiserver-service.test-infra",
					},
					Pipeline: []PipelineStep{
						{
							Name:    "pddl",
							Timeout: 3600,
							CaseRef: CaseRef{
								CaseID: 3,
								Arguments: []CaseParamBind{
									{Name: "target", Value: "{{testbed.tidb}}"},
								},
							},
						},
						{
							Name:    "sddl",
							Timeout: 3600,
							CaseRef: CaseRef{
								CaseID: 3,
								Arguments: []CaseParamBind{
									{Name: "target", Value: "{{testbed.tidb}}"},
								},
							},
						},
					},
				},
			},
			{
				ID: 4,
				TestPlanBase: TestPlanBase{
					Name:     "param-inject",
					FailFast: true,
					Variables: []PlanVarDecl{
						{Name: "version", Default: "master"},
						{Name: "txn-mode", Default: "pessimistic"},
					},
					TestBed: TestBed{
						Name: "testbed-",
						Items: []TestBedItem{
							{
								ResourceHandle: resource.ResourceHandle{Name: "tidb", Type: resource.ResourceTypeTiDBCluster},
								Spec:           json.RawMessage(`{"version": "{{variables.version}}"}`),
							},
						},
						ResourcePool: "default",
						RmsUrl:       "http://rms-apiserver-service.test-infra",
					},
					Pipeline: []PipelineStep{
						{
							Name: "bank",
							CaseRef: CaseRef{
								CaseID: 5,
								Arguments: []CaseParamBind{
									{Name: "tidb", Value: "{{testbed.tidb}}"},
									{Name: "dsn", Value: "{{testbed.items.tidb.tidb.dsn}}"},
									{Name: "spec", Value: "{{testbed.items.tidb.tidb.spec}}-{{testbed.items.tidb.tidb.dsn}}"},
								},
							},
						},
						{
							Name: "bank2",
							CaseRef: CaseRef{
								CaseID: 5,
								Arguments: []CaseParamBind{
									{Name: "tidb", Value: "{{testbed.tidb}}"},
									{Name: "dsn", Value: "{{testbed.items.tidb.tidb.dsn}}"},
									{Name: "spec", Value: "{{testbed.items.tidb.tidb.spec}}"},
								},
							},
						},
					},
				},
			},
			{
				ID: 5,
				TestPlanBase: TestPlanBase{
					Name:     "steps",
					FailFast: false,
					Variables: []PlanVarDecl{
						{Name: "version", Default: "master"},
						{Name: "txn-mode", Default: "pessimistic"},
					},
					TestBed: TestBed{
						Name: "testbed-",
						Items: []TestBedItem{
							{
								ResourceHandle: resource.ResourceHandle{Name: "tidb", Type: resource.ResourceTypeTiDBCluster},
								Spec:           json.RawMessage(`{"version": "{{variables.version}}"}`),
							},
						},
						ResourcePool: "default",
						RmsUrl:       "http://rms-apiserver-service.test-infra",
					},
					Steps: []PipelineStep{
						{
							Name: "bank",
							CaseRef: CaseRef{
								CaseID: 5,
								Arguments: []CaseParamBind{
									{Name: "tidb", Value: "{{testbed.tidb}}"},
									{Name: "dsn", Value: "{{testbed.items.tidb.tidb.dsn}}"},
									{Name: "spec", Value: "{{testbed.items.tidb.tidb.spec}}-{{testbed.items.tidb.tidb.dsn}}"},
								},
							},
						},
						{
							Name: "bank2",
							CaseRef: CaseRef{
								CaseID: 5,
								Arguments: []CaseParamBind{
									{Name: "tidb", Value: "{{testbed.tidb}}"},
									{Name: "dsn", Value: "{{testbed.items.tidb.tidb.dsn}}"},
									{Name: "spec", Value: "{{testbed.items.tidb.tidb.spec}}"},
								},
							},
						},
					},
				},
			},
		},
		0,
	)
}

func mustDumpJson(t *testing.T, wfSpec *wf.WorkflowSpec, out string) {
	f, err := os.OpenFile(out, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0644)
	require.NoError(t, err)
	defer f.Close()
	enc := json.NewEncoder(f)
	enc.SetEscapeHTML(false)
	enc.SetIndent("", "  ")
	require.NoError(t, enc.Encode(wfSpec))
}

func TestCompileArgo(t *testing.T) {
	for _, tt := range []struct {
		name string
		plan PlanRef
	}{
		{"debug-ddl", PlanRef{PlanID: 1, VariableBinds: []PlanVarBind{
			{Name: "version", Value: "v5.0.0"},
		}}},
		{"txn-suite", PlanRef{PlanID: 2, VariableBinds: []PlanVarBind{
			{Name: "txn-mode", Value: "optimistic"},
		}}},
		{"release-test", PlanRef{PlanID: 3, VariableBinds: []PlanVarBind{
			{Name: "version", Value: "v5.0.1-pre"},
		}}},
		{"param-inject", PlanRef{PlanID: 4, VariableBinds: []PlanVarBind{}}},
		{"steps", PlanRef{PlanID: 5, VariableBinds: []PlanVarBind{}}},
	} {
		t.Run(tt.name, func(t *testing.T) {
			act, err := CompileArgo(newTestContext(), tt.plan)
			require.NoError(t, err)

			arg := wf.Parameter{Name: setupImageParamName, Value: wf.AnyStringPtr("setup:latest")}
			act.Arguments.Parameters = append(act.Arguments.Parameters, arg)

			conds, err := validate.ValidateWorkflow(nil, nil, &wf.Workflow{Spec: *act}, validate.ValidateOpts{
				Lint:                       true,
				WorkflowTemplateValidation: true,
			})
			require.NoError(t, err)
			require.Empty(t, conds)

			if os.Getenv("DUMP_FIXTURES") == "1" {
				mustDumpJson(t, act, fmt.Sprintf("fixtures/%s.json", tt.name))
			}

			expStr, err := ioutil.ReadFile(fmt.Sprintf("fixtures/%s.json", tt.name))
			require.NoError(t, err)
			exp := &wf.WorkflowSpec{}
			err = json.Unmarshal(expStr, exp)
			require.NoError(t, err)

			compareAndClearSpecialFields(t, exp, act)
			require.EqualValues(t, exp, act)
		})
	}
}

func getSetupTemplate(w *wf.WorkflowSpec) *wf.Template {
	for i := 0; i < len(w.Templates); i++ {
		if strings.HasSuffix(w.Templates[i].Name, "-setup") {
			return &w.Templates[i]
		}
	}
	return nil
}

func compareAndClearSpecialFields(t *testing.T, exp, act *wf.WorkflowSpec) {

	expSetup := getSetupTemplate(exp).Container
	actSetup := getSetupTemplate(act).Container

	expSetupArgs := strings.Split(expSetup.Args[0], "\n")
	actSetupArgs := strings.Split(actSetup.Args[0], "\n")
	require.ElementsMatch(t, expSetupArgs, actSetupArgs)
	expSetup.Args[0] = ""
	actSetup.Args[0] = ""

	for i := 0; i < len(expSetup.Env); i++ {
		require.Equal(t, expSetup.Env[i].Name, actSetup.Env[i].Name)
		if strings.HasPrefix(expSetup.Env[i].Name, "TESTBED_") {
			require.JSONEq(t, expSetup.Env[i].Value, actSetup.Env[i].Value)
			expSetup.Env[i].Value = ""
			actSetup.Env[i].Value = ""
		}
	}

	for i := 0; i < len(exp.Templates); i++ {
		expi := &exp.Templates[i]
		acti := &act.Templates[i]
		require.ElementsMatch(t, expi.Inputs.Parameters, acti.Inputs.Parameters)
		expi.Inputs.Parameters = nil
		acti.Inputs.Parameters = nil
		require.ElementsMatch(t, expi.Outputs.Parameters, acti.Outputs.Parameters)
		expi.Outputs.Parameters = nil
		acti.Outputs.Parameters = nil
		if expi.DAG != nil {
			for j := 0; j < len(expi.DAG.Tasks); j++ {
				expArg := &expi.DAG.Tasks[j].Arguments
				actArg := &acti.DAG.Tasks[j].Arguments
				require.ElementsMatch(t, expArg.Parameters, actArg.Parameters)
				expArg.Parameters = nil
				actArg.Parameters = nil
			}
		}
		if expi.Steps != nil {
			for j := 0; j < len(expi.Steps); j++ {
				expArg := &expi.Steps[j].Steps[0].Arguments
				actArg := &acti.Steps[j].Steps[0].Arguments
				require.ElementsMatch(t, expArg.Parameters, actArg.Parameters)
				expArg.Parameters = nil
				actArg.Parameters = nil
			}
		}
	}
}

func TestRecursiveSubstitute(t *testing.T) {
	for i, tt := range []struct {
		input      interface{}
		output     interface{}
		substitute func(string) (string, error)
		ok         bool
	}{
		{ok: true},
		{input: "foo", output: "foo", ok: true},
		{
			input: "foo", output: "bar", ok: true,
			substitute: func(s string) (string, error) {
				return "bar", nil
			},
		},
		{
			input:  map[string]interface{}{"foo": "bar", "bar": "foo"},
			output: map[string]interface{}{"foo": "bar", "bar": "bar"},
			ok:     true,
			substitute: func(s string) (string, error) {
				return "bar", nil
			},
		},
		{
			input:  []interface{}{"foo", "bar"},
			output: []interface{}{"bar", "bar"},
			ok:     true,
			substitute: func(s string) (string, error) {
				return "bar", nil
			},
		},
		{
			input:  []interface{}{"foo", map[string]interface{}{"foo": "bar", "bar": "foo"}},
			output: []interface{}{"bar", map[string]interface{}{"foo": "bar", "bar": "bar"}},
			ok:     true,
			substitute: func(s string) (string, error) {
				return "bar", nil
			},
		},
		{
			input:  []interface{}{"foo", map[string]interface{}{"foo": "bar", "bar": "foo"}},
			output: []interface{}{"bar", map[string]interface{}{"foo": "bar", "bar": "bar"}},
			substitute: func(s string) (string, error) {
				if s == "foo" {
					return "", errors.New("oops")
				}
				return "bar", nil
			},
		},
	} {
		t.Run(strconv.Itoa(i), func(t *testing.T) {
			out, err := recursiveSubstitute(tt.input, tt.substitute)
			if tt.ok {
				require.NoError(t, err)
				require.Equal(t, tt.output, out)
			} else {
				require.Error(t, err)
			}
		})
	}
}

func TestIsCaseVersionSupported(t *testing.T) {
	for _, testItem := range []struct {
		name             string
		branch           string
		supportedVersion []string
		expected         bool
		expectedError    error
	}{
		{
			name:             "release",
			branch:           "release-5.4",
			supportedVersion: []string{">=5.4, <5.5"},
			expected:         true,
			expectedError:    nil,
		},
		{
			name:             "release",
			branch:           "v5.1.3-hotfix-20220126",
			supportedVersion: []string{">v5.0.0"},
			expected:         true,
			expectedError:    nil,
		},
		{
			name:             "and",
			branch:           "v5.6",
			supportedVersion: []string{">5.4, <5.5"},
			expected:         false,
			expectedError:    nil,
		},
		{
			name:             "or",
			branch:           "v5.6.1",
			supportedVersion: []string{">5.4, <5.5", ">=5.6"},
			expected:         true,
			expectedError:    nil,
		},
		{
			name:             "v5.4-ge",
			branch:           "v5.4.1",
			supportedVersion: []string{">5.4.1"},
			expected:         false,
			expectedError:    nil,
		},
		{
			name:             "v5.4",
			branch:           "v5.4",
			supportedVersion: []string{"=5.4"},
			expected:         true,
			expectedError:    nil,
		},
		{
			name:             "hotfix",
			branch:           "v5.1.0-hotfix",
			supportedVersion: []string{">=5.1.0"},
			expected:         true,
			expectedError:    nil,
		},
		{
			name:             "verByName@v5.0.0",
			branch:           "",
			supportedVersion: []string{">=5.1.0"},
			expected:         false,
			expectedError:    nil,
		},
		{
			name:             "verByName@5.0.1",
			branch:           "",
			supportedVersion: []string{">=5.0.0"},
			expected:         true,
			expectedError:    nil,
		},
		{
			name:             "badVersion@va.b.c",
			branch:           "",
			supportedVersion: []string{">=5.0.0"},
			expected:         false,
			expectedError:    fmt.Errorf("Malformed version: va.b.c"),
		},
	} {
		t.Run(testItem.name, func(t *testing.T) {
			ctx := NewMockContext([]TestCase{
				{
					ID: 1,
					TestCaseMeta: TestCaseMeta{
						SupportedVersions: testItem.supportedVersion,
					},
				}}, nil, 0)
			supported, err := isCaseVersionSupported(ctx, &CaseRef{CaseID: 1},
				&TestPlan{
					TestPlanBase:   TestPlanBase{Name: testItem.name},
					GitHubPlanInfo: GitHubPlanInfo{Branch: testItem.branch},
				})
			require.Equal(t, testItem.expectedError, err)
			require.Equal(t, testItem.expected, supported)
		})
	}
}