package impl

import (
	"errors"
	"fmt"
	"os"
	"sync"

	"github.com/go-sql-driver/mysql"
	"github.com/shiningrush/fastflow"
	mysqlkeeper "github.com/shiningrush/fastflow/keeper/mysql"
	"github.com/shiningrush/fastflow/pkg/entity"
	"github.com/shiningrush/fastflow/pkg/entity/run"
	"github.com/shiningrush/fastflow/pkg/mod"
	"github.com/shiningrush/fastflow/pkg/utils/data"
	mysqlstore "github.com/shiningrush/fastflow/store/mysql"
	"github.com/tidbcloud/serverless-service/internal/config"
	appworkflow "github.com/tidbcloud/serverless-service/internal/workflow"
)

const (
	parserWorkersCnt   = 32
	executorWorkersCnt = 128
)

type workflowImpl struct {
	store     mod.Store
	keeper    mod.Keeper
	commander mod.Commander

	startOnce sync.Once
	startErr  error
}

func NewWorkflow() (appworkflow.Workflow, error) {
	isExecutor := config.IsWorkflowMode()
	// MigrationSwitch is always true so that *any* pod that initializes
	// fastflow (http/grpc/worker/br-worker write DAGs via RunDagWithTags;
	// workflow runs the executor) creates the fastflow tables on first
	// boot if they don't exist. The underlying CREATE TABLE IF NOT EXISTS
	// is idempotent; race between replicas is harmless.
	store, err := initStore(true)
	if err != nil {
		return nil, err
	}

	keeper, err := initKeeper(isExecutor)
	if err != nil {
		store.Close()
		return nil, err
	}

	commander := initCommander()
	return &workflowImpl{
		store:     store,
		keeper:    keeper,
		commander: commander,
	}, nil
}

func initStore(migrate bool) (mod.Store, error) {
	cfg := config.Global().MysqlConfig
	store := mysqlstore.NewStore(&mysqlstore.StoreOption{
		MySQLConfig: &mysql.Config{
			Addr:   fmt.Sprintf("%s:%d", cfg.Host, cfg.Port),
			User:   cfg.User,
			Passwd: cfg.Password,
			DBName: cfg.Database,
		},
		MigrationSwitch: migrate,
	})
	if err := store.Init(); err != nil {
		return nil, err
	}
	mod.SetStore(store)
	return store, nil
}

func initKeeper(isExecutor bool) (mod.Keeper, error) {
	cfg := config.Global().MysqlConfig
	keeper := mysqlkeeper.NewKeeper(&mysqlkeeper.KeeperOption{
		Key: workflowWorkerKey(isExecutor),
		MySQLConfig: &mysql.Config{
			Addr:   fmt.Sprintf("%s:%d", cfg.Host, cfg.Port),
			User:   cfg.User,
			Passwd: cfg.Password,
			DBName: cfg.Database,
		},
		// Same reasoning as initStore: always migrate so non-executor pods
		// can also read fastflow heartbeats / leases.
		MigrationSwitch: true,
		WatcherFlag:     !isExecutor,
	})
	if err := keeper.Init(); err != nil {
		return nil, err
	}
	mod.SetKeeper(keeper)
	return keeper, nil
}

func initCommander() mod.Commander {
	commander := &mod.DefCommander{}
	mod.SetCommander(commander)
	return commander
}

func workflowWorkerKey(isExecutor bool) string {
	if !isExecutor {
		return fmt.Sprintf("%s-%s", appworkflow.WatcherName, config.GetWorkMode().String())
	}
	if podName := os.Getenv("POD_NAME"); podName != "" {
		return podName
	}
	if hostName, err := os.Hostname(); err == nil && hostName != "" {
		return hostName
	}
	return fmt.Sprintf("serverless-service-%s", config.GetWorkMode().String())
}

func (impl *workflowImpl) Close() {
	mod.SetKeeper(nil)
	mod.SetStore(nil)
	mod.SetCommander(nil)
	impl.store.Close()
	impl.keeper.Close()
}

func (impl *workflowImpl) RegisterDag(dag *entity.Dag, actions []run.Action) error {
	fastflow.RegisterAction(actions)
	oldDag, err := impl.store.GetDag(dag.ID)
	if err != nil {
		if errors.Is(err, data.ErrDataNotFound) {
			return impl.store.CreateDag(dag)
		}
		return err
	}
	if oldDag == nil {
		return impl.store.CreateDag(dag)
	}
	return impl.store.UpdateDag(dag)
}

func (impl *workflowImpl) RunDagWithTags(
	dagID string,
	vars map[string]string,
	tags map[string]string,
) (*entity.DagInstance, error) {
	return impl.commander.RunDagWithTags(dagID, vars, tags)
}

func (impl *workflowImpl) RetryDagIns(dagInsID string, ops ...mod.CommandOptSetter) error {
	return impl.commander.RetryDagIns(dagInsID, ops...)
}

func (impl *workflowImpl) CancelDagIns(dagInsID string, ops ...mod.CommandOptSetter) error {
	return impl.commander.CancelDagIns(dagInsID, ops...)
}

func (impl *workflowImpl) ListDagIns(input *mod.ListDagInstanceInput) ([]*entity.DagInstance, error) {
	return impl.store.ListDagInstance(input)
}

func (impl *workflowImpl) StartExecutor() error {
	if !config.IsWorkflowMode() {
		return nil
	}
	impl.startOnce.Do(func() {
		impl.startErr = fastflow.Init(&fastflow.InitialOption{
			Keeper:            impl.keeper,
			Store:             impl.store,
			ParserWorkersCnt:  parserWorkersCnt,
			ExecutorWorkerCnt: executorWorkersCnt,
		})
	})
	return impl.startErr
}