package workflow

import (
	"context"
	"errors"
	"fmt"
	"log"
	"os"
	"strings"
	"sync"
	"time"

	"github.com/go-sql-driver/mysql"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/samber/lo"
	"github.com/tidbcloud/fastflow"
	mysqlkeeper "github.com/tidbcloud/fastflow/keeper/mysql"
	"github.com/tidbcloud/fastflow/pkg/entity"
	"github.com/tidbcloud/fastflow/pkg/entity/run"
	"github.com/tidbcloud/fastflow/pkg/exporter"
	"github.com/tidbcloud/fastflow/pkg/mod"
	"github.com/tidbcloud/fastflow/pkg/utils/data"
	mysqlstore "github.com/tidbcloud/fastflow/store/mysql"
	"github.com/tidbcloud/pkgv2/db"
	"github.com/tidbcloud/tiinf/monitor/tilog"
	"go.uber.org/zap"
	"gorm.io/gorm"
	"gorm.io/gorm/logger"

	"github.com/tidbcloud/tidb-management-service/utilities/grpc/metadata"
	"github.com/tidbcloud/tidb-management-service/utilities/idgenerator"
)

var (
	dagInstanceRunningStatus = []entity.DagInstanceStatus{
		entity.DagInstanceStatusInit,
		entity.DagInstanceStatusScheduled,
		entity.DagInstanceStatusRunning,
		entity.DagInstanceStatusFailed,
	}

	taskInstanceRunningStatus = []entity.TaskInstanceStatus{
		entity.TaskInstanceStatusInit,
		entity.TaskInstanceStatusRunning,
		entity.TaskInstanceStatusEnding,
		entity.TaskInstanceStatusFailed,
		entity.TaskInstanceStatusRetrying,
	}

	gormConfig = &gorm.Config{
		Logger: logger.New(log.New(os.Stdout, "\r\n", log.LstdFlags), logger.Config{
			SlowThreshold:             300 * time.Millisecond,
			LogLevel:                  logger.Warn,
			IgnoreRecordNotFoundError: true,
			Colorful:                  true,
		}),
	}

	executorTimeout = time.Second * 60

	registerDagMu sync.Mutex
)

type workflowImpl struct {
	wd          *watchDog
	provider    db.Provider
	commander   mod.Commander
	mysqlCfg    *mysql.Config
	id          ID
	idGenerator idgenerator.IDGenerator
	mode        Mode
	competitive CompetitiveOption
}

func NewWorkflow(
	id ID,
	cfg *mysql.Config,
	provider db.TransProvider,
	idGenerator idgenerator.IDGenerator,
	opts ...NewWorkflowOption,
) (Workflow, error) {
	if id == "" {
		return nil, errors.New("need work flow id")
	}
	log.Printf("start workflow, id: %s", id)

	commander := &mod.DefCommander{}
	mod.SetCommander(commander)

	wfCfg := &NewWorkflowConfig{}
	for _, opt := range opts {
		if opt != nil {
			opt(wfCfg)
		}
	}
	if wfCfg.WorkerMode == "" {
		wfCfg.WorkerMode = ModeDispatcher
	}
	wd := newWatchDog(provider, wfCfg.WorkerMode)

	// Start the store.
	if err := newStore(cfg); err != nil {
		return nil, err
	}

	return &workflowImpl{
		wd:          wd,
		provider:    provider,
		commander:   commander,
		mysqlCfg:    cfg,
		id:          id,
		idGenerator: idGenerator,
		mode:        wfCfg.WorkerMode,
		competitive: wfCfg.Competitive,
	}, nil
}

func newStore(cfg *mysql.Config) error {
	store := mysqlstore.NewStore(&mysqlstore.StoreOption{
		MySQLConfig:     cfg,
		MigrationSwitch: true,
		GormConfig:      gormConfig,
	})
	if err := store.Init(); err != nil {
		return err
	}
	if err := validateFastflowHasCmdSchema(cfg); err != nil {
		store.Close()
		return err
	}
	mod.SetStore(store)

	return nil
}

func (impl *workflowImpl) RunDag(ctx context.Context, dagId, orgId, projectId string, specVar map[string]string) error {
	return impl.runDag(ctx, dagId, orgId, projectId, "", specVar, impl.provider)
}

func (impl *workflowImpl) IsLeader() bool {
	if impl.mode == ModeCompetitive {
		return true
	}
	return mod.GetKeeper().IsLeader()
}

func (impl *workflowImpl) RunDagWithBizId(
	ctx context.Context,
	dagId, orgId, projectId, bizId string,
	specVar map[string]string,
) error {
	return impl.runDag(ctx, dagId, orgId, projectId, bizId, specVar, impl.provider)
}

func (impl *workflowImpl) getDag(ctx context.Context, dagId string) (*entity.Dag, error) {
	dag := &entity.Dag{}
	err := impl.provider.ReadDB(ctx).Transaction(func(tx *gorm.DB) error {
		return tx.Where("id = ?", dagId).First(dag).Error
	})
	if err != nil {
		if errors.Is(err, gorm.ErrRecordNotFound) {
			return nil, fmt.Errorf("key[ %s ] not found: %w", dagId, data.ErrDataNotFound)
		}
		tilog.Error(ctx, "get dag %s failed: %s", zap.String("dag", dagId), zap.Error(err))
		return nil, fmt.Errorf("get dag failed: %w", err)
	}
	return dag, nil
}

func (impl *workflowImpl) runDag(
	ctx context.Context,
	dagId, orgId, projectId, bizId string,
	specVar map[string]string,
	provider db.Provider,
) error {
	dag, err := impl.getDag(ctx, dagId)
	if err != nil {
		return err
	}

	if _, ok := specVar[KeyTraceId]; !ok {
		specVar[KeyTraceId] = GetTraceId(ctx)
	}

	if _, ok := specVar[KeyRequestMeta]; !ok {
		requestMetaData, err := metadata.MarshalRequestMetaFromCtx(ctx)
		if err != nil {
			return err
		}

		specVar[KeyRequestMeta] = string(requestMetaData)
	}

	specVar[KeyStartTimeInSec] = fmt.Sprintf("%d", time.Now().Unix())
	dagIns, err := dag.Run(entity.TriggerManually, specVar)
	if err != nil {
		return err
	}

	if dagIns.Vars == nil {
		dagIns.Vars = make(entity.DagInstanceVars)
	}
	if value, ok := dagIns.Vars[KeyComponentOperations]; !ok || value.Value == "" {
		if newValue, ok := specVar[KeyComponentOperations]; ok {
			dagIns.Vars[KeyComponentOperations] = entity.DagInstanceVar{
				Value: newValue,
			}
		}
	}

	tx := provider.WriteDB(ctx)

	dagIns.Initial()
	if dagIns.ShareData == nil {
		dagIns.ShareData = &entity.ShareData{}
	}
	if dagIns.ShareData.Dict == nil {
		dagIns.ShareData.Dict = map[string]string{}
	}
	if err := tx.Create(dagIns).Error; err != nil {
		return err
	}

	if orgId != "" {
		dagIns.Tags = append(dagIns.Tags, entity.DagInstanceTag{
			Key:   OrgIdKey,
			Value: orgId,
		})
	}

	if projectId != "" {
		dagIns.Tags = append(dagIns.Tags, entity.DagInstanceTag{
			Key:   ProjectKey,
			Value: projectId,
		})
	}

	if bizId != "" {
		dagIns.Tags = append(dagIns.Tags, entity.DagInstanceTag{
			Key:   BizIdKey,
			Value: bizId,
		})
	}

	for _, tag := range dagIns.Tags {
		tag.Initial()
		tag.DagInsId = dagIns.ID
		err := tx.Create(tag).Error
		if err != nil {
			return err
		}
	}
	return nil
}

func (impl *workflowImpl) RegisterDag(dag *entity.Dag, actions []run.Action) error {
	if strings.HasPrefix(dag.Name, DynamicDagPrefix) {
		return errors.New("duplicate dynamic dag prefix")
	}
	dag.Vars[KeyStartTimeInSec] = entity.DagVar{}
	dag.Vars[KeyTraceId] = entity.DagVar{}
	dag.Vars[KeyRequestMeta] = entity.DagVar{}
	return registerDag(dag, actions)
}

func (impl *workflowImpl) DeleteDagInstances(ctx context.Context, bizId string, excludeDags map[string]bool) error {
	tx := impl.provider.WriteDB(ctx)

	instances, err := mod.GetStore().ListDagInstance(&mod.ListDagInstanceInput{
		Status: dagInstanceRunningStatus,
		Tags: map[string]string{
			BizIdKey: bizId,
		},
	})
	if err != nil {
		return err
	}

	for i := range instances {
		if v, ok := excludeDags[instances[0].DagID]; ok && v {
			continue
		}
		// double check
		_, ok := lo.Find(instances[0].Tags, func(item entity.DagInstanceTag) bool {
			return item.Key == BizIdKey && item.Value == bizId
		})
		if !ok {
			return fmt.Errorf("unexpect instance %s tag", instances[0].ID)
		}
		instances[i].Success()
		if err := tx.Delete(&entity.DagInstance{}, instances[i].ID).Error; err != nil {
			return err
		}
	}

	return nil
}

func (impl *workflowImpl) DeleteDagInstancesWithDagId(ctx context.Context, bizId, dagId string) error {
	tx := impl.provider.WriteDB(ctx)

	instances, err := mod.GetStore().ListDagInstance(&mod.ListDagInstanceInput{
		Status: dagInstanceRunningStatus,
		Tags: map[string]string{
			BizIdKey: bizId,
		},
	})
	if err != nil {
		return err
	}

	for i := range instances {
		if instances[0].DagID != dagId {
			continue
		}
		// double check
		_, ok := lo.Find(instances[0].Tags, func(item entity.DagInstanceTag) bool {
			return item.Key == BizIdKey && item.Value == bizId
		})
		if !ok {
			return fmt.Errorf("unexpect instance %s tag", instances[0].ID)
		}
		instances[i].Success()
		if err := tx.Delete(&entity.DagInstance{}, instances[i].ID).Error; err != nil {
			return err
		}
	}
	return nil
}

func (impl *workflowImpl) ListRunningDagInstances(
	ctx context.Context,
	dagIds []string,
	bizId string,
) ([]*entity.DagInstance, error) {
	if len(dagIds) == 0 || bizId == "" {
		return make([]*entity.DagInstance, 0), nil
	}

	tx := impl.provider.ReadDB(ctx)

	var instances []*entity.DagInstance
	if err := tx.Model(&entity.DagInstance{}).
		Where("status in (?) and dag_id in (?)", dagInstanceRunningStatus, dagIds).
		Joins("join dag_instance_tags on dag_instances.id = dag_instance_tags.dag_ins_id").
		Where("dag_instance_tags.`key` = ? and dag_instance_tags.value=?", "biz-id", bizId).
		Find(&instances).Error; err != nil {
		return nil, err
	}

	if err := fillInstanceTags(ctx, tx, instances); err != nil {
		return nil, err
	}

	return instances, nil
}

func (impl *workflowImpl) ListRunningDagInstancesByBizIds(
	ctx context.Context,
	dagIds []string,
	bizIds []string,
) ([]*entity.DagInstance, error) {
	if len(dagIds) == 0 || len(bizIds) == 0 {
		return make([]*entity.DagInstance, 0), nil
	}

	tx := impl.provider.ReadDB(ctx)

	var instances []*entity.DagInstance
	if err := tx.Model(&entity.DagInstance{}).
		Where("status in (?) and dag_id in (?)", dagInstanceRunningStatus, dagIds).
		Joins("join dag_instance_tags on dag_instances.id = dag_instance_tags.dag_ins_id").
		Where("dag_instance_tags.`key` = ? and dag_instance_tags.value in (?)", BizIdKey, bizIds).
		Find(&instances).Error; err != nil {
		return nil, err
	}

	if err := fillInstanceTags(ctx, tx, instances); err != nil {
		return nil, err
	}

	return instances, nil
}

func (impl *workflowImpl) ListRunningDagInstancesByNamePrefix(
	ctx context.Context,
	namePrefix, bizId string,
) ([]*entity.DagInstance, error) {
	if namePrefix == "" || bizId == "" {
		return make([]*entity.DagInstance, 0), nil
	}

	tx := impl.provider.ReadDB(ctx)

	var instances []*entity.DagInstance
	if err := tx.Model(&entity.DagInstance{}).
		Where("status in (?) and dag_id like ?", dagInstanceRunningStatus, namePrefix+"%").
		Joins("join dag_instance_tags on dag_instances.id = dag_instance_tags.dag_ins_id").
		Where("dag_instance_tags.`key` = ? and dag_instance_tags.value=?", "biz-id", bizId).
		Find(&instances).Error; err != nil {
		return nil, err
	}

	if err := fillInstanceTags(ctx, tx, instances); err != nil {
		return nil, err
	}

	return instances, nil
}

func (impl *workflowImpl) ListRunningDagInstancesByOrgId(
	ctx context.Context,
	dagIds []string,
	orgId string,
) ([]*entity.DagInstance, error) {
	if len(dagIds) == 0 || orgId == "" {
		return make([]*entity.DagInstance, 0), nil
	}

	tx := impl.provider.ReadDB(ctx)

	var instances []*entity.DagInstance
	if err := tx.Model(&entity.DagInstance{}).
		Where("status in (?) and dag_id in (?)", dagInstanceRunningStatus, dagIds).
		Joins("join dag_instance_tags on dag_instances.id = dag_instance_tags.dag_ins_id").
		Where("dag_instance_tags.`key` = ? and dag_instance_tags.value=?", OrgIdKey, orgId).
		Find(&instances).Error; err != nil {
		return nil, err
	}

	if err := fillInstanceTags(ctx, tx, instances); err != nil {
		return nil, err
	}

	return instances, nil
}

func fillInstanceTags(ctx context.Context, tx *gorm.DB, dagInstances []*entity.DagInstance) error {
	if len(dagInstances) == 0 {
		return nil
	}

	dagInsIds := make([]string, len(dagInstances))
	for i, dagInstance := range dagInstances {
		dagInsIds[i] = dagInstance.ID
	}

	var tags []entity.DagInstanceTag
	if err := tx.Where("dag_ins_id IN ?", dagInsIds).Find(&tags).Error; err != nil {
		tilog.Error(ctx, "list dag instance tags error", zap.Error(err))
		return err
	}

	tagMap := make(map[string][]entity.DagInstanceTag, len(dagInstances))
	for i := range tags {
		tag := tags[i]
		tagMap[tag.DagInsId] = append(tagMap[tag.DagInsId], tag)
	}

	for _, dagInstance := range dagInstances {
		if t, ok := tagMap[dagInstance.ID]; ok {
			dagInstance.Tags = t
		}
	}
	return nil
}

// ListAllRunningDagInstances returns running dag instances.
// The key of the result maps is the biz id. In dedicated global service, the biz id is the cluster id.
func (impl *workflowImpl) ListAllRunningDagInstances() ([]*entity.DagInstance, error) {
	dagInstances, err := mod.GetStore().ListDagInstance(&mod.ListDagInstanceInput{
		Status: dagInstanceRunningStatus,
	})
	if err != nil {
		return nil, err
	}
	return dagInstances, nil
}

// ListAllRunningTaskInstances returns running task instances. The key of the result maps is the dag instance id.
func (impl *workflowImpl) ListAllRunningTaskInstances() ([]*entity.TaskInstance, error) {
	taskInstances, err := mod.GetStore().ListTaskInstance(&mod.ListTaskInstanceInput{
		Status: taskInstanceRunningStatus,
	})
	if err != nil {
		return nil, err
	}
	return taskInstances, nil
}

func registerDag(dag *entity.Dag, actions []run.Action) error {
	registerDagMu.Lock()
	defer registerDagMu.Unlock()

	fastflow.RegisterAction(actions)
	oldDag, err := mod.GetStore().GetDag(dag.ID)
	if errors.Is(err, data.ErrDataNotFound) {
		if err := mod.GetStore().CreateDag(dag); err != nil {
			return err
		}
	}
	if oldDag != nil {
		if err := mod.GetStore().UpdateDag(dag); err != nil {
			return err
		}
	}

	return nil
}

// Start starts the workflow.
//
// Some of important default configurations are (preset in `fastflow.Init()`):
//
//	Dispatcher and watchdog are goroutine in leader that schedules dag instance to worker and ensure fairness.
//		opt.DagScheduleTimeout = 15 * time.Second
//	One dag instance will be scheduled to one worker.
//	Parser is goroutine in worker that poll store for dag instances scheduled to the worker and forward them to executor.
//		opt.ParserWorkersCnt = 100
//	All actions in the dag instance will be executed on the same worker and across executor goroutines.
//	Executor is goroutine in worker that execute actions workload.
//		opt.ExecutorTimeout = 30 * time.Second
//		opt.ExecutorWorkerCnt = 1000
func (impl *workflowImpl) Start() error {
	effectiveMode := impl.mode
	if effectiveMode == "" {
		effectiveMode = ModeDispatcher
	}
	impl.wd.mode = effectiveMode

	// Start the keeper.
	keeper := mysqlkeeper.NewKeeper(&mysqlkeeper.KeeperOption{
		Key:             string(impl.id),
		MySQLConfig:     impl.mysqlCfg,
		MigrationSwitch: true,
		GormConfig:      gormConfig,
		WatcherFlag:     effectiveMode == ModeCompetitive,
	})
	if err := keeper.Init(); err != nil {
		return err
	}
	mod.SetKeeper(keeper)

	initOpt := &fastflow.InitialOption{
		Keeper: mod.GetKeeper(),
		Store:  mod.GetStore(),
		Mode:   mod.WorkerMode(effectiveMode),
		Competitive: mod.CompetitiveOption{
			ClaimLimit:        impl.competitive.ClaimLimit,
			ClaimInterval:     impl.competitive.ClaimInterval,
			LeaseStaleTTL:     impl.competitive.LeaseStaleTTL,
			LeaseReclaimLimit: impl.competitive.LeaseReclaimLimit,
		},
		ExecutorTimeout:   executorTimeout,
		ExecutorWorkerCnt: 100,
	}
	err := fastflow.Init(initOpt)
	if err != nil {
		return err
	}

	// wait for fastflow initialization to complete
	// then start the watch dog
	// this is to avoid problems with data race
	// fastflow.Init will SetStore
	// watch dog will GetStore
	go func() {
		impl.wd.start()
	}()

	return nil
}

func (impl *workflowImpl) Close() {
	mod.GetKeeper().Close()
	mod.GetStore().Close()
	impl.wd.close()
}

func (impl *workflowImpl) SetDagInstanceLifecycleHook(hooks entity.DagInstanceLifecycleHook) {
	entity.HookDagInstance = hooks
}

func (impl *workflowImpl) RegisterMetrics(registry prometheus.Registerer) {
	c := newWorkflowCollector(impl)
	fastflowExecutorColl := &exporter.ExecutorCollector{}
	fastflowLeaderColl := &exporter.LeaderCollector{}
	registry.MustRegister(c, fastflowExecutorColl, fastflowLeaderColl)
}