package dispatcher

import (
	"context"
	"errors"
	"reflect"
	"testing"
	"time"
	"unsafe"

	"github.com/pingcap/failpoint"
	"github.com/stretchr/testify/require"

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

func setUnexportedBool(t *testing.T, field reflect.Value, value bool) {
	t.Helper()
	require.True(t, field.IsValid())
	reflect.NewAt(field.Type(), unsafe.Pointer(field.UnsafeAddr())).Elem().SetBool(value)
}

func setUnexportedValue(t *testing.T, field reflect.Value, value any) {
	t.Helper()
	require.True(t, field.IsValid())
	writable := reflect.NewAt(field.Type(), unsafe.Pointer(field.UnsafeAddr())).Elem()
	candidate := reflect.ValueOf(value)
	require.True(t, candidate.IsValid())
	require.True(t, candidate.Type().AssignableTo(field.Type()))
	writable.Set(candidate)
}

func TestContextRecoveryHigherFrontierCompletesProductionLifecycle(t *testing.T) {
	stubRuntimeBinaries(t, "opencode")
	initTokenForTest(t)
	t.Setenv(profile.ConfigDirEnvKey, t.TempDir())
	require.NoError(t, failpoint.Enable("dispatcher-new-agent-driver", "return(true)"))
	require.NoError(t, failpoint.Enable("dispatcher-new-agent-process", "return(true)"))

	workDir := t.TempDir()
	const agentID = "agent-recovery-production-lifecycle"
	d := mustNewDispatcher("http://server.test")
	sess := &interactiveSession{
		sessionKey: agentID,
		agentID:    agentID,
		agent: agent.Config{
			ID:      agentID,
			Name:    "Recovery lifecycle agent",
			Runtime: string(runtimeconfig.TypeOpenCode),
		},
		runtime: runtimeconfig.Runtime{
			Type:   runtimeconfig.TypeOpenCode,
			Binary: "opencode",
		},
		binding: conversationBinding{
			channelID:       "ops",
			serverAgentID:   agentID,
			targetID:        "#ops:deadbeef",
			targetScopeType: "channel",
		},
		workDir:             workDir,
		ephemeral:           true,
		durableDelivery:     true,
		suppressCheckpoints: true,
	}
	d.sessionManager.Set(agentID, sess)

	var (
		launchCount int
		launched    []*agent.Process
		drivers     []*blockingDispatcherManagedDriver
		launchWakes []agent.WakeInfo
		launchModes []promptMode
		wakeEvents  []agent.WakeStatusEvent
	)
	d.contextRecoveryLaunchObserver = func(mode promptMode, msg types.InboxMessage) {
		launchModes = append(launchModes, mode)
		require.Equal(t, "context_recovery", msg.MessageType)
	}
	SetNewAgentDriverHook(func(runtimeconfig.RuntimeType) (drivercontract.Driver, error) {
		driver := newBlockingDispatcherManagedDriver()
		drivers = append(drivers, driver)
		return driver, nil
	})
	SetNewAgentProcessHook(func(
		ctx context.Context,
		driver drivercontract.Driver,
		cfg drivercontract.LaunchConfig,
		options agent.ProcessOptions,
	) (*agent.Process, error) {
		launchCount++
		if options.Wake != nil {
			launchWakes = append(launchWakes, *options.Wake)
		}
		originalWakeStatus := options.OnWakeStatus
		options.OnWakeStatus = func(event agent.WakeStatusEvent) {
			wakeEvents = append(wakeEvents, event)
			if originalWakeStatus != nil {
				originalWakeStatus(event)
			}
		}
		process, err := agent.NewProcess(ctx, driver, cfg, options)
		if err == nil {
			launched = append(launched, process)
		}
		return process, err
	})
	t.Cleanup(func() {
		d.contextRecoveryLaunchObserver = nil
		SetNewAgentDriverHook(nil)
		SetNewAgentProcessHook(nil)
		require.NoError(t, failpoint.Disable("dispatcher-new-agent-driver"))
		require.NoError(t, failpoint.Disable("dispatcher-new-agent-process"))
		sess.retired.Store(true)
		sess.mu.Lock()
		sess.process = nil
		sess.mu.Unlock()
		for index, process := range launched {
			drivers[index].allowExit()
			process.Stop()
			require.True(t, process.WaitForExit(time.Second))
		}
		d.Close()
	})

	oldWake := agent.WakeInfo{
		MessageID:     "wake-old",
		Target:        "#ops:deadbeef",
		TargetKey:     "ops",
		TargetSeq:     3,
		DeliveryClass: agent.DeliveryClassVisibleReply,
	}
	oldProcess := &agent.Process{AgentID: agentID, SessionID: "old-process"}
	setProcessExitSummary(t, oldProcess, 0, false)
	oldValue := reflect.ValueOf(oldProcess).Elem()
	setUnexportedBool(t, oldValue.FieldByName("contextWindowExceeded"), true)
	setUnexportedBool(t, oldValue.FieldByName("sleeping"), true)
	setUnexportedValue(t, oldValue.FieldByName("contextRecoveryWakes"), []agent.WakeInfo{oldWake})
	sess.mu.Lock()
	sess.process = oldProcess
	sess.mu.Unlock()

	removalFault := errors.New("injected unknown context-recovery marker removal")
	removalCalls := 0
	d.contextRecoveryMarkerRemoval = func(string, string) error {
		removalCalls++
		return removalFault
	}

	// This is the real Process-exit continuation path. It captures the old
	// wake, calls startContextRecoveryProcessWithContextAndReservation, launches
	// and surfaces seq=3, then returns restarted=true with the injected removal
	// error. The handler must not fall through and remove the live replacement.
	d.handleProcessExit(sess, oldProcess)
	require.Equal(t, 1, launchCount)
	require.Equal(t, 1, removalCalls)
	require.Equal(t, []promptMode{promptModeContextRecovery}, launchModes)
	require.Len(t, launchWakes, 1)
	require.Equal(t, int64(3), launchWakes[0].TargetSeq)
	require.Equal(t, "ops", launchWakes[0].TargetKey)
	oldFrontier, ok := launched[0].MessageFreshnessSnapshot("ops")
	require.True(t, ok)
	require.Equal(t, int64(3), oldFrontier.ObservedSeq)
	state, err := loadSessionResumeState(workDir)
	require.NoError(t, err)
	require.Len(t, state.DeferredDeliveries, 1)
	oldMarker, ok := contextRecoveryMarkerTarget(state.DeferredDeliveries[0])
	require.True(t, ok)
	require.True(t, state.DeferredDeliveries[0].ReplayPending)
	require.Equal(t, int64(3), oldMarker.TargetSeq)
	require.Equal(t, contextRecoveryContinuationID(agentID, oldWake), oldMarker.MessageID)

	// Capture a still higher same-target frontier while seq=3's process is live.
	higherWakes := []agent.WakeInfo{
		{MessageID: "wake-five", Target: "#ops:deadbeef", TargetKey: "ops", TargetSeq: 5, DeliveryClass: agent.DeliveryClassVisibleReply},
		{MessageID: "wake-seven", Target: "#ops:deadbeef", TargetKey: "ops", TargetSeq: 7, DeliveryClass: agent.DeliveryClassVisibleReply},
	}
	require.NoError(t, d.enqueueContextRecoveryContinuation(sess, higherWakes))
	state, err = loadSessionResumeState(workDir)
	require.NoError(t, err)
	require.Len(t, state.DeferredDeliveries, 2)
	var higherMarker types.InboxMessage
	for _, entry := range state.DeferredDeliveries {
		marker, markerOK := contextRecoveryMarkerTarget(entry)
		require.True(t, markerOK)
		if marker.TargetSeq == 7 {
			higherMarker = marker
		}
	}
	require.Equal(t, contextRecoveryContinuationID(agentID, higherWakes[1]), higherMarker.MessageID)
	require.Equal(t, int64(7), higherMarker.TargetSeq)
	require.Equal(t, 1, launchCount, "the old launch remains the only launch before its process boundary")

	// The old process boundary releases only seq=3. The production removal seam
	// is disabled now, so seq=7 uses the real durable removal transaction and
	// starts as a fresh context-recovery process rather than generic replay.
	d.contextRecoveryMarkerRemoval = nil
	setProcessExitSummary(t, launched[0], 0, false)
	d.handleProcessExit(sess, launched[0])
	require.Equal(t, 2, launchCount)
	require.Equal(t, []promptMode{promptModeContextRecovery, promptModeContextRecovery}, launchModes)
	require.Len(t, launchWakes, 2)
	require.Equal(t, higherMarker.MessageID, launchWakes[1].MessageID)
	require.Equal(t, int64(7), launchWakes[1].TargetSeq)
	newFrontier, ok := launched[1].MessageFreshnessSnapshot("ops")
	require.True(t, ok)
	require.Equal(t, int64(7), newFrontier.ObservedSeq, "the higher marker was surfaced by the new process")

	state, err = loadSessionResumeState(workDir)
	require.NoError(t, err)
	for _, entry := range state.DeferredDeliveries {
		marker, markerOK := contextRecoveryMarkerTarget(entry)
		if markerOK {
			t.Fatalf("recovery marker survived successful second launch: %+v", marker)
		}
	}

	// A lower server frontier and a held response cannot settle seq=7. The
	// exact higher SeenUpToSeq then settles it through the real Process API.
	require.False(t, launched[1].ConfirmVisibleReplyForTarget("ops", 5, "reply-low"))
	newFrontier, ok = launched[1].MessageFreshnessSnapshot("ops")
	require.True(t, ok)
	require.Equal(t, int64(7), newFrontier.ObservedSeq)
	launched[1].RememberFreshnessHeld("ops", 5, 7)
	require.False(t, launched[1].ConfirmVisibleReplyForTarget("ops", 5, "reply-held"))
	require.True(t, launched[1].ConfirmVisibleReplyForTarget("ops", 7, "reply-seven"))
	require.False(t, launched[1].ConfirmVisibleReplyForTarget("ops", 7, "reply-duplicate"))
	require.Len(t, wakeEvents, 3)
	require.Equal(t, oldMarker.MessageID, wakeEvents[0].MessageID)
	require.Equal(t, higherMarker.MessageID, wakeEvents[1].MessageID)
	require.Equal(t, higherMarker.MessageID, wakeEvents[2].MessageID)
	require.Equal(t, consts.WakeStatusReplySent, wakeEvents[2].Status)
}