package manager

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

	"github.com/stretchr/testify/require"

	"github.com/tidbcloud/AgentTeam/desktop/internal/profile"
	"github.com/tidbcloud/AgentTeam/desktop/internal/types"
)

type lifecycleTestConnection struct {
	id        string
	startFunc func(context.Context) error
	stopFunc  func(context.Context) error
	starts    atomic.Int32
	stops     atomic.Int32
}

func (c *lifecycleTestConnection) ID() string    { return c.id }
func (c *lifecycleTestConnection) Enabled() bool { return true }
func (c *lifecycleTestConnection) Start(ctx context.Context) error {
	c.starts.Add(1)
	if c.startFunc != nil {
		return c.startFunc(ctx)
	}
	return nil
}
func (c *lifecycleTestConnection) Stop(ctx context.Context) error {
	c.stops.Add(1)
	if c.stopFunc != nil {
		return c.stopFunc(ctx)
	}
	return nil
}
func (c *lifecycleTestConnection) Send(context.Context, types.Metadata, types.OutboundMessage) error {
	return nil
}

func TestManagerShutdownFencesAdmittedStartAndCompensatesRuntime(t *testing.T) {
	startEntered := make(chan struct{})
	releaseStart := make(chan struct{})
	var startEnteredOnce sync.Once
	connection := &lifecycleTestConnection{
		id: "late-start",
		startFunc: func(context.Context) error {
			startEnteredOnce.Do(func() { close(startEntered) })
			<-releaseStart // Deliberately ignore cancellation at the publication race.
			return nil
		},
	}
	constructorCalls := atomic.Int32{}
	manager := New(nil, func(profile.AgentTeamServer, types.ServerConnectionContext) (types.ServerConnection, error) {
		constructorCalls.Add(1)
		return connection, nil
	})
	cfg := profile.AgentTeamServer{ID: connection.id, Enabled: true}

	startDone := make(chan error, 1)
	go func() { startDone <- manager.Reconcile(nil, &cfg) }()
	select {
	case <-startEntered:
	case <-time.After(time.Second):
		t.Fatal("admitted start did not reach the blocked runtime Start")
	}

	shutdown := manager.BeginShutdown()
	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancel()
	shutdown.StartDrain(ctx)

	// New work is rejected before construction once the terminal fence is shut.
	require.ErrorIs(t, manager.Reconcile(nil, &cfg), ErrManagerClosed)
	require.Equal(t, int32(1), constructorCalls.Load())

	close(releaseStart)
	select {
	case err := <-startDone:
		require.ErrorIs(t, err, ErrManagerClosed)
	case <-time.After(time.Second):
		t.Fatal("superseded start did not finish")
	}
	require.NoError(t, shutdown.Wait(ctx))
	require.Equal(t, int32(1), connection.starts.Load())
	require.Equal(t, int32(1), connection.stops.Load(), "publication loser must synchronously compensating-stop")
	require.Empty(t, manager.ServerConnections())

	lateRegister := &lifecycleTestConnection{id: "late-register"}
	require.ErrorIs(t, manager.TryRegister(lateRegister), ErrManagerClosed)
	require.Equal(t, int32(1), lateRegister.stops.Load())
	require.Empty(t, manager.ServerConnections())
}

func TestManagerSameIDNewerStartSupersedesAndStopsOlderRuntime(t *testing.T) {
	firstEntered := make(chan struct{})
	secondEntered := make(chan struct{})
	releaseFirst := make(chan struct{})
	releaseSecond := make(chan struct{})
	first := &lifecycleTestConnection{
		id: "same-id",
		startFunc: func(context.Context) error {
			close(firstEntered)
			<-releaseFirst
			return nil
		},
	}
	second := &lifecycleTestConnection{
		id: "same-id",
		startFunc: func(context.Context) error {
			close(secondEntered)
			<-releaseSecond
			return nil
		},
	}
	manager := New(nil, func(server profile.AgentTeamServer, _ types.ServerConnectionContext) (types.ServerConnection, error) {
		if server.Name == "first" {
			return first, nil
		}
		return second, nil
	})
	firstCfg := profile.AgentTeamServer{ID: "same-id", Name: "first", Enabled: true}
	secondCfg := profile.AgentTeamServer{ID: "same-id", Name: "second", Enabled: true}

	firstDone := make(chan error, 1)
	go func() { firstDone <- manager.Reconcile(nil, &firstCfg) }()
	<-firstEntered
	secondDone := make(chan error, 1)
	go func() { secondDone <- manager.Reconcile(nil, &secondCfg) }()
	<-secondEntered

	// The newer mutation publishes first. The older start must not overwrite it
	// when it eventually completes.
	close(releaseSecond)
	require.NoError(t, <-secondDone)
	close(releaseFirst)
	require.ErrorIs(t, <-firstDone, ErrServerConnectionStartSuperseded)

	published, ok := manager.ServerConnection("same-id")
	require.True(t, ok)
	require.Same(t, second, published)
	require.Equal(t, int32(1), first.stops.Load())
	require.Equal(t, int32(0), second.stops.Load())

	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
	defer cancel()
	manager.Stop(ctx)
	require.Equal(t, int32(1), second.stops.Load())
}

func TestManagerReconcileSupersededFailureCannotRollbackOverNewerRuntime(t *testing.T) {
	const id = "reconcile-rollback-generation"
	failedConstructionEntered := make(chan struct{})
	releaseFailedConstruction := make(chan struct{})
	oldRuntime := &lifecycleTestConnection{id: id}
	rollbackRuntime := &lifecycleTestConnection{id: id}
	winnerRuntime := &lifecycleTestConnection{id: id}

	manager := New(nil, func(server profile.AgentTeamServer, _ types.ServerConnectionContext) (types.ServerConnection, error) {
		switch server.ServerURL {
		case "https://candidate-b.test":
			close(failedConstructionEntered)
			<-releaseFailedConstruction
			return nil, errors.New("candidate B construction failed")
		case "https://winner-c.test":
			return winnerRuntime, nil
		case "https://previous-a.test":
			return rollbackRuntime, nil
		default:
			return nil, errors.New("unexpected server config")
		}
	})
	require.NoError(t, manager.TryRegister(oldRuntime))

	previousA := profile.AgentTeamServer{ID: id, Enabled: true, ServerURL: "https://previous-a.test"}
	candidateB := profile.AgentTeamServer{ID: id, Enabled: true, ServerURL: "https://candidate-b.test"}
	winnerC := profile.AgentTeamServer{ID: id, Enabled: true, ServerURL: "https://winner-c.test"}

	olderDone := make(chan error, 1)
	go func() { olderDone <- manager.Reconcile(&previousA, &candidateB) }()
	select {
	case <-failedConstructionEntered:
	case <-time.After(time.Second):
		t.Fatal("older reconcile did not reach blocked candidate construction")
	}

	// The newer mutation publishes C while B is still admitted. Releasing B
	// makes the older operation observe supersession. It must not allocate a new
	// generation for rollback A and overwrite C.
	require.NoError(t, manager.Reconcile(&candidateB, &winnerC))
	published, ok := manager.ServerConnection(id)
	require.True(t, ok)
	require.Same(t, winnerRuntime, published)

	close(releaseFailedConstruction)
	select {
	case err := <-olderDone:
		require.ErrorIs(t, err, ErrServerConnectionStartSuperseded)
	case <-time.After(time.Second):
		t.Fatal("older reconcile did not finish after candidate failure")
	}

	published, ok = manager.ServerConnection(id)
	require.True(t, ok)
	require.Same(t, winnerRuntime, published, "stale rollback overwrote the newer runtime")
	require.Zero(t, rollbackRuntime.starts.Load(), "stale rollback runtime was started")
	require.Zero(t, winnerRuntime.stops.Load(), "newer runtime was stopped by stale rollback")

	manager.Stop(context.Background())
}

func TestManagerReconcileBlockedOldStopCannotPublishAfterNewerWinner(t *testing.T) {
	const id = "reconcile-blocked-stop-generation"
	oldStopEntered := make(chan struct{})
	releaseOldStop := make(chan struct{})
	oldRuntime := &lifecycleTestConnection{id: id}
	oldRuntime.stopFunc = func(context.Context) error {
		switch oldRuntime.stops.Load() {
		case 1:
			close(oldStopEntered)
			<-releaseOldStop
			return nil
		default:
			// The newer reconcile may also stop the same published runtime. Let
			// that request complete first so the older request resumes after C is
			// already the authoritative runtime winner.
			return nil
		}
	}
	candidateBRuntime := &lifecycleTestConnection{id: id}
	winnerCRuntime := &lifecycleTestConnection{id: id}
	manager := New(nil, func(server profile.AgentTeamServer, _ types.ServerConnectionContext) (types.ServerConnection, error) {
		switch server.ServerURL {
		case "https://candidate-b.test":
			return candidateBRuntime, nil
		case "https://winner-c.test":
			return winnerCRuntime, nil
		default:
			return nil, errors.New("unexpected server config")
		}
	})
	require.NoError(t, manager.TryRegister(oldRuntime))

	previousA := profile.AgentTeamServer{ID: id, Enabled: true, ServerURL: "https://previous-a.test"}
	candidateB := profile.AgentTeamServer{ID: id, Enabled: true, ServerURL: "https://candidate-b.test"}
	winnerC := profile.AgentTeamServer{ID: id, Enabled: true, ServerURL: "https://winner-c.test"}

	olderDone := make(chan error, 1)
	go func() { olderDone <- manager.Reconcile(&previousA, &candidateB) }()
	select {
	case <-oldStopEntered:
	case <-time.After(time.Second):
		t.Fatal("older reconcile did not block while stopping A")
	}

	require.NoError(t, manager.Reconcile(&candidateB, &winnerC))
	published, ok := manager.ServerConnection(id)
	require.True(t, ok)
	require.Same(t, winnerCRuntime, published)

	close(releaseOldStop)
	select {
	case err := <-olderDone:
		require.ErrorIs(t, err, ErrServerConnectionStartSuperseded)
	case <-time.After(time.Second):
		t.Fatal("older reconcile did not retire after its stale stop completed")
	}

	published, ok = manager.ServerConnection(id)
	require.True(t, ok)
	require.Same(t, winnerCRuntime, published, "older request published B after newer winner C")
	require.Zero(t, candidateBRuntime.starts.Load(), "stale B runtime was started after its operation was superseded")
	require.Zero(t, winnerCRuntime.stops.Load(), "stale operation stopped the newer winner")

	manager.Stop(context.Background())
}

func TestManagerEquivalentNewerUpdateRepairsRuntimeBehindBlockedOldStop(t *testing.T) {
	const id = "reconcile-equivalent-winner"
	oldStopEntered := make(chan struct{})
	releaseOldStop := make(chan struct{})
	oldRuntime := &lifecycleTestConnection{id: id}
	oldRuntime.stopFunc = func(context.Context) error {
		if oldRuntime.stops.Load() == 1 {
			close(oldStopEntered)
			<-releaseOldStop
		}
		return nil
	}
	staleBRuntime := &lifecycleTestConnection{id: id}
	winnerBRuntime := &lifecycleTestConnection{id: id}
	manager := New(nil, func(server profile.AgentTeamServer, _ types.ServerConnectionContext) (types.ServerConnection, error) {
		switch server.Name {
		case "initial-a":
			return oldRuntime, nil
		case "candidate-b":
			return staleBRuntime, nil
		case "winner-b-name-only":
			return winnerBRuntime, nil
		default:
			return nil, errors.New("unexpected server config")
		}
	})
	previousA := profile.AgentTeamServer{ID: id, Enabled: true, Name: "initial-a", ServerURL: "https://a.test"}
	candidateB := profile.AgentTeamServer{ID: id, Enabled: true, Name: "candidate-b", ServerURL: "https://b.test"}
	winnerB := candidateB
	winnerB.Name = "winner-b-name-only" // Runtime-equivalent profile delta.
	require.NoError(t, manager.Reconcile(nil, &previousA))

	olderDone := make(chan error, 1)
	go func() { olderDone <- manager.Reconcile(&previousA, &candidateB) }()
	select {
	case <-oldStopEntered:
	case <-time.After(time.Second):
		t.Fatal("older reconcile did not block while stopping A")
	}

	// Although candidateB -> winnerB is runtime-equivalent, the currently
	// published runtime still belongs to A. The newer operation must repair B,
	// not trust map presence and return while the older Stop is still in flight.
	require.NoError(t, manager.Reconcile(&candidateB, &winnerB))
	published, ok := manager.ServerConnection(id)
	require.True(t, ok)
	require.Same(t, winnerBRuntime, published)

	close(releaseOldStop)
	require.ErrorIs(t, <-olderDone, ErrServerConnectionStartSuperseded)
	published, ok = manager.ServerConnection(id)
	require.True(t, ok)
	require.Same(t, winnerBRuntime, published)
	require.Zero(t, staleBRuntime.starts.Load())
	require.Zero(t, winnerBRuntime.stops.Load())

	manager.Stop(context.Background())
}

func TestManagerPrepareDeleteCompensationCannotRestoreOverNewerRuntime(t *testing.T) {
	const id = "delete-rollback-generation"
	oldRuntime := &lifecycleTestConnection{id: id}
	rollbackRuntime := &lifecycleTestConnection{id: id}
	winnerRuntime := &lifecycleTestConnection{id: id}
	manager := New(nil, func(server profile.AgentTeamServer, _ types.ServerConnectionContext) (types.ServerConnection, error) {
		switch server.ServerURL {
		case "https://previous-a.test":
			return rollbackRuntime, nil
		case "https://winner-c.test":
			return winnerRuntime, nil
		default:
			return nil, errors.New("unexpected server config")
		}
	})
	require.NoError(t, manager.TryRegister(oldRuntime))

	previousA := profile.AgentTeamServer{ID: id, Enabled: true, ServerURL: "https://previous-a.test"}
	rollback, err := manager.PrepareDelete(&previousA)
	require.NoError(t, err)
	require.Equal(t, int32(1), oldRuntime.stops.Load())
	require.Empty(t, manager.ServerConnections())

	winnerC := profile.AgentTeamServer{ID: id, Enabled: true, ServerURL: "https://winner-c.test"}
	require.NoError(t, manager.Reconcile(nil, &winnerC))
	require.ErrorIs(t, rollback(), ErrServerConnectionStartSuperseded)

	published, ok := manager.ServerConnection(id)
	require.True(t, ok)
	require.Same(t, winnerRuntime, published)
	require.Zero(t, rollbackRuntime.starts.Load(), "stale delete compensation started the old runtime")
	require.Zero(t, winnerRuntime.stops.Load(), "stale delete compensation stopped the newer runtime")

	manager.Stop(context.Background())
}

func TestManagerPrepareDeleteStopsStrayRuntimeForDisabledSnapshot(t *testing.T) {
	const id = "delete-disabled-stray-runtime"
	strayRuntime := &lifecycleTestConnection{id: id}
	manager := New(nil, nil)
	require.NoError(t, manager.TryRegister(strayRuntime))

	rollback, err := manager.PrepareDelete(&profile.AgentTeamServer{ID: id, Enabled: false})
	require.NoError(t, err)
	require.NotNil(t, rollback)
	require.Equal(t, int32(1), strayRuntime.stops.Load(), "disabled profile snapshot still owns cleanup by server ID")
	_, ok := manager.ServerConnection(id)
	require.False(t, ok)

	require.NoError(t, rollback())
	require.Zero(t, strayRuntime.starts.Load(), "disabled delete rollback unexpectedly restarted a runtime")
}

func TestManagerDisableInvalidatesInFlightSameIDStart(t *testing.T) {
	startEntered := make(chan struct{})
	releaseStart := make(chan struct{})
	connection := &lifecycleTestConnection{
		id: "disable-race",
		startFunc: func(context.Context) error {
			close(startEntered)
			<-releaseStart
			return nil
		},
	}
	manager := New(nil, func(profile.AgentTeamServer, types.ServerConnectionContext) (types.ServerConnection, error) {
		return connection, nil
	})
	cfg := profile.AgentTeamServer{ID: connection.id, Enabled: true}
	startDone := make(chan error, 1)
	go func() { startDone <- manager.Reconcile(nil, &cfg) }()
	<-startEntered

	require.NoError(t, manager.Disable(connection.id))
	close(releaseStart)
	require.ErrorIs(t, <-startDone, ErrServerConnectionStartSuperseded)
	require.Equal(t, int32(1), connection.stops.Load())
	_, ok := manager.ServerConnection(connection.id)
	require.False(t, ok)
}

func TestManagerConcurrentStopSharesOneTerminalDrain(t *testing.T) {
	stopEntered := make(chan struct{})
	releaseStop := make(chan struct{})
	var stopEnteredOnce sync.Once
	connection := &lifecycleTestConnection{
		id: "concurrent-stop",
		stopFunc: func(context.Context) error {
			stopEnteredOnce.Do(func() { close(stopEntered) })
			<-releaseStop
			return nil
		},
	}
	manager := New(nil, nil)
	require.NoError(t, manager.TryRegister(connection))

	const callers = 32
	var callersWG sync.WaitGroup
	callersWG.Add(callers)
	for range callers {
		go func() {
			defer callersWG.Done()
			manager.Stop(context.Background())
		}()
	}
	select {
	case <-stopEntered:
	case <-time.After(time.Second):
		t.Fatal("terminal drain did not start")
	}
	require.Equal(t, int32(1), connection.stops.Load())
	require.Empty(t, manager.ServerConnections(), "shutdown must atomically detach published runtimes")

	late := &lifecycleTestConnection{id: "rejected-after-stop"}
	err := manager.TryRegister(late)
	require.True(t, errors.Is(err, ErrManagerClosed))
	require.Equal(t, int32(1), late.stops.Load())

	close(releaseStop)
	done := make(chan struct{})
	go func() {
		callersWG.Wait()
		close(done)
	}()
	select {
	case <-done:
	case <-time.After(time.Second):
		t.Fatal("concurrent Stop callers did not join the shared drain")
	}
	require.Equal(t, int32(1), connection.stops.Load())
}

func TestManagerBoundedShutdownKeepsExactDrainAliveForLaterStrongJoin(t *testing.T) {
	secondStopEntered := make(chan struct{})
	releaseSecondStop := make(chan struct{})
	connection := &lifecycleTestConnection{id: "bounded-then-strong"}
	connection.stopFunc = func(ctx context.Context) error {
		switch connection.stops.Load() {
		case 1:
			<-ctx.Done()
			return ctx.Err()
		case 2:
			close(secondStopEntered)
			<-releaseSecondStop
			return nil
		default:
			return nil
		}
	}
	manager := New(nil, nil)
	require.NoError(t, manager.TryRegister(connection))
	shutdown := manager.BeginShutdown()

	boundedCtx, cancelBounded := context.WithTimeout(context.Background(), 40*time.Millisecond)
	defer cancelBounded()
	require.ErrorIs(t, shutdown.Wait(boundedCtx), context.DeadlineExceeded)
	select {
	case <-secondStopEntered:
	case <-time.After(time.Second):
		t.Fatal("bounded Stop failure did not continue with an exact background join")
	}

	strongDone := make(chan error, 1)
	go func() { strongDone <- shutdown.Wait(context.Background()) }()
	select {
	case err := <-strongDone:
		t.Fatalf("strong waiter crossed the still-open second Stop boundary: %v", err)
	case <-time.After(30 * time.Millisecond):
	}

	close(releaseSecondStop)
	select {
	case err := <-strongDone:
		require.NoError(t, err)
	case <-time.After(time.Second):
		t.Fatal("strong waiter did not join the continued drain")
	}
	require.Equal(t, int32(2), connection.stops.Load(), "one bounded attempt plus one exact strong continuation")
}

func TestManagerTerminalAPIsConsistentlyRejectAfterShutdown(t *testing.T) {
	manager := New(nil, nil)
	manager.BeginShutdown()

	require.ErrorIs(t, manager.Start(context.Background()), ErrManagerClosed)
	require.ErrorIs(t, manager.Enable("missing"), ErrManagerClosed)
	require.ErrorIs(t, manager.Disable("missing"), ErrManagerClosed)
	require.ErrorIs(t, manager.Reconcile(nil, nil), ErrManagerClosed)
}

func TestManagerTryRegisterIsIdentityIdempotentAndRejectsDifferentDuplicate(t *testing.T) {
	manager := New(nil, nil)
	registered := &lifecycleTestConnection{id: "identity"}
	require.NoError(t, manager.TryRegister(registered))
	require.NoError(t, manager.TryRegister(registered))
	require.Equal(t, int32(0), registered.stops.Load())

	duplicate := &lifecycleTestConnection{id: "identity"}
	require.ErrorIs(t, manager.TryRegister(duplicate), ErrServerConnectionAlreadyRegistered)
	require.Equal(t, int32(1), duplicate.stops.Load())
	current, ok := manager.ServerConnection(" identity ")
	require.True(t, ok)
	require.Same(t, registered, current)
	require.Equal(t, int32(0), registered.stops.Load())

	manager.Stop(context.Background())
	require.Equal(t, int32(1), registered.stops.Load())
}

func TestManagerRejectsEmptyServerConnectionID(t *testing.T) {
	manager := New(nil, func(profile.AgentTeamServer, types.ServerConnectionContext) (types.ServerConnection, error) {
		t.Fatal("empty ID must be rejected before construction")
		return nil, nil
	})
	invalid := &lifecycleTestConnection{id: "  "}
	require.ErrorIs(t, manager.TryRegister(invalid), ErrInvalidServerConnectionID)
	require.Equal(t, int32(1), invalid.stops.Load())
	require.ErrorIs(t, manager.Enable("  "), ErrInvalidServerConnectionID)
	require.ErrorIs(t, manager.Reconcile(nil, &profile.AgentTeamServer{ID: "  ", Enabled: true}), ErrInvalidServerConnectionID)
}

func TestManagerPublishedRuntimeOwnsLifetimeContextUntilTerminalMutation(t *testing.T) {
	for _, test := range []struct {
		name string
		stop func(*Manager, string)
	}{
		{
			name: "disable",
			stop: func(manager *Manager, id string) {
				require.NoError(t, manager.Disable(id))
			},
		},
		{
			name: "shutdown",
			stop: func(manager *Manager, _ string) {
				shutdown := manager.BeginShutdown()
				require.NoError(t, shutdown.Wait(context.Background()))
			},
		},
	} {
		t.Run(test.name, func(t *testing.T) {
			const id = "lifetime-" + "context"
			runtimeCtxSeen := make(chan context.Context, 1)
			connection := &lifecycleTestConnection{
				id: id,
				startFunc: func(ctx context.Context) error {
					runtimeCtxSeen <- ctx
					return nil
				},
			}
			manager := New(nil, func(profile.AgentTeamServer, types.ServerConnectionContext) (types.ServerConnection, error) {
				return connection, nil
			})
			cfg := profile.AgentTeamServer{ID: id, Enabled: true}
			require.NoError(t, manager.Reconcile(nil, &cfg))
			runtimeCtx := <-runtimeCtxSeen

			select {
			case <-runtimeCtx.Done():
				t.Fatalf("successful publication returned with an already-canceled runtime context: %v", runtimeCtx.Err())
			default:
			}

			test.stop(manager, id)
			select {
			case <-runtimeCtx.Done():
			case <-time.After(time.Second):
				t.Fatal("terminal mutation did not cancel the published runtime lifetime")
			}
			require.Equal(t, int32(1), connection.stops.Load())
		})
	}
}

func TestManagerPublishedRuntimeDetachesStartupCallerCancellation(t *testing.T) {
	const id = "caller-detach"
	runtimeCtxSeen := make(chan context.Context, 1)
	connection := &lifecycleTestConnection{
		id: id,
		startFunc: func(ctx context.Context) error {
			runtimeCtxSeen <- ctx
			return nil
		},
	}
	manager := New(nil, func(profile.AgentTeamServer, types.ServerConnectionContext) (types.ServerConnection, error) {
		return connection, nil
	})
	callerCtx, cancelCaller := context.WithCancel(context.Background())
	require.NoError(t, manager.startConfiguredServerConnection(callerCtx, profile.AgentTeamServer{ID: id, Enabled: true}))
	runtimeCtx := <-runtimeCtxSeen

	cancelCaller()
	select {
	case <-runtimeCtx.Done():
		t.Fatalf("published runtime remained owned by the completed startup caller: %v", runtimeCtx.Err())
	case <-time.After(30 * time.Millisecond):
	}

	manager.Stop(context.Background())
	select {
	case <-runtimeCtx.Done():
	case <-time.After(time.Second):
		t.Fatal("manager shutdown did not cancel detached runtime lifetime")
	}
}