package agent

import (
	"context"
	"errors"
	"sync/atomic"
	"testing"

	"github.com/stretchr/testify/require"

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

type resumeSafetyMutationDriver struct {
	deliveryTestDriver
	compactionCalls atomic.Int32
	restoreCalls    atomic.Int32
	deliverCalls    atomic.Int32
	interruptCalls  atomic.Int32
}

func (d *resumeSafetyMutationDriver) EncodeStdinMessage(text, _ string) []byte {
	return []byte(text + "\n")
}

func (d *resumeSafetyMutationDriver) StartIdleCompaction(context.Context, string) (bool, error) {
	d.compactionCalls.Add(1)
	return true, nil
}

func (d *resumeSafetyMutationDriver) PrepareRecoveryPoint() (string, string, bool, error) {
	return rewindcaps.StrategyCodexAppServerTurns, "recovery-ref", true, nil
}

func (d *resumeSafetyMutationDriver) RestoreRecoveryPoint(context.Context, string, string) error {
	d.restoreCalls.Add(1)
	return nil
}

func (d *resumeSafetyMutationDriver) SupportsAckedBusyRuntimeInputDelivery() bool { return true }

func (d *resumeSafetyMutationDriver) DeliverAckedBusyRuntimeInput(context.Context, drivercontract.RuntimeInput, string) error {
	d.deliverCalls.Add(1)
	return nil
}

func (d *resumeSafetyMutationDriver) DeliverRuntimeInput(context.Context, drivercontract.RuntimeInput, string) error {
	d.deliverCalls.Add(1)
	return nil
}

func (d *resumeSafetyMutationDriver) InterruptActiveRun(context.Context, string) error {
	d.interruptCalls.Add(1)
	return nil
}

type resumeSafetyCountingWriter struct{ writes atomic.Int32 }

func (w *resumeSafetyCountingWriter) Write(p []byte) (int, error) {
	w.writes.Add(1)
	return len(p), nil
}

func (w *resumeSafetyCountingWriter) Close() error { return nil }

func armFailingUnsafeResumeSafetyPersistence(t *testing.T, ap *Process) {
	t.Helper()
	failUnsafe := false
	require.NoError(t, ap.SetOnResumeSafety(func(_ uint64, safe bool, _ string) error {
		if !safe && failUnsafe {
			return errors.New("injected resume-state fsync failure")
		}
		return nil
	}))
	require.NoError(t, ap.markResumeSafety(true), "establish the previously durable safe boundary")
	failUnsafe = true
}

func TestResumeSafetyPersistenceFailureBlocksIdleCompactionBeforeDriverMutation(t *testing.T) {
	driver := &resumeSafetyMutationDriver{}
	ap := &Process{
		AgentID:               "agent-compaction-preflight",
		SessionID:             "thread-1",
		processRuntimeBinding: processRuntimeBinding{driver: driver},
		processLifecycleState: processLifecycleState{state: processStateRunning, isIdle: true},
	}
	armFailingUnsafeResumeSafetyPersistence(t, ap)

	require.False(t, ap.startIdleCompactionAfterTurnEnd())
	require.Zero(t, driver.compactionCalls.Load())
	require.True(t, ap.IsIdle())
}

func TestResumeSafetyPersistenceFailureBlocksLiveRestoreBeforeDriverMutation(t *testing.T) {
	driver := &resumeSafetyMutationDriver{}
	ap := &Process{
		AgentID:               "agent-restore-preflight",
		SessionID:             "thread-1",
		processRuntimeBinding: processRuntimeBinding{driver: driver},
		processLifecycleState: processLifecycleState{state: processStateRunning, isIdle: true},
	}
	armFailingUnsafeResumeSafetyPersistence(t, ap)

	err := ap.RestoreRecoveryPoint(context.Background(), rewindcaps.StrategyCodexAppServerTurns, "recovery-ref")

	require.ErrorContains(t, err, "resume-unsafe")
	require.Zero(t, driver.restoreCalls.Load())
	require.True(t, ap.IsIdle())
}

func TestResumeSafetyPersistenceFailureBlocksTransientRetryRestoreAndRedelivery(t *testing.T) {
	state := transientRuntimeRetryState{
		MessageID:   "msg-1",
		SessionID:   "thread-1",
		Strategy:    rewindcaps.StrategyCodexAppServerTurns,
		RecoveryRef: "recovery-ref",
		Input:       testRuntimeInputWithConfig(drivercontract.RuntimeInput{ID: "msg-1"}, "retry body"),
	}
	for _, tc := range []struct {
		name   string
		invoke func(*Process) error
	}{
		{name: "redelivery", invoke: func(ap *Process) error {
			return ap.redeliverForTransientRuntimeRetry(context.Background(), state)
		}},
		{name: "restore then redelivery", invoke: func(ap *Process) error {
			return ap.restoreLiveForTransientRuntimeRetry(context.Background(), state)
		}},
	} {
		t.Run(tc.name, func(t *testing.T) {
			driver := &resumeSafetyMutationDriver{}
			writer := &resumeSafetyCountingWriter{}
			ap := &Process{
				AgentID:   "agent-retry-preflight",
				SessionID: "thread-1",
				processRuntimeBinding: processRuntimeBinding{
					driver:    driver,
					stdinPipe: writer,
				},
				processLifecycleState: processLifecycleState{state: processStateRunning, isIdle: true},
			}
			armFailingUnsafeResumeSafetyPersistence(t, ap)

			err := tc.invoke(ap)

			require.ErrorContains(t, err, "resume-unsafe")
			require.Zero(t, driver.restoreCalls.Load())
			require.Zero(t, driver.deliverCalls.Load())
			require.Zero(t, writer.writes.Load())
		})
	}
}

func TestResumeSafetyPersistenceFailureBlocksFinalVisibleReplyCheckWrite(t *testing.T) {
	driver := &resumeSafetyMutationDriver{}
	writer := &resumeSafetyCountingWriter{}
	ap := &Process{
		AgentID:   "agent-final-check-preflight",
		SessionID: "thread-1",
		processRuntimeBinding: processRuntimeBinding{
			driver:    driver,
			stdinPipe: writer,
		},
		processLifecycleState: processLifecycleState{state: processStateRunning, isIdle: true},
		processWakeState: processWakeState{
			pendingWake: &pendingDelivery{Wake: WakeInfo{MessageID: "msg-1", Target: "#ops"}, Status: pendingDeliveryAwaitingCheck},
		},
	}
	armFailingUnsafeResumeSafetyPersistence(t, ap)

	require.False(t, ap.injectFinalVisibleReplyCheck())
	require.Zero(t, writer.writes.Load())
	ap.mu.Lock()
	ap.mu.Unlock()
}

func TestSafeResumeSafetyPersistenceFailureRemainsFailClosedWithoutBreakingProcess(t *testing.T) {
	driver := &resumeSafetyMutationDriver{}
	ap := &Process{
		AgentID:               "agent-safe-persist",
		SessionID:             "thread-1",
		processRuntimeBinding: processRuntimeBinding{driver: driver},
	}
	var durableTransitions []bool
	require.NoError(t, ap.SetOnResumeSafety(func(_ uint64, safe bool, _ string) error {
		if safe {
			return errors.New("injected safe promotion fsync failure")
		}
		durableTransitions = append(durableTransitions, false)
		return nil
	}))

	require.NoError(t, ap.markResumeSafety(false))
	require.NoError(t, ap.markResumeSafety(true), "safe promotion failure must not fail runtime work")
	require.Equal(t, []bool{false}, durableTransitions, "durable state remains at the unsafe boundary")
	ap.mu.Lock()
	require.True(t, ap.resumeSafety, "local quiescence may still be tracked for a later retry")
	ap.mu.Unlock()
}