package dispatcher

import (
	"context"
	"testing"

	"github.com/stretchr/testify/require"

	"github.com/tidbcloud/AgentTeam/desktop/internal/agent"
	"github.com/tidbcloud/AgentTeam/desktop/internal/messagesend"
)

// bearerContext builds a context carrying the given raw bearer's identity,
// exactly as the MCP access-control middleware does.
func bearerContext(token string) context.Context {
	return messagesend.WithBearerID(context.Background(), token)
}

func TestSendScopeFailsClosedWithoutBearer(t *testing.T) {
	d, err := New("", Dependencies{})
	require.NoError(t, err)
	_, ok := d.SendScope(context.Background(), "agent-a")
	require.False(t, ok, "no caller identity must fail closed")
}

func TestSendScopeFailsClosedForUnknownAgent(t *testing.T) {
	d, err := New("", Dependencies{})
	require.NoError(t, err)
	_, ok := d.SendScope(bearerContext("token-a"), "agent-unknown")
	require.False(t, ok)
}

func TestSendScopeServesPendingLaunchWindowThenPublished(t *testing.T) {
	d, err := New("", Dependencies{})
	require.NoError(t, err)
	process := &agent.Process{}
	identity := messagesend.TokenIdentity("token-fresh")

	// Pending launch window: the registry binds the fresh bearer to its
	// in-flight Process via the start reservation.
	res, _, ok := d.sessionManager.BeginStart("agent-pending", nil)
	require.True(t, ok)
	d.bindPendingScope(identity, process, res)

	// A non-live process must fail the pending window closed regardless of
	// reservation currency.
	process.Stop()

	_, pendingOK := d.SendScope(bearerContext("token-fresh"), "agent-pending")
	require.False(t, pendingOK,
		"pending window must fail closed for a stopped process")

	// Publish: the session takes over with the same token identity.
	d.sessionManager.FinishStart(res)
	sess := &interactiveSession{
		sessionKey: "agent-pending",
		agentID:    "agent-pending",
		mcpToken:   "token-fresh",
	}
	d.sessionManager.Set("agent-pending", sess)
	sess.mu.Lock()
	sess.process = process
	sess.mu.Unlock()

	// A stale bearer from a previous generation must fail closed.
	_, staleOK := d.SendScope(bearerContext("token-old"), "agent-pending")
	require.False(t, staleOK, "stale bearer must not read the published scope")
}

func TestSendScopePendingBindingDiesWithReservation(t *testing.T) {
	d, err := New("", Dependencies{})
	require.NoError(t, err)
	process := &agent.Process{}
	identity := messagesend.TokenIdentity("token-invalidate")

	res, _, ok := d.sessionManager.BeginStart("agent-invalidate", nil)
	require.True(t, ok)
	d.bindPendingScope(identity, process, res)

	d.sessionManager.Delete("agent-invalidate")
	_, ok = d.SendScope(bearerContext("token-invalidate"), "agent-invalidate")
	require.False(t, ok,
		"invalidated reservation must fail the pending binding closed even before process liveness is considered")
}

func TestSendScopeServesPendingAndPublishedHappyPaths(t *testing.T) {
	d, err := New("", Dependencies{})
	require.NoError(t, err)

	// Build a live-like Process with a scope target via the package-internal
	// constructor path: use the agent package's running state via a stub is
	// not accessible here, so exercise the published path with a real agent
	// Process created through the exported API surface.
	sess := &interactiveSession{sessionKey: "agent-happy", agentID: "agent-happy", mcpToken: "token-happy"}
	d.sessionManager.Set("agent-happy", sess)

	// Stale bearer first (negative control).
	_, staleOK := d.SendScope(bearerContext("token-other"), "agent-happy")
	require.False(t, staleOK)

	// Correct bearer: no process published yet -> fail closed.
	_, noProcess := d.SendScope(bearerContext("token-happy"), "agent-happy")
	require.False(t, noProcess)
}

func TestSendScopePublishedServesLiveProcessScope(t *testing.T) {
	d, err := New("", Dependencies{})
	require.NoError(t, err)

	process := NewDispatcherScopeFixture(t)
	sess := &interactiveSession{sessionKey: "agent-live", agentID: "agent-live", mcpToken: "token-live"}
	d.sessionManager.Set("agent-live", sess)
	sess.mu.Lock()
	sess.process = process
	sess.mu.Unlock()

	scope, ok := d.SendScope(bearerContext("token-live"), "agent-live")
	require.True(t, ok, "correct bearer with a live process must read the scope")
	require.True(t, scope.ContainsTarget("#fixture-turn"))

	// Stop after publish: fail closed.
	process.Stop()
	_, stopped := d.SendScope(bearerContext("token-live"), "agent-live")
	require.False(t, stopped)
}

func TestSendScopePendingWindowServesLiveProcessScope(t *testing.T) {
	d, err := New("", Dependencies{})
	require.NoError(t, err)

	process := NewDispatcherScopeFixture(t)
	identity := messagesend.TokenIdentity("token-pending-live")

	res, _, ok := d.sessionManager.BeginStart("agent-pending-live", nil)
	require.True(t, ok)
	d.bindPendingScope(identity, process, res)

	// Positive read through the real pending path: live process + current
	// reservation + correct bearer.
	scope, served := d.SendScope(bearerContext("token-pending-live"), "agent-pending-live")
	require.True(t, served, "pending window must serve a live process under a current reservation")
	require.True(t, scope.ContainsTarget("#fixture-turn"))

	// A different bearer fails closed even while the binding is live.
	_, wrongBearer := d.SendScope(bearerContext("token-someone-else"), "agent-pending-live")
	require.False(t, wrongBearer)
	d.sessionManager.FinishStart(res)
}

func TestSendScopeSnapshotStopBarrier(t *testing.T) {
	d, err := New("", Dependencies{})
	require.NoError(t, err)

	process := NewDispatcherScopeFixture(t)
	sess := &interactiveSession{sessionKey: "agent-barrier", agentID: "agent-barrier", mcpToken: "token-barrier"}
	d.sessionManager.Set("agent-barrier", sess)
	sess.mu.Lock()
	sess.process = process
	sess.mu.Unlock()

	// Take the scope and generation snapshot the way the guard does, then
	// Stop the process BEFORE the final validation: the sandwich must fail.
	scope, live, generation := process.SendScopeWithGeneration()
	require.True(t, live)
	require.True(t, scope.ContainsTarget("#fixture-turn"))

	process.Stop()
	require.False(t, process.ValidateSendScope(generation),
		"post-snapshot Stop must fail the final validation")
	_, stopped := d.SendScope(bearerContext("token-barrier"), "agent-barrier")
	require.False(t, stopped, "dispatcher read after Stop must fail closed")
}

func TestSendScopeReplacementBarrier(t *testing.T) {
	d, err := New("", Dependencies{})
	require.NoError(t, err)

	process := NewDispatcherScopeFixture(t)
	sess := &interactiveSession{sessionKey: "agent-replace", agentID: "agent-replace", mcpToken: "token-old"}
	d.sessionManager.Set("agent-replace", sess)
	sess.mu.Lock()
	sess.process = process
	sess.mu.Unlock()

	_, ok := d.SendScope(bearerContext("token-old"), "agent-replace")
	require.True(t, ok, "pre-replacement read succeeds")

	// Replacement: the session is superseded with a new token + process.
	replacement := NewDispatcherScopeFixture(t)
	newSess := &interactiveSession{sessionKey: "agent-replace", agentID: "agent-replace", mcpToken: "token-new"}
	d.sessionManager.Set("agent-replace", newSess)
	newSess.mu.Lock()
	newSess.process = replacement
	newSess.mu.Unlock()

	_, stale := d.SendScope(bearerContext("token-old"), "agent-replace")
	require.False(t, stale, "old bearer must fail closed after replacement")
	_, fresh := d.SendScope(bearerContext("token-new"), "agent-replace")
	require.True(t, fresh, "new bearer reads the replacement scope")
}