package workflow

import (
	"context"

	"github.com/go-sql-driver/mysql"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/tidbcloud/fastflow/pkg/entity"
	"github.com/tidbcloud/fastflow/pkg/entity/run"
	"github.com/tidbcloud/fastflow/pkg/mod"
	"github.com/tidbcloud/fastflow/store"
	"github.com/tidbcloud/pkgv2/db"

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

type workflowOnlyDbImpl struct {
	provider db.Provider
	realImpl *workflowImpl
}

func NewWorkflowOnlyDb(
	cfg *mysql.Config,
	provider db.Provider,
	workerId uint16,
	idGenerator idgenerator.IDGenerator,
) (Workflow, error) {
	if err := newStore(cfg); err != nil {
		return nil, err
	}

	store.InitFlakeGenerator(workerId)

	return &workflowOnlyDbImpl{
		provider: provider,
		realImpl: &workflowImpl{provider: provider, idGenerator: idGenerator},
	}, nil
}

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

func (impl *workflowOnlyDbImpl) IsLeader() bool {
	return true
}

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

func (impl *workflowOnlyDbImpl) RegisterDag(dag *entity.Dag, actions []run.Action) error {
	return registerDag(dag, actions)
}

func (impl *workflowOnlyDbImpl) Start() error {
	return nil
}

func (impl *workflowOnlyDbImpl) DeleteDagInstances(
	ctx context.Context,
	bizId string,
	excludeDags map[string]bool,
) error {
	return impl.realImpl.DeleteDagInstances(ctx, bizId, excludeDags)
}

func (impl *workflowOnlyDbImpl) DeleteDagInstancesWithDagId(ctx context.Context, bizId, dagId string) error {
	return impl.realImpl.DeleteDagInstancesWithDagId(ctx, bizId, dagId)
}

func (impl *workflowOnlyDbImpl) Close() {
	mod.GetStore().Close()
}

func (impl *workflowOnlyDbImpl) ListRunningDagInstances(
	ctx context.Context,
	dagIds []string,
	bizId string,
) ([]*entity.DagInstance, error) {
	return impl.realImpl.ListRunningDagInstances(ctx, dagIds, bizId)
}

func (impl *workflowOnlyDbImpl) ListRunningDagInstancesByBizIds(
	ctx context.Context,
	dagIds []string,
	bizIds []string,
) ([]*entity.DagInstance, error) {
	return impl.realImpl.ListRunningDagInstancesByBizIds(ctx, dagIds, bizIds)
}

func (impl *workflowOnlyDbImpl) ListAllRunningDagInstances() ([]*entity.DagInstance, error) {
	return impl.realImpl.ListAllRunningDagInstances()
}

func (impl *workflowOnlyDbImpl) ListAllRunningTaskInstances() ([]*entity.TaskInstance, error) {
	return impl.realImpl.ListAllRunningTaskInstances()
}

func (impl *workflowOnlyDbImpl) SetDagInstanceLifecycleHook(hooks entity.DagInstanceLifecycleHook) {
	impl.realImpl.SetDagInstanceLifecycleHook(hooks)
}

func (impl *workflowOnlyDbImpl) RegisterDynamicDagThenRun(
	ctx context.Context,
	prefix string,
	opt *CreateDynamicDagOption,
) error {
	return impl.realImpl.RegisterDynamicDagThenRun(ctx, prefix, opt)
}

func (impl *workflowOnlyDbImpl) ListRunningDagInstancesByNamePrefix(
	ctx context.Context,
	namePrefix, bizId string,
) ([]*entity.DagInstance, error) {
	return impl.realImpl.ListRunningDagInstancesByNamePrefix(ctx, namePrefix, bizId)
}

func (impl *workflowOnlyDbImpl) ListRunningDagInstancesByOrgId(
	ctx context.Context,
	dagIds []string,
	orgId string,
) ([]*entity.DagInstance, error) {
	return impl.realImpl.ListRunningDagInstancesByOrgId(ctx, dagIds, orgId)
}

func (impl *workflowOnlyDbImpl) RegisterMetrics(registry prometheus.Registerer) {
	impl.realImpl.RegisterMetrics(registry)
}
