package common

import (
	"bytes"
	"context"
	"database/sql"
	"encoding/json"
	"fmt"
	"strconv"
	"time"

	"github.com/pingcap/test-infra/common/model/tcms/taskruns"

	"k8s.io/klog/v2"

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

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

type ArgoJobProcessor interface {
	Process(ctx context.Context, job *ArgoJob) error
}

type ArgoJobStore interface {
	PutJob(ctx context.Context, job *ArgoJob, action OnDupAction) error
	GetJob(ctx context.Context, id ID) (*ArgoJob, error)
	EditJob(ctx context.Context, id ID, p ArgoJobProcessor) error
	ListJobs(ctx context.Context, conds ListJobConditions) ([]*ArgoJob, error)
}

type ArgoJobStatus string

const (
	ArgoTestJobPending    ArgoJobStatus = "PENDING"
	ArgoTestJobCancelling ArgoJobStatus = "CANCELLING"
	ArgoTestJobRunning    ArgoJobStatus = "RUNNING"
	ArgoTestJobDone       ArgoJobStatus = "DONE"
)

type OnDupAction int

const (
	OnDupUpdate OnDupAction = iota
	OnDupIgnore
	OnDupUpdateTcmsFields
)

type ListJobConditions struct {
	Limit  int
	Status ArgoJobStatus
}

func buildListJobSQL(conds ListJobConditions) (string, []interface{}) {
	var args []interface{}
	buf := bytes.NewBufferString("select plan_exec, plan, workflow, status, message from test_job")
	if len(conds.Status) > 0 {
		buf.WriteString(" where status = ?")
		args = append(args, conds.Status)
	}
	buf.WriteString(" order by created_at")
	if conds.Limit > 0 {
		buf.WriteString(" limit " + strconv.Itoa(conds.Limit))
	}
	return buf.String(), args
}

type ArgoJob struct {
	PlanExec     *TestPlanExecution
	Plan         *TestPlan
	Workflow     *wf.Workflow
	Status       ArgoJobStatus
	EPriority    ExecPriority
	Message      string
	TestbedSetup *taskruns.TestbedSetup
}

func (job *ArgoJob) ID() ID { return job.PlanExec.ID }

func (job *ArgoJob) UpdateStatus(w *wf.Workflow) *ArgoJob {
	if w != nil {
		job.Workflow = w
	}
	if w == nil && job.PlanExec.Status == ExecutionStatusPrepareBackoff {
		return job
	}
	if job.Status == ArgoTestJobDone || job.Workflow == nil {
		return job
	}
	phase := wf.WorkflowPhase(job.Workflow.Labels["workflows.argoproj.io/phase"])
	if phase == wf.WorkflowUnknown || phase == wf.WorkflowPending || phase == wf.WorkflowRunning {
		job.Status = ArgoTestJobRunning
	} else {
		job.Status = ArgoTestJobDone
	}
	return job
}

func (job *ArgoJob) String() string {
	return fmt.Sprintf("job-%v", job.PlanExec.ID)
}

func NewJobStore(db *sql.DB) ArgoJobStore {
	return &sqlJobStore{db: db}
}

type sqlJobStore struct {
	db *sql.DB
}

func (store *sqlJobStore) PutJob(ctx context.Context, job *ArgoJob, action OnDupAction) error {
	return insertJob(ctx, job, action, store.db.ExecContext)
}

func (store *sqlJobStore) GetJob(ctx context.Context, id ID) (*ArgoJob, error) {
	row := store.db.QueryRowContext(ctx, "select plan_exec, plan, workflow, status, message from test_job where id = ?", id)
	job, err := scanJob(row.Scan)
	if err != nil {
		return nil, err
	}
	return job.UpdateStatus(nil), nil
}

func (store *sqlJobStore) EditJob(ctx context.Context, id ID, p ArgoJobProcessor) error {
	if p == nil {
		klog.Warning("edit job by noop")
		return nil
	}
	tx, err := store.db.Begin()
	if err != nil {
		return errors.WithStack(err)
	}
	msg := ""
	defer func() {
		tx.Rollback()
		if len(msg) > 0 {
			store.db.ExecContext(ctx, "update test_job set message = ?, updated_at = unix_timestamp() where id = ?", msg, id)
		}
	}()

	row := tx.QueryRowContext(ctx, "select plan_exec, plan, workflow, status, message from test_job where id = ? for update", id)
	job, err := scanJob(row.Scan)
	if err != nil {
		return err
	}
	err = p.Process(ctx, job)
	if err != nil {
		msg = "process error: " + err.Error()
		return err
	}
	err = insertJob(ctx, job, OnDupUpdate, tx.ExecContext)
	if err != nil {
		msg = "insert error: " + err.Error()
		return err
	}
	err = tx.Commit()
	if err != nil {
		msg = "commit error: " + err.Error()
		return errors.WithStack(err)
	}
	return nil
}

func (store *sqlJobStore) ListJobs(ctx context.Context, conds ListJobConditions) ([]*ArgoJob, error) {
	query, args := buildListJobSQL(conds)
	jobs := make([]*ArgoJob, 0, 8)
	rows, err := store.db.QueryContext(ctx, query, args...)
	if err != nil {
		return nil, errors.WithStack(err)
	}
	defer rows.Close()
	for rows.Next() {
		job, err := scanJob(rows.Scan)
		if err != nil {
			return nil, err
		}
		jobs = append(jobs, job.UpdateStatus(nil))
	}
	if err = rows.Err(); err != nil {
		return nil, errors.WithStack(err)
	}
	return jobs, nil
}

func scanJob(scanFn func(dest ...interface{}) error) (*ArgoJob, error) {
	var (
		planExecBytes []byte
		planBytes     []byte
		workflowBytes []byte
		job           = &ArgoJob{PlanExec: new(TestPlanExecution), Plan: new(TestPlan)}
	)
	err := scanFn(&planExecBytes, &planBytes, &workflowBytes, &job.Status, &job.Message)
	if err != nil {
		return nil, errors.WithStack(err)
	}
	if err = json.Unmarshal(planExecBytes, job.PlanExec); err != nil {
		klog.Errorf("unmarshal plan execution (%q): %+v ", string(planExecBytes), err)
		return nil, errors.WithStack(err)
	}
	if err = json.Unmarshal(planBytes, job.Plan); err != nil {
		klog.Errorf("unmarshal plan  (%q): %+v ", string(planBytes), err)
		return nil, errors.WithStack(err)
	}
	if err = json.Unmarshal(workflowBytes, &job.Workflow); err != nil {
		klog.Errorf("unmarshal workflow (%q): %+v ", string(workflowBytes), err)
		return nil, errors.WithStack(err)
	}
	return job, nil
}

func insertJob(ctx context.Context, job *ArgoJob, onDupAction OnDupAction, execFn func(ctx context.Context, query string, args ...interface{}) (sql.Result, error)) error {
	if job.PlanExec == nil {
		return errors.New("invalid test job: plan execution are required")
	}
	id := job.PlanExec.ID
	planExecBytes, err := json.Marshal(job.PlanExec)
	if err != nil {
		return errors.WithStack(err)
	}
	planBytes, err := json.Marshal(job.Plan)
	if err != nil {
		return errors.WithStack(err)
	}
	workflowBytes, err := json.Marshal(job.Workflow)
	if err != nil {
		return errors.WithStack(err)
	}
	epoch := time.Now().Unix()

	switch onDupAction {
	case OnDupIgnore:
		_, err = execFn(ctx, "insert ignore test_job(id, plan_exec, plan, workflow, status, message, created_at, updated_at) values (?, ?, ?, ?, ?, ?, ?, ?) ",
			id, planExecBytes, planBytes, workflowBytes, job.Status, job.Message, epoch, epoch)
	case OnDupUpdate:
		_, err = execFn(ctx, "insert into test_job(id, plan_exec, plan, workflow, status, message, created_at, updated_at) values (?,?, ?, ?, ?, ?, ?, ?) "+
			"on duplicate key update plan_exec=values(plan_exec), plan=values(plan), workflow=values(workflow), status=values(status), message=values(message), updated_at=values(updated_at)",
			id, planExecBytes, planBytes, workflowBytes, job.Status, job.Message, epoch, epoch)
	case OnDupUpdateTcmsFields:
		_, err = execFn(ctx, "insert into test_job(id, plan_exec, plan, workflow, status, message, created_at, updated_at) values (?,?, ?, ?, ?, ?, ?, ?) "+
			"on duplicate key update plan_exec=values(plan_exec), plan=values(plan), status=values(status), updated_at=values(updated_at)",
			id, planExecBytes, planBytes, workflowBytes, job.Status, job.Message, epoch, epoch)
	default:
		err = fmt.Errorf("unknown onDupAction: %d", int(onDupAction))
	}
	if err != nil {
		return errors.WithStack(err)
	}
	return nil
}