package dispatcher

import (
	"context"
	"errors"
	"os"
	"testing"
	"time"

	"github.com/stretchr/testify/require"

	"github.com/tidbcloud/AgentTeam/desktop/internal/agent"
	"github.com/tidbcloud/AgentTeam/desktop/internal/api/exec"
	"github.com/tidbcloud/AgentTeam/desktop/internal/token"
	"github.com/tidbcloud/AgentTeam/desktop/internal/wire/apiclient"
)

func TestCurrentProcessExitRevokesExactExecRuntimeAuthorization(t *testing.T) {
	manager, err := token.NewKeyManager(t.TempDir() + "/config.yaml")
	require.NoError(t, err)
	d := &Dispatcher{execAuthorizations: exec.NewRuntimeAuthorizationStore(manager), sessionManager: NewSessionManager()}
	d.initServices()
	lease, err := d.execAuthorizations.Create(exec.RuntimeAuthorizationConfig{
		ServerID: "server-1", AgentID: "agent-1", SessionID: "session-1", RunID: "runtime-1",
		WorkDir: t.TempDir(), Server: &apiclient.Client{},
		ContextExpires: time.Now().Add(time.Hour), TokenTTL: time.Hour,
	})
	require.NoError(t, err)
	process := newExecAuthorizationTestProcess(t, lease.AuthorizationID)
	session := &interactiveSession{
		sessionKey: "agent-1", agentID: "agent-1", process: process,
		ephemeral: true, durableDelivery: false,
	}
	d.sessionManager.Set("agent-1", session)

	d.handleProcessExit(session, process)

	_, err = manager.ValidateExecToken("agent-1", "server-1", lease.Token)
	require.Error(t, err)
}

func TestStaleProcessExitStillRevokesItsExactExecAuthorization(t *testing.T) {
	t.Setenv("AGENTTEAM_CONFIG_DIR", t.TempDir())
	manager, err := token.NewKeyManager(t.TempDir() + "/config.yaml")
	require.NoError(t, err)
	d := &Dispatcher{
		serverURL: "http://127.0.0.1:3456", execAuthorizations: exec.NewRuntimeAuthorizationStore(manager),
		sessionManager: NewSessionManager(),
	}
	d.initServices()
	cfg, err := d.activateExec(
		context.Background(), execTestLaunchContext("agentteam-main"), "session-1",
		&execResolved{api: &apiclient.Client{AgentID: "server-agent"}},
	)
	require.NoError(t, err)
	tokenValue := readExecCredentialForTest(t, cfg.TokenFile)
	exiting := newExecAuthorizationTestProcess(t, cfg.AuthorizationID)
	replacement := &agent.Process{AgentID: "agent-1", SessionID: "replacement"}
	session := &interactiveSession{
		sessionKey: "agent-1", agentID: "agent-1", process: replacement,
		ephemeral: true, durableDelivery: false,
	}
	d.sessionManager.Set("agent-1", session)

	d.handleProcessExit(session, exiting)

	_, err = manager.ValidateExecToken("agent-1", "agentteam-main", tokenValue)
	require.Error(t, err)
	require.Eventually(t, func() bool {
		_, statErr := os.Stat(cfg.TokenFile)
		return errors.Is(statErr, os.ErrNotExist)
	}, time.Second, time.Millisecond)
	require.Same(t, replacement, session.process, "stale exit must not disturb the current Process")
}

func TestProcessExitRevokesExactExecAuthorizationBeforeSessionLifecycleLock(t *testing.T) {
	d := &Dispatcher{execAuthorizations: &exec.RuntimeAuthorizationStore{}}
	d.initServices()
	revoked := make(chan string, 1)
	d.authorizationRevoke = func(authorizationID string) <-chan struct{} {
		revoked <- authorizationID
		return closedExecDrain()
	}
	exiting := newExecAuthorizationTestProcess(t, "departed-authorization")
	replacement := &agent.Process{AgentID: "agent-1", SessionID: "replacement"}
	session := &interactiveSession{
		sessionKey: "agent-1", agentID: "agent-1", process: replacement,
		ephemeral: true, durableDelivery: false,
	}

	// A session mutation may hold this lock while the Process waiter enters its
	// exit callback. Exact Process-scoped authority must already be revoked,
	// rather than waiting behind that unrelated lifecycle operation.
	session.mu.Lock()
	done := make(chan struct{})
	go func() {
		d.handleProcessExit(session, exiting)
		close(done)
	}()
	select {
	case authorizationID := <-revoked:
		require.Equal(t, "departed-authorization", authorizationID)
	case <-time.After(time.Second):
		session.mu.Unlock()
		t.Fatal("Process exit waited on the session lock before revoking its Exec authorization")
	}
	select {
	case <-done:
		session.mu.Unlock()
		t.Fatal("Process exit unexpectedly crossed the held session lifecycle lock")
	default:
	}
	session.mu.Unlock()
	select {
	case <-done:
	case <-time.After(time.Second):
		t.Fatal("Process exit did not finish after the session lock was released")
	}
	require.Same(t, replacement, session.process, "stale exit must preserve the replacement Process")
}