package manager

import (
	"context"
	"errors"
	"fmt"
	"strings"
	"sync"
	"testing"
	"time"

	"github.com/stretchr/testify/require"
	"go.yaml.in/yaml/v3"

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

type managerTestRecorder struct {
	mu         sync.Mutex
	starts     int
	stops      int
	signatures []string
}

func (r *managerTestRecorder) recordStart(signature string) {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.starts++
	r.signatures = append(r.signatures, signature)
}

func (r *managerTestRecorder) recordStop() {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.stops++
}

func (r *managerTestRecorder) snapshot() (starts int, stops int, signatures []string) {
	r.mu.Lock()
	defer r.mu.Unlock()
	return r.starts, r.stops, append([]string(nil), r.signatures...)
}

type managerTestServer struct {
	id        string
	enabled   bool
	signature string
	mu        sync.Mutex
	identity  profile.AgentTeamServer
	recorder  *managerTestRecorder
	startErr  error
	stopErr   error
	stopFunc  func(context.Context) error
	applyErr  error
	onApply   func()
	onStart   func()
	startCtx  context.Context
	connected chan struct{}
}

func (c *managerTestServer) ID() string    { return c.id }
func (c *managerTestServer) Enabled() bool { return c.enabled }

func (c *managerTestServer) Start(ctx context.Context) error {
	if !c.enabled {
		return nil
	}
	c.startCtx = ctx
	c.recorder.recordStart(c.signature)
	if c.onStart != nil {
		c.onStart()
	}
	return c.startErr
}

func (c *managerTestServer) Stop(ctx context.Context) error {
	c.recorder.recordStop()
	if c.stopFunc != nil {
		return c.stopFunc(ctx)
	}
	return c.stopErr
}

func (c *managerTestServer) Send(context.Context, types.Metadata, types.OutboundMessage) error {
	return nil
}

func (c *managerTestServer) WaitConnected(ctx context.Context) error {
	if c.connected == nil {
		return nil
	}
	select {
	case <-c.connected:
		return nil
	case <-ctx.Done():
		return ctx.Err()
	}
}

func (c *managerTestServer) ServerIdentitySnapshot() profile.AgentTeamServer {
	c.mu.Lock()
	defer c.mu.Unlock()
	return c.identity
}

func (c *managerTestServer) ApplyServerBinding(previous, current profile.AgentTeamServer) error {
	if c.onApply != nil {
		c.onApply()
	}
	if c.applyErr != nil {
		return c.applyErr
	}
	c.mu.Lock()
	defer c.mu.Unlock()
	if strings.TrimSpace(c.identity.WorkspaceID) != strings.TrimSpace(previous.WorkspaceID) {
		return errors.New("stale workspace binding")
	}
	c.identity.WorkspaceID = strings.TrimSpace(current.WorkspaceID)
	c.identity.MachineID = strings.TrimSpace(current.MachineID)
	return nil
}

func TestManagerReconcile_UpdateWithoutRuntimeChangeSkipsRestart(t *testing.T) {
	manager := New(nil, newManagerTestServer)
	initial := newManagerTestProfileServer("ch-stable", "agent-a", map[string]any{"mode": "alpha"})
	updated := initial
	updated.Name = "Updated Name"
	updated.Description = "Updated Description"

	require.NoError(t, manager.Reconcile(nil, &initial))
	require.NoError(t, manager.Reconcile(&initial, &updated))

	recorder := managerTestRecorderFor(initial.ID)
	starts, stops, signatures := recorder.snapshot()
	require.Equal(t, 1, starts)
	require.Equal(t, 0, stops)
	require.Equal(t, []string{managerTestSignature(initial)}, signatures)
}

func TestManagerWaitForServerConnectionUsesRuntimeReadiness(t *testing.T) {
	connected := make(chan struct{})
	manager := New(nil, func(chCfg profile.AgentTeamServer, _ types.ServerConnectionContext) (types.ServerConnection, error) {
		return &managerTestServer{
			id:        chCfg.ID,
			enabled:   true,
			identity:  chCfg,
			recorder:  managerTestRecorderFor(chCfg.ID),
			connected: connected,
		}, nil
	})
	server := newManagerTestProfileServer("ch-readiness", "agent-a", map[string]any{"mode": "alpha"})
	require.NoError(t, manager.Reconcile(nil, &server))

	done := make(chan error, 1)
	go func() {
		done <- manager.WaitForServerConnection(context.Background(), server.ID)
	}()
	select {
	case err := <-done:
		t.Fatalf("readiness returned before connection: %v", err)
	case <-time.After(20 * time.Millisecond):
	}
	close(connected)
	require.NoError(t, <-done)
}

func TestManagerReconcile_UpdateWithRuntimeChangeRestartsServer(t *testing.T) {
	manager := New(nil, newManagerTestServer)
	initial := newManagerTestProfileServer("ch-restart", "agent-a", map[string]any{"mode": "alpha"})
	updated := newManagerTestProfileServer("ch-restart", "agent-a", map[string]any{"mode": "beta"})

	require.NoError(t, manager.Reconcile(nil, &initial))
	require.NoError(t, manager.Reconcile(&initial, &updated))

	recorder := managerTestRecorderFor(initial.ID)
	starts, stops, signatures := recorder.snapshot()
	require.Equal(t, 2, starts)
	require.Equal(t, 1, stops)
	require.Equal(t, []string{managerTestSignature(initial), managerTestSignature(updated)}, signatures)
}

func TestManagerReconcile_WorkspaceBindingAppliesWithoutRestart(t *testing.T) {
	manager := New(nil, newManagerTestServer)
	initial := newManagerTestProfileServer("ch-binding", "agent-a", map[string]any{"mode": "alpha"})
	initial.WorkspaceID = "wsp-old"
	updated := initial
	updated.WorkspaceID = "wsp-new"

	require.NoError(t, manager.Reconcile(nil, &initial))
	runtime, ok := manager.ServerConnection(initial.ID)
	require.True(t, ok)
	require.NoError(t, manager.Reconcile(&initial, &updated))

	starts, stops, _ := managerTestRecorderFor(initial.ID).snapshot()
	require.Equal(t, 1, starts)
	require.Equal(t, 0, stops)
	require.Same(t, runtime, mustManagerRuntime(t, manager, initial.ID))
	require.Equal(t, "wsp-new", runtime.(*managerTestServer).ServerIdentitySnapshot().WorkspaceID)
}

func TestManagerReconcile_ConnectionStatusDoesNotRestart(t *testing.T) {
	manager := New(nil, newManagerTestServer)
	initial := newManagerTestProfileServer("ch-status", "agent-a", map[string]any{"mode": "alpha"})
	updated := initial
	updated.ConnectionStatus = "invalid_machine_api_key"

	require.NoError(t, manager.Reconcile(nil, &initial))
	require.NoError(t, manager.Reconcile(&initial, &updated))

	starts, stops, _ := managerTestRecorderFor(initial.ID).snapshot()
	require.Equal(t, 1, starts)
	require.Equal(t, 0, stops)
}

func TestManagerReconcile_WorkspaceBindingFailureKeepsKnownRuntimeConfig(t *testing.T) {
	manager := New(nil, newManagerTestServer)
	initial := newManagerTestProfileServer("ch-binding-fail", "agent-a", map[string]any{"mode": "alpha"})
	initial.WorkspaceID = "wsp-old"
	failed := initial
	failed.WorkspaceID = "wsp-failed"
	winner := initial
	winner.WorkspaceID = "wsp-winner"

	require.NoError(t, manager.Reconcile(nil, &initial))
	runtime := mustManagerRuntime(t, manager, initial.ID).(*managerTestServer)
	runtime.applyErr = errors.New("apply failed")
	require.ErrorContains(t, manager.Reconcile(&initial, &failed), "apply failed")

	runtime.applyErr = nil
	require.NoError(t, manager.Reconcile(&initial, &winner))
	require.Equal(t, "wsp-winner", runtime.ServerIdentitySnapshot().WorkspaceID)
	starts, stops, _ := managerTestRecorderFor(initial.ID).snapshot()
	require.Equal(t, 1, starts)
	require.Equal(t, 0, stops)
}

func TestManagerReconcile_DeleteStopsAndUnregistersRunningServer(t *testing.T) {
	manager := New(nil, newManagerTestServer)
	initial := newManagerTestProfileServer("ch-delete", "agent-a", map[string]any{"mode": "alpha"})

	require.NoError(t, manager.Reconcile(nil, &initial))
	_, ok := manager.ServerConnection(initial.ID)
	require.True(t, ok)

	require.NoError(t, manager.Reconcile(&initial, nil))

	recorder := managerTestRecorderFor(initial.ID)
	starts, stops, _ := recorder.snapshot()
	require.Equal(t, 1, starts)
	require.Equal(t, 1, stops)
	_, ok = manager.ServerConnection(initial.ID)
	require.False(t, ok)
}

func TestManagerReconcile_StopFailureKeepsOldRuntimeAndDoesNotStartReplacement(t *testing.T) {
	manager := New(nil, newManagerTestServer)
	initial := newManagerTestProfileServer("ch-stop-failure", "agent-a", map[string]any{"mode": "alpha"})
	updated := newManagerTestProfileServer("ch-stop-failure", "agent-a", map[string]any{"mode": "beta"})

	require.NoError(t, manager.Reconcile(nil, &initial))
	runtime, ok := manager.ServerConnection(initial.ID)
	require.True(t, ok)
	runtime.(*managerTestServer).stopErr = errors.New("stop timed out")

	err := manager.Reconcile(&initial, &updated)
	require.ErrorContains(t, err, "stop timed out")

	current, ok := manager.ServerConnection(initial.ID)
	require.True(t, ok)
	require.Same(t, runtime, current)
	starts, stops, signatures := managerTestRecorderFor(initial.ID).snapshot()
	require.Equal(t, 1, starts)
	require.Equal(t, 1, stops)
	require.Equal(t, []string{managerTestSignature(initial)}, signatures)
}

func TestManagerRepairServerStopsResidualRuntimeBeforeStartingPersistedIdentity(t *testing.T) {
	manager := New(nil, newManagerTestServer)
	residual := newManagerTestProfileServer("ch-repair", "agent-stale", map[string]any{"mode": "stale"})
	persisted := newManagerTestProfileServer("ch-repair", "agent-current", map[string]any{"mode": "current"})

	require.NoError(t, manager.Reconcile(nil, &residual))
	require.NoError(t, manager.RepairServer(context.Background(), persisted))

	current, ok := manager.ServerConnection(persisted.ID)
	require.True(t, ok)
	require.Equal(t, managerTestSignature(persisted), current.(*managerTestServer).signature)
	starts, stops, signatures := managerTestRecorderFor(persisted.ID).snapshot()
	require.Equal(t, 2, starts)
	require.Equal(t, 1, stops)
	require.Equal(t, []string{managerTestSignature(residual), managerTestSignature(persisted)}, signatures)
}

func TestManagerRepairServerStopFailureKeepsResidualRuntime(t *testing.T) {
	manager := New(nil, newManagerTestServer)
	residual := newManagerTestProfileServer("ch-repair-stop", "agent-stale", map[string]any{"mode": "stale"})
	persisted := newManagerTestProfileServer("ch-repair-stop", "agent-current", map[string]any{"mode": "current"})

	require.NoError(t, manager.Reconcile(nil, &residual))
	runtime, ok := manager.ServerConnection(residual.ID)
	require.True(t, ok)
	runtime.(*managerTestServer).stopErr = errors.New("stop timed out")

	err := manager.RepairServer(context.Background(), persisted)
	require.ErrorContains(t, err, "stop timed out")
	current, ok := manager.ServerConnection(residual.ID)
	require.True(t, ok)
	require.Same(t, runtime, current)
	starts, stops, _ := managerTestRecorderFor(residual.ID).snapshot()
	require.Equal(t, 1, starts)
	require.Equal(t, 1, stops)
}

func TestManagerRepairServerStartFailureRemainsRegisteredWhenCleanupFails(t *testing.T) {
	manager := New(nil, func(chCfg profile.AgentTeamServer, _ types.ServerConnectionContext) (types.ServerConnection, error) {
		return &managerTestServer{
			id:        chCfg.ID,
			enabled:   chCfg.Enabled,
			signature: managerTestSignature(chCfg),
			identity:  chCfg,
			recorder:  managerTestRecorderFor(chCfg.ID),
			startErr:  errors.New("start failed"),
			stopErr:   errors.New("cleanup failed"),
		}, nil
	})
	persisted := newManagerTestProfileServer("ch-repair-start", "agent-current", map[string]any{"mode": "current"})

	err := manager.RepairServer(context.Background(), persisted)
	require.ErrorContains(t, err, "start failed")
	require.ErrorContains(t, err, "cleanup failed")
	_, ok := manager.ServerConnection(persisted.ID)
	require.True(t, ok, "unconfirmed cleanup must not drop the only handle to the residual runtime")
}

func TestManagerRepairServerRejectsWrongRuntimeIdentity(t *testing.T) {
	manager := New(nil, func(chCfg profile.AgentTeamServer, _ types.ServerConnectionContext) (types.ServerConnection, error) {
		runtime, err := newManagerTestServer(chCfg, nil)
		runtime.(*managerTestServer).identity.WorkspaceID = "wsp-wrong"
		return runtime, err
	})
	persisted := newManagerTestProfileServer("ch-repair-identity", "agent-current", map[string]any{"mode": "current"})
	persisted.WorkspaceID = "wsp-current"

	err := manager.RepairServer(context.Background(), persisted)
	require.ErrorContains(t, err, "identity verification failed")
	_, ok := manager.ServerConnection(persisted.ID)
	require.False(t, ok)
}

func TestManagerReconcile_StopHonorsCallerDeadline(t *testing.T) {
	manager := New(nil, newManagerTestServer)
	initial := newManagerTestProfileServer("ch-deadline", "agent-a", map[string]any{"mode": "alpha"})
	updated := newManagerTestProfileServer("ch-deadline", "agent-a", map[string]any{"mode": "beta"})
	require.NoError(t, manager.Reconcile(nil, &initial))
	runtime := mustManagerRuntime(t, manager, initial.ID).(*managerTestServer)
	runtime.stopFunc = func(ctx context.Context) error {
		<-ctx.Done()
		return ctx.Err()
	}

	ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond)
	defer cancel()
	err := manager.ReconcileContext(ctx, &initial, &updated)
	require.ErrorIs(t, err, context.DeadlineExceeded)
}

func TestManagerRepairServerStopHonorsCallerDeadline(t *testing.T) {
	manager := New(nil, newManagerTestServer)
	initial := newManagerTestProfileServer("ch-repair-deadline", "agent-a", map[string]any{"mode": "alpha"})
	require.NoError(t, manager.Reconcile(nil, &initial))
	runtime := mustManagerRuntime(t, manager, initial.ID).(*managerTestServer)
	runtime.stopFunc = func(ctx context.Context) error {
		<-ctx.Done()
		return ctx.Err()
	}

	ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond)
	defer cancel()
	err := manager.RepairServer(ctx, initial)
	require.ErrorIs(t, err, context.DeadlineExceeded)
}

func TestManagerRepairServerRejectsConcurrentReplacement(t *testing.T) {
	manager := New(nil, newManagerTestServer)
	persisted := newManagerTestProfileServer("ch-repair-race", "agent-current", map[string]any{"mode": "current"})
	replacement := &managerTestServer{
		id:        persisted.ID,
		enabled:   true,
		signature: "concurrent",
		identity:  persisted,
		recorder:  managerTestRecorderFor(persisted.ID),
	}
	constructed := false
	manager.newServerConnection = func(chCfg profile.AgentTeamServer, _ types.ServerConnectionContext) (types.ServerConnection, error) {
		if constructed {
			return newManagerTestServer(chCfg, nil)
		}
		constructed = true
		runtime, err := newManagerTestServer(chCfg, nil)
		if err == nil {
			runtime.(*managerTestServer).onStart = func() {
				manager.Register(replacement)
			}
		}
		return runtime, err
	}

	err := manager.RepairServer(context.Background(), persisted)
	require.ErrorContains(t, err, "superseded")
	current, ok := manager.ServerConnection(persisted.ID)
	require.True(t, ok)
	require.Same(t, replacement, current)
}

func newManagerTestServer(chCfg profile.AgentTeamServer, _ types.ServerConnectionContext) (types.ServerConnection, error) {
	return &managerTestServer{
		id:        chCfg.ID,
		enabled:   chCfg.Enabled,
		signature: managerTestSignature(chCfg),
		identity:  chCfg,
		recorder:  managerTestRecorderFor(chCfg.ID),
	}, nil
}

func mustManagerRuntime(t *testing.T, manager *Manager, id string) types.ServerConnection {
	t.Helper()
	runtime, ok := manager.ServerConnection(id)
	require.True(t, ok)
	return runtime
}

func managerTestRecorderFor(id string) *managerTestRecorder {
	recorder, _ := managerTestRecorders.LoadOrStore(id, &managerTestRecorder{})
	return recorder.(*managerTestRecorder)
}

var managerTestRecorders sync.Map

func newManagerTestProfileServer(id, agentID string, config map[string]any) profile.AgentTeamServer {
	raw, err := yaml.Marshal(config)
	if err != nil {
		panic(fmt.Sprintf("marshal test server config: %v", err))
	}
	return profile.AgentTeamServer{
		ID:      id,
		Enabled: true,
		AgentID: agentID,
		Config:  raw,
	}
}

func managerTestSignature(ch profile.AgentTeamServer) string {
	return fmt.Sprintf("%s|%s", ch.AgentID, string(ch.Config))
}

func TestServerConnectionEquivalent_RuntimeFields(t *testing.T) {
	equiv := func(prev, curr profile.AgentTeamServer) bool {
		return serverConnectionEquivalent(prev, curr)
	}

	require.True(t, equiv(
		profile.AgentTeamServer{ID: "ch", AgentID: "agt-a", Config: []byte("a: 1\n")},
		profile.AgentTeamServer{ID: "ch", AgentID: "agt-a", Config: []byte("a: 1\n")},
	))
	require.False(t, equiv(
		profile.AgentTeamServer{ID: "ch", AgentID: "agt-a"},
		profile.AgentTeamServer{ID: "ch", AgentID: "agt-b"},
	))
	require.False(t, equiv(
		profile.AgentTeamServer{ID: "ch", AgentID: "agt-a", Config: []byte("a: 1\n")},
		profile.AgentTeamServer{ID: "ch", AgentID: "agt-a", Config: []byte("a: 2\n")},
	))
}