package schedule

import (
	"context"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"path"
	"strconv"
	"strings"
	"time"

	. "github.com/pingcap/test-infra/common/model"
	"github.com/pingcap/test-infra/services/tps/pkg/common"
	"github.com/pkg/errors"
	kerrors "k8s.io/apimachinery/pkg/api/errors"
	"k8s.io/klog/v2"

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

type Config struct {
	ServiceAccountName string
	ArtifactRepoRef    *wfapi.ArtifactRepositoryRef
	Endpoints          map[common.EndpointType]string
}

type StatusSyncer struct {
	TCMS common.TCMS
	Argo ArgoClientMgr
}

func (p *StatusSyncer) Process(ctx context.Context, job *common.ArgoJob) error {
	argo, err := p.Argo.ClientFor(job)
	if err != nil {
		return err
	}
	wf, err := argo.Get(ctx, job.Workflow.Name)
	if err != nil {
		if kerrors.IsNotFound(err) {
			klog.Warningf("mark job done since workflow %q is gone", job.Workflow.Name)
			return SetCanceledTestJob(ctx, job, err.Error(), fmt.Sprintf("workflow %q is gone, may be cancelled manually.", job.Workflow.Name), p.TCMS)
		}
		return errors.WithStack(err)
	}
	if wf.Status.Phase == wfapi.WorkflowUnknown || wf.Status.Phase == wfapi.WorkflowPending {
		return nil
	}

	oldStatus := job.PlanExec.Status
	job.UpdateStatus(wf)
	klog.InfoS("syncing job status", "name", job, "status", job.Status)

	planExec := job.PlanExec
	// update plan-execution and testbed info
	// PREPARING ----> RUNNING, ERROR, PREPARE BACKOFF
	if planExec.Status == ExecutionStatusPreparing {
		if planExec.Meta == nil {
			planExec.Meta = map[string]string{}
		}
		planExec.Meta["argo.namespace"] = wf.Namespace
		planExec.Meta["argo.name"] = wf.Name
		planExec.ArgoURL = p.genPlanArgoURL(argo.URL(), wf.Namespace, wf.Name)

		setupStatus, setupFinishedAt, setupLog := getSetupStatus(wf)
		// Update TestBed Status to TCMS
		tb := &TestBedInfo{
			PlanExecutionID: job.PlanExec.ID,
		}
		if setupLog != nil {
			tb.SetUpLogURL = setupLog.String()
		}
		switch setupStatus {
		case SetupSuccess:
			tb.Status = TestBedInfoStatusReady
			tb.InitiatedAt = setupFinishedAt
			job.Message = "running"
			planExec.Status = ExecutionStatusRunning
		case SetupFailRetryable:
			tb.Status = TestBedInfoStatusError
			job.Message = fmt.Sprintf("setUp retryable failed: argo job %s", wf.Name)
			planExec.RetryCnt++
			planExec.Status = ExecutionStatusPrepareBackoff
		case SetupFailNonRetryable:
			tb.Status = TestBedInfoStatusError
			job.Message = "env error"
			planExec.Status = ExecutionStatusError
		default:
			// plan-execution status not changed
			return nil
		}

		// update testbed status to tcms
		if testbedName, ok := wf.Annotations["testbed"]; !ok {
			klog.Errorf("job.Plan.TestBed.Name %s: cannot find field 'testbed' in argo annotations", job.Plan.TestBed.Name)
		} else if testbedName != job.Plan.TestBed.Name {
			klog.Errorf("job.Plan.TestBed.Name %s is not equal to testbedName from argo %s", job.Plan.TestBed.Name, testbedName)
		} else {
			tb.Name = testbedName
			// PutTestBedInfo will sync from RMS when first set to READY
			err = p.TCMS.PutTestBedInfo(tb)
			if err != nil {
				klog.Errorf("PutTestBedInfo() failed: %s", err.Error())
			}
		}
	}

	// update case-executions
	// exclude plan-execution ExecutionStatusPrepareBackoff
	if planExec.Status != ExecutionStatusPrepareBackoff {
		// update (status changed) case-executions
		if err = updateCaseStatus(job, p.TCMS, wf, readCaseStatus); err != nil {
			klog.ErrorS(err, "failed to update case executions status", "workflow", wf.Name)
		}
	}

	if !planExec.Status.IsCompleted() && planExec.Status != ExecutionStatusPrepareBackoff {
		// update plan-execution
		// COMPLETED: FAILURE, ERROR, CANCELLED
		if job.Status == common.ArgoTestJobDone {
			switch wf.Status.Phase {
			case wfapi.WorkflowSucceeded:
				job.Message = "passed"
				planExec.Status = ExecutionStatusSuccess
			case wfapi.WorkflowFailed:
				if existEnvPodsFailed(wf) {
					job.Message = "env error"
					planExec.Status = ExecutionStatusError
				} else {
					job.Message = "failed"
					planExec.Status = ExecutionStatusFailure
				}
			case wfapi.WorkflowError:
				job.Message = "sys error"
				planExec.Status = ExecutionStatusError
			}
		}
	}

	// When there is a status change, we update plan-execution to tcms
	if planExec.Status != oldStatus {
		if !wf.Status.StartedAt.IsZero() {
			startedAt := Timestamp(wf.Status.StartedAt.Unix())
			planExec.StartedAt = &startedAt
		}
		if !wf.Status.FinishedAt.IsZero() && planExec.Status != ExecutionStatusPrepareBackoff {
			completedAt := Timestamp(wf.Status.FinishedAt.Unix())
			planExec.CompletedAt = &completedAt
		}

		// update plan-execution to tcms
		planExec.Output = appendExecOutput(planExec.Output, job.Message)
		job.Message = ""
		if err := p.TCMS.PutPlanExecution(planExec); err != nil {
			klog.ErrorS(err, "failed to sync plan-execution status", "workflow", wf.Name)
		}

		if planExec.Status.IsCompleted() {
			err = p.TCMS.SyncTestBed(job.PlanExec.ID, job.Plan.TestBed.Name)
			if err != nil {
				klog.ErrorS(err, "failed to sync testbeds status", "testbedName", job.Plan.TestBed.Name)
			}
		}
	}
	return nil
}

func (p StatusSyncer) genPlanArgoURL(base, namespace, name string) string {
	baseURL, _ := url.Parse(base)
	baseURL.Path = path.Join(baseURL.Path, "workflows", namespace, name)
	return baseURL.String()
}

func existEnvPodsFailed(wf *wfapi.Workflow) bool {
	templates := make(map[string]*wfapi.Template, len(wf.Spec.Templates))
	for i := range wf.Spec.Templates {
		templates[wf.Spec.Templates[i].Name] = &wf.Spec.Templates[i]
	}
	for _, node := range wf.Status.Nodes {
		if node.Type != wfapi.NodeTypePod {
			continue
		}
		template := templates[node.TemplateName]
		role := template.Metadata.Labels["role"]
		if role == "setUp" || role == "tearDown" {
			if node.Phase.FailedOrError() {
				return true
			}
		}
	}
	return false
}

type SetupStatus int

const (
	SetupUnknown SetupStatus = iota
	SetupSuccess
	SetupFailRetryable
	SetupFailNonRetryable
)

func getSetupStatus(wf *wfapi.Workflow) (SetupStatus, *time.Time, *url.URL) {
	templates := make(map[string]*wfapi.Template, len(wf.Spec.Templates))
	for i := range wf.Spec.Templates {
		templates[wf.Spec.Templates[i].Name] = &wf.Spec.Templates[i]
	}
	var baseUrl string
	if wf.Status.ArtifactRepositoryRef != nil &&
		wf.Status.ArtifactRepositoryRef.ArtifactRepository != nil &&
		wf.Status.ArtifactRepositoryRef.ArtifactRepository.S3 != nil {
		s3 := wf.Status.ArtifactRepositoryRef.ArtifactRepository.S3
		if u, err := url.Parse(fmt.Sprintf("http://%s/%s", s3.Endpoint, s3.Bucket)); err == nil {
			baseUrl = u.String()
		}
	}
	for _, node := range wf.Status.Nodes {
		if node.Type != wfapi.NodeTypePod {
			continue
		}
		template := templates[node.TemplateName]
		if podRole := template.Metadata.Labels["role"]; podRole != "setUp" {
			continue
		}
		var logUrl *url.URL
		if node.Outputs != nil {
			for _, artifact := range node.Outputs.Artifacts {
				if artifact.S3 == nil {
					continue
				}
				s3 := artifact.S3
				if s3.Endpoint != "" && s3.Bucket != "" {
					u, err := url.Parse(fmt.Sprintf("http://%s/%s/%s", s3.Endpoint, s3.Bucket, s3.Key))
					if err != nil {
						continue
					}
					logUrl = u
				} else if baseUrl != "" {
					u, err := url.Parse(fmt.Sprintf("%s/%s", baseUrl, s3.Key))
					if err != nil {
						continue
					}
					logUrl = u
				}
				if artifact.Name == "main-logs" {
					break
				}
			}
		}

		switch node.Phase {
		case wfapi.NodeSucceeded:
			return SetupSuccess, &node.FinishedAt.Time, logUrl
		case wfapi.NodeError:
			return SetupFailNonRetryable, &node.FinishedAt.Time, logUrl
		case wfapi.NodeFailed:
			if exitCode := node.Outputs.ExitCode; exitCode != nil && *exitCode == "124" {
				return SetupFailRetryable, &node.FinishedAt.Time, logUrl
			}
			return SetupFailNonRetryable, &node.FinishedAt.Time, logUrl
		default:
			return SetupUnknown, &node.FinishedAt.Time, logUrl
		}
	}

	// never setup but workflow finished
	if wf.Status.Phase.Completed() {
		return SetupFailNonRetryable, &wf.Status.FinishedAt.Time, nil
	}

	return SetupUnknown, nil, nil
}

func readCaseStatus(node wfapi.NodeStatus) (ExecutionStatus, bool) {
	if node.Phase == wfapi.NodePending {
		return ExecutionStatusPending, true
	}
	if node.Phase == wfapi.NodeRunning {
		return ExecutionStatusRunning, true
	}
	if node.Phase == wfapi.NodeSkipped || node.Phase == wfapi.NodeOmitted {
		return ExecutionStatusSkipped, true
	}
	if node.Phase == wfapi.NodeFailed && strings.Contains(node.Message, "deadline") {
		return ExecutionStatusTimeout, true
	}
	if node.Outputs == nil || node.Outputs.ExitCode == nil {
		return "", false
	}
	switch *node.Outputs.ExitCode {
	case "0":
		return ExecutionStatusSuccess, true
	case "3":
		return ExecutionStatusSkipped, true
	case "4":
		return ExecutionStatusTimeout, true
	case "5":
		return ExecutionStatusCancelled, true
	default:
		return ExecutionStatusFailure, true
	}
}

func getNodeLogs(node wfapi.NodeStatus, baseURl string) map[string]*url.URL {
	if node.Outputs == nil {
		klog.Warningf("cannot find output")
		return nil
	}
	res := make(map[string]*url.URL)
	for _, artifact := range node.Outputs.Artifacts {
		if artifact.S3 == nil {
			continue
		}
		s3 := artifact.S3
		if s3.Endpoint != "" && s3.Bucket != "" {
			u, err := url.Parse(fmt.Sprintf("http://%s/%s/%s", s3.Endpoint, s3.Bucket, s3.Key))
			if err != nil {
				continue
			}
			res["log."+artifact.Name] = u
		} else if baseURl != "" {
			u, err := url.Parse(fmt.Sprintf("%s/%s", baseURl, s3.Key))
			if err != nil {
				continue
			}
			res["log."+artifact.Name] = u
		}
	}
	return res
}

func readLog(logUrl url.URL) string {
	resp, err := http.Get(logUrl.String())
	if err != nil {
		klog.Warningf("request output from %s: %v", logUrl.String(), err)
		return ""
	}
	defer resp.Body.Close()
	output, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		klog.Warningf("read output from %s: %v", logUrl.String(), err)
		return ""
	}
	return string(output)
}

type NaiveScheduler struct {
	TCMS common.TCMS
	Argo ArgoClientMgr
}

func (p *NaiveScheduler) Process(ctx context.Context, job *common.ArgoJob) error {
	client, err := p.Argo.ClientFor(job)
	if err != nil {
		return err
	}
	job = p.Argo.BindClient(job, client)
	err = submitTestJob(ctx, job, p.TCMS, client)
	if err != nil {
		klog.ErrorS(err, "submit job fail")
		e := SetErrorTestJob(ctx, job, err.Error(), p.TCMS)
		if e != nil {
			klog.ErrorS(e, "set skip fail")
		}
	}
	return nil
}

type Metric interface {
	Get() float64
	Inc(val float64)
}

type Canceler struct {
	TCMS common.TCMS
	Argo ArgoClientMgr
}

func updateCaseStatus(job *common.ArgoJob, tcms common.TCMS, wf *wfapi.Workflow, statusGetter func(node wfapi.NodeStatus) (ExecutionStatus, bool)) error {
	templates := make(map[string]*wfapi.Template, len(wf.Spec.Templates))
	for i := range wf.Spec.Templates {
		templates[wf.Spec.Templates[i].Name] = &wf.Spec.Templates[i]
	}
	for _, node := range wf.Status.Nodes {
		if node.Type != wfapi.NodeTypePod && node.Type != wfapi.NodeTypeSkipped {
			continue
		}

		template := templates[node.TemplateName]
		execID, err := strconv.ParseUint(template.Metadata.Labels["caseExecID"], 10, 64)
		if err != nil {
			if s := template.Metadata.Labels["caseExecID"]; len(s) > 0 {
				klog.Warningf("invalid case-execution id")
			}
			continue
		}

		exec, err := tcms.GetCaseExecution(ID(execID))
		if err != nil {
			klog.ErrorS(err, "get case-execution for update", "id", execID)
			continue
		}

		status, ok := statusGetter(node)
		if !ok {
			klog.InfoS("case-execution is not ready for syncing", "id", execID)
			continue
		}
		if exec.Status == status {
			continue
		}
		exec.Status = status

		if !node.StartedAt.IsZero() {
			startedAt := Timestamp(node.StartedAt.Unix())
			exec.StartedAt = &startedAt
		}
		exec.Meta = map[string]string{
			"argo.workflow": node.BoundaryID,
			"argo.pod":      node.ID,
			"argo.name":     node.Name,
		}
		exec.ArgoURL = genCaseArgoUrl(job, node.ID)
		if exec.Status.IsCompleted() {
			var baseUrl string
			if wf.Status.ArtifactRepositoryRef != nil &&
				wf.Status.ArtifactRepositoryRef.ArtifactRepository != nil &&
				wf.Status.ArtifactRepositoryRef.ArtifactRepository.S3 != nil {
				s3 := wf.Status.ArtifactRepositoryRef.ArtifactRepository.S3
				if u, err := url.Parse(fmt.Sprintf("http://%s/%s", s3.Endpoint, s3.Bucket)); err == nil {
					baseUrl = u.String()
				}
			}
			logLinks := getNodeLogs(node, baseUrl)
			for k, v := range logLinks {
				exec.Meta[k] = v.String()
			}
			if !node.FinishedAt.IsZero() {
				completedAt := Timestamp(node.FinishedAt.Unix())
				exec.CompletedAt = &completedAt
			}
			if outputLink, ok := logLinks["log.output"]; ok {
				exec.Output = readLog(*outputLink)
			}
		}

		if err := tcms.PutCaseExecution(exec); err != nil {
			klog.ErrorS(err, "set case-execution")
		}
	}
	return nil
}

func genCaseArgoUrl(job *common.ArgoJob, nodeId string) string {
	if job == nil {
		return ""
	}
	if job.PlanExec.ArgoURL != "" {
		baseURL, _ := url.Parse(job.PlanExec.ArgoURL)
		query := baseURL.Query()
		query.Set("nodeId", nodeId)
		baseURL.RawQuery = query.Encode()
		return baseURL.String()
	} else {
		return ""
	}

}

func (p *Canceler) doProcess(ctx context.Context, job *common.ArgoJob) error {
	if job.Workflow == nil {
		// Argo workflow have not been started yet
		// Nothing need to clean up
		return nil
	}
	argo, err := p.Argo.ClientFor(job)
	if err != nil {
		return err
	}
	wf, err := argo.Stop(ctx, job.Workflow.Name)
	if err != nil {
		return err
	}
	job.Workflow = wf
	// Handle testbed
	testbed, err := p.TCMS.GetTestbedInfo(wf.Annotations["testbed"])
	if err != nil {
		klog.ErrorS(err, "failed to get testbed", "testbed name", wf.Annotations["testbed"])
	} else {
		if testbed.Status == "INITIATING" {
			// PutTestBedInfo will release testbed automatically when set to ERROR
			tb := &TestBedInfo{
				Status:          "ERROR",
				PlanExecutionID: job.PlanExec.ID,
				Name:            testbed.Name,
			}
			err = p.TCMS.PutTestBedInfo(tb)
			if err != nil {
				klog.Errorf("PutTestBedInfo() failed: %s", err.Error())
			}
		}
	}

	// todo ignore update case may fail, because cases status is badly ensured, fix later
	err = updateCaseStatus(job, p.TCMS, wf, func(_ wfapi.NodeStatus) (ExecutionStatus, bool) {
		return ExecutionStatusCancelled, true
	})
	if err != nil {
		klog.ErrorS(err, "failed to update case executions status", "workflow", wf.Name)
	}
	return nil
}

func (p *Canceler) Process(ctx context.Context, job *common.ArgoJob) error {
	err := p.doProcess(ctx, job)
	if err != nil {
		klog.ErrorS(err, "can not cancel job", "execID", job.ID())
		return err
	}
	if err := SetCanceledTestJob(ctx, job, job.Message, "", p.TCMS); err != nil {
		klog.ErrorS(err, "cannot set plan-execution after cancelling workflow", "execID", job.PlanExec.ID)
		return err
	}
	return nil
}