package agent

import (
	"context"
	"sync"
	"testing"
	"time"

	"github.com/stretchr/testify/require"

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

type admissionBlocker struct {
	started chan struct{}
	once    sync.Once
}

func newAdmissionBlocker() admissionBlocker {
	return admissionBlocker{started: make(chan struct{})}
}

func (b *admissionBlocker) block(ctx context.Context) error {
	b.once.Do(func() { close(b.started) })
	<-ctx.Done()
	return ctx.Err()
}

type blockingAckedAdmissionDriver struct {
	busyDeliveryStubDriver
	admissionBlocker
}

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

func (d *blockingAckedAdmissionDriver) DeliverAckedBusyRuntimeInput(ctx context.Context, _ drivercontract.RuntimeInput, _ string) error {
	return d.block(ctx)
}

type blockingProtocolAdmissionDriver struct {
	busyDeliveryStubDriver
	admissionBlocker
}

func (d *blockingProtocolAdmissionDriver) SupportsBusyStdinDelivery() bool { return true }

func (d *blockingProtocolAdmissionDriver) DeliverBusyStdinInput(ctx context.Context, _ drivercontract.RuntimeInput, _ string) error {
	return d.block(ctx)
}

type blockingInterruptAdmissionDriver struct {
	busyDeliveryStubDriver
	admissionBlocker
}

func (d *blockingInterruptAdmissionDriver) InterruptActiveRun(ctx context.Context, _ string) error {
	return d.block(ctx)
}

type blockingRecoveryAdmissionDriver struct {
	busyDeliveryStubDriver
	admissionBlocker
}

func (d *blockingRecoveryAdmissionDriver) Name() string { return "codex" }

func (d *blockingRecoveryAdmissionDriver) PrepareRecoveryPoint() (string, string, bool, error) {
	return "codex_app_server_turns", "recovery-ref", true, nil
}

func (d *blockingRecoveryAdmissionDriver) RestoreRecoveryPoint(ctx context.Context, _, _ string) error {
	return d.block(ctx)
}

func newAdmissionTestProcess(driver drivercontract.Driver, idle bool) *Process {
	return &Process{
		AgentID:   "agent-admission-lifetime",
		SessionID: "session-1",
		processRuntimeBinding: processRuntimeBinding{
			driver: driver,
			cancel: func() {},
		},
		processLifecycleState: processLifecycleState{
			state:  processStateRunning,
			isIdle: idle,
		},
		processRecoveryState: processRecoveryState{disableRecovery: true},
	}
}

func requireStopCancelsAdmission(t *testing.T, ap *Process, started <-chan struct{}, requireCallError bool, invoke func() error) {
	t.Helper()
	callDone := make(chan error, 1)
	go func() { callDone <- invoke() }()
	select {
	case <-started:
	case err := <-callDone:
		t.Fatalf("runtime mutation returned before entering the blocking driver: %v", err)
	case <-time.After(time.Second):
		t.Fatal("runtime mutation did not start")
	}

	stopDone := make(chan struct{})
	go func() {
		ap.Stop()
		close(stopDone)
	}()
	select {
	case <-stopDone:
	case <-time.After(time.Second):
		t.Fatal("Stop did not cancel and drain the in-flight runtime mutation")
	}
	select {
	case err := <-callDone:
		if requireCallError {
			require.Error(t, err)
		}
	case <-time.After(time.Second):
		t.Fatal("runtime mutation did not observe Process admission cancellation")
	}
	require.False(t, ap.AcceptingDeliveries())
}

func TestStopCancelsBackgroundProtocolDeliveryAdmission(t *testing.T) {
	driver := &blockingProtocolAdmissionDriver{admissionBlocker: newAdmissionBlocker()}
	ap := newAdmissionTestProcess(driver, false)
	requireStopCancelsAdmission(t, ap, driver.started, false, func() error {
		return ap.DeliverForHandoff(context.Background(), types.InboxMessage{
			MessageID: "message-1",
			Target:    "#ops:thread",
			Content:   "deliver",
		})
	})
}

func TestStopCancelsBackgroundInterruptAdmission(t *testing.T) {
	driver := &blockingInterruptAdmissionDriver{admissionBlocker: newAdmissionBlocker()}
	ap := newAdmissionTestProcess(driver, false)
	requireStopCancelsAdmission(t, ap, driver.started, true, func() error {
		_, err := ap.InterruptActiveRun(context.Background())
		return err
	})
}

func TestStopCancelsBackgroundRecoveryAdmission(t *testing.T) {
	driver := &blockingRecoveryAdmissionDriver{admissionBlocker: newAdmissionBlocker()}
	ap := newAdmissionTestProcess(driver, true)
	requireStopCancelsAdmission(t, ap, driver.started, true, func() error {
		return ap.RestoreRecoveryPoint(context.Background(), "codex_app_server_turns", "recovery-ref")
	})
}

func TestConcurrentStopsWaitForSameDeliveryGateDrain(t *testing.T) {
	ap := newAdmissionTestProcess(&busyDeliveryStubDriver{}, false)
	require.NoError(t, ap.acquireDeliveryGateForDrain(context.Background()))

	stopOne := make(chan struct{})
	stopTwo := make(chan struct{})
	go func() { ap.Stop(); close(stopOne) }()
	require.Eventually(t, func() bool { return !ap.AcceptingDeliveries() }, time.Second, time.Millisecond)
	go func() { ap.Stop(); close(stopTwo) }()

	for name, done := range map[string]<-chan struct{}{"owner": stopOne, "concurrent": stopTwo} {
		select {
		case <-done:
			t.Fatalf("%s Stop returned before the pre-existing gate lease drained", name)
		case <-time.After(20 * time.Millisecond):
		}
	}
	ap.releaseDeliveryGate()
	for name, done := range map[string]<-chan struct{}{"owner": stopOne, "concurrent": stopTwo} {
		select {
		case <-done:
		case <-time.After(time.Second):
			t.Fatalf("%s Stop did not complete after the gate lease drained", name)
		}
	}
}

func TestStopTriggeredChildExitDoesNotCorruptDeliveryGate(t *testing.T) {
	waitEntered := make(chan struct{})
	waitRelease := make(chan struct{})
	var releaseOnce sync.Once
	ap := newAdmissionTestProcess(&busyDeliveryStubDriver{}, false)
	ap.exited = make(chan struct{})
	ap.waitFunc = func() error {
		close(waitEntered)
		<-waitRelease
		return nil
	}
	ap.terminateFunc = func() { releaseOnce.Do(func() { close(waitRelease) }) }
	go ap.waitProcess()
	<-waitEntered

	stopDone := make(chan struct{})
	go func() { ap.Stop(); close(stopDone) }()
	select {
	case <-stopDone:
	case <-time.After(time.Second):
		t.Fatal("Stop deadlocked with waitProcess delivery-gate drain")
	}
	require.True(t, ap.WaitForExit(time.Second), "waitProcess must release the raw gate and signal exit after Stop cancellation")

	// The gate remains a valid binary semaphore: a canceled admission lifetime
	// must not make waitProcess release a token it never acquired.
	ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond)
	defer cancel()
	require.NoError(t, ap.acquireDeliveryGateForDrain(ctx))
	ap.releaseDeliveryGate()
}