package dispatcher

import (
	"context"
	"sync"

	"github.com/pingcap/failpoint"

	"github.com/tidbcloud/AgentTeam/desktop/internal/agent"
	"github.com/tidbcloud/AgentTeam/desktop/internal/agent/drivercontract"
	"github.com/tidbcloud/AgentTeam/desktop/internal/agent/runtimeconfig"
)

// Test seams gated behind pingcap/failpoint. Production call sites invoke
// agent.NewDriver / agent.NewProcess directly; tests enable the
// "dispatcher-new-agent-driver" / "dispatcher-new-agent-process" failpoints
// and register fakes via SetNewAgentDriverHook / SetNewAgentProcessHook.
// The call sites check overrideNewAgentDriver / overrideNewAgentProcess
// before running the real constructor so NewProcess side effects (spawning
// a child process) are not triggered during tests.

type newAgentDriverFunc func(runtimeconfig.RuntimeType) (drivercontract.Driver, error)

type newAgentProcessFunc func(
	ctx context.Context,
	driver drivercontract.Driver,
	cfg drivercontract.LaunchConfig,
	options agent.ProcessOptions,
) (*agent.Process, error)

var (
	hookMu              sync.RWMutex
	newAgentDriverHook  newAgentDriverFunc
	newAgentProcessHook newAgentProcessFunc
)

// SetNewAgentDriverHook installs a test hook used in place of
// agent.NewDriver while the "dispatcher-new-agent-driver" failpoint is
// enabled. Passing nil clears the hook.
func SetNewAgentDriverHook(fn newAgentDriverFunc) {
	hookMu.Lock()
	defer hookMu.Unlock()
	newAgentDriverHook = fn
}

// SetNewAgentProcessHook installs a test hook used in place of
// agent.NewProcess while the "dispatcher-new-agent-process" failpoint is
// enabled. Passing nil clears the hook.
func SetNewAgentProcessHook(fn newAgentProcessFunc) {
	hookMu.Lock()
	defer hookMu.Unlock()
	newAgentProcessHook = fn
}

// overrideNewAgentDriver returns a non-nil hook only when the failpoint is
// enabled AND a hook has been registered. Call sites short-circuit to this
// hook instead of invoking agent.NewDriver.
func overrideNewAgentDriver() (newAgentDriverFunc, bool) {
	if _, err := failpoint.Eval("dispatcher-new-agent-driver"); err != nil {
		return nil, false
	}
	hookMu.RLock()
	defer hookMu.RUnlock()
	if newAgentDriverHook == nil {
		return nil, false
	}
	return newAgentDriverHook, true
}

// overrideNewAgentProcess returns a non-nil hook only when the failpoint is
// enabled AND a hook has been registered. Call sites short-circuit to this
// hook instead of invoking agent.NewProcess.
func overrideNewAgentProcess() (newAgentProcessFunc, bool) {
	if _, err := failpoint.Eval("dispatcher-new-agent-process"); err != nil {
		return nil, false
	}
	hookMu.RLock()
	defer hookMu.RUnlock()
	if newAgentProcessHook == nil {
		return nil, false
	}
	return newAgentProcessHook, true
}