package daemon

import (
	"context"
	"errors"
	"net"
	"net/http"
	"testing"
	"time"

	"github.com/stretchr/testify/require"

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

type shutdownTestTransition struct{}

func (shutdownTestTransition) Complete()   {}
func (shutdownTestTransition) FailClosed() {}

type shutdownTestRuntimeParticipant struct{}

func (shutdownTestRuntimeParticipant) BeginServerTransition(string) profile.ServerTransition {
	return shutdownTestTransition{}
}

func (shutdownTestRuntimeParticipant) ReconcileServer(
	context.Context,
	*profile.AgentTeamServer,
	*profile.AgentTeamServer,
) error {
	return nil
}

func (shutdownTestRuntimeParticipant) RepairServer(
	context.Context,
	profile.AgentTeamServer,
) error {
	return nil
}

func (shutdownTestRuntimeParticipant) InvalidateServer(string) {}

func TestShutdownDrainsAcceptedHandlersBeforeReturningServerMutationsToOfflineMode(t *testing.T) {
	configDir := t.TempDir()
	t.Setenv(profile.ConfigDirEnvKey, configDir)
	servers := profile.Default().Servers()
	server := &profile.AgentTeamServer{
		ID:      "shutdown-active-handler",
		Enabled: true,
		Name:    "Before shutdown",
	}
	require.NoError(t, servers.Add(server))
	require.NoError(t, servers.PrepareRuntimeMode())
	require.NoError(t, servers.ActivateRuntimeMode(shutdownTestRuntimeParticipant{}))
	t.Cleanup(servers.DeactivateRuntimeMode)

	handlerEntered := make(chan struct{})
	releaseHandler := make(chan struct{})
	patchResult := make(chan error, 1)
	handler := http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) {
		close(handlerEntered)
		<-releaseHandler
		nextName := "After shutdown"
		_, err := servers.Patch(server.ID, profile.ServerUpdate{Name: &nextName})
		patchResult <- err
		writer.WriteHeader(http.StatusNoContent)
	})
	listener, err := net.Listen("tcp", "127.0.0.1:0")
	require.NoError(t, err)
	apiServer := &http.Server{Handler: handler}
	go func() {
		_ = apiServer.Serve(listener)
	}()
	t.Cleanup(func() { _ = apiServer.Close() })

	requestDone := make(chan error, 1)
	go func() {
		response, requestErr := http.Get("http://" + listener.Addr().String())
		if response != nil {
			_ = response.Body.Close()
		}
		requestDone <- requestErr
	}()
	<-handlerEntered

	d := &Daemon{
		apiServer:          apiServer,
		shutdownAPIServer:  apiServer.Shutdown,
		apiShutdownTimeout: time.Second,
	}
	shutdownDone := make(chan struct{})
	go func() {
		d.shutdown()
		close(shutdownDone)
	}()

	require.Eventually(t, func() bool {
		return servers.LifecycleMode() == profile.ServerLifecyclePreparing
	}, time.Second, time.Millisecond)
	select {
	case <-shutdownDone:
		t.Fatal("shutdown returned before the accepted handler exited")
	default:
	}

	close(releaseHandler)
	require.ErrorIs(t, <-patchResult, profile.ErrServerRuntimePreparing)
	require.NoError(t, <-requestDone)
	select {
	case <-shutdownDone:
	case <-time.After(time.Second):
		t.Fatal("shutdown did not return after the accepted handler exited")
	}

	require.Equal(t, profile.ServerLifecycleOffline, servers.LifecycleMode())
	current, ok := servers.Get(server.ID)
	require.True(t, ok)
	require.Equal(t, "Before shutdown", current.Name)
}

func TestShutdownTimeoutKeepsServerMutationsFailClosed(t *testing.T) {
	configDir := t.TempDir()
	t.Setenv(profile.ConfigDirEnvKey, configDir)
	servers := profile.Default().Servers()
	require.NoError(t, servers.PrepareRuntimeMode())
	require.NoError(t, servers.ActivateRuntimeMode(shutdownTestRuntimeParticipant{}))
	t.Cleanup(servers.DeactivateRuntimeMode)

	d := &Daemon{
		apiShutdownTimeout: time.Millisecond,
		shutdownAPIServer: func(ctx context.Context) error {
			<-ctx.Done()
			return ctx.Err()
		},
	}
	d.shutdown()

	require.Equal(t, profile.ServerLifecyclePreparing, servers.LifecycleMode())
	_, err := servers.Patch("missing", profile.ServerUpdate{})
	require.ErrorIs(t, err, profile.ErrServerRuntimePreparing)
}

func TestShutdownCancelsGatewayOpsBeforeHTTPDrain(t *testing.T) {
	configDir := t.TempDir()
	t.Setenv(profile.ConfigDirEnvKey, configDir)
	servers := profile.Default().Servers()
	require.NoError(t, servers.PrepareRuntimeMode())
	require.NoError(t, servers.ActivateRuntimeMode(shutdownTestRuntimeParticipant{}))
	t.Cleanup(servers.DeactivateRuntimeMode)

	registry := streamregistry.New()
	entry, err := registry.Register(context.Background(), "session-a", 7, "server-a")
	require.NoError(t, err)
	t.Cleanup(func() {
		entry.Cancel(streamregistry.CancelError{Reason: streamregistry.ReasonDaemonShutdown})
		entry.Release()
	})

	handlerEntered := make(chan struct{})
	handlerExited := make(chan struct{})
	handler := http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		close(handlerEntered)
		<-entry.Context().Done()
		entry.Release()
		close(handlerExited)
		w.WriteHeader(http.StatusNoContent)
	})
	listener, err := net.Listen("tcp", "127.0.0.1:0")
	require.NoError(t, err)
	apiServer := &http.Server{Handler: handler}
	go func() { _ = apiServer.Serve(listener) }()
	t.Cleanup(func() { _ = apiServer.Close() })

	requestDone := make(chan struct{})
	go func() {
		defer close(requestDone)
		resp, _ := http.Get("http://" + listener.Addr().String())
		if resp != nil {
			_ = resp.Body.Close()
		}
	}()
	select {
	case <-handlerEntered:
	case <-time.After(time.Second):
		t.Fatal("HTTP handler did not start")
	}

	d := &Daemon{
		streamRegistry:     registry,
		apiServer:          apiServer,
		shutdownAPIServer:  apiServer.Shutdown,
		apiShutdownTimeout: 200 * time.Millisecond,
	}
	d.shutdown()

	// The registry cancel must happen before HTTP Shutdown waits for the
	// accepted handler. The handler releases the entry only after observing the
	// typed shutdown cause, allowing both drains to complete and runtime mode
	// to transition Offline.
	require.ErrorIs(t, context.Cause(entry.Context()), streamregistry.ReasonDaemonShutdown)
	select {
	case <-handlerExited:
	case <-time.After(time.Second):
		t.Fatal("gateway operation did not exit after daemon shutdown cancel")
	}
	select {
	case <-requestDone:
	case <-time.After(time.Second):
		t.Fatal("HTTP request did not drain")
	}
	require.Zero(t, registry.ActiveGatewayOps())
	require.Equal(t, profile.ServerLifecycleOffline, servers.LifecycleMode())
}

func TestShutdownRegistryTimeoutStillDrainsHTTPAndKeepsPreparing(t *testing.T) {
	configDir := t.TempDir()
	t.Setenv(profile.ConfigDirEnvKey, configDir)
	servers := profile.Default().Servers()
	require.NoError(t, servers.PrepareRuntimeMode())
	require.NoError(t, servers.ActivateRuntimeMode(shutdownTestRuntimeParticipant{}))
	t.Cleanup(servers.DeactivateRuntimeMode)

	registry := streamregistry.New()
	entry, err := registry.Register(context.Background(), "session-stuck", 9, "server-a")
	require.NoError(t, err)
	t.Cleanup(entry.Release)

	listener, err := net.Listen("tcp", "127.0.0.1:0")
	require.NoError(t, err)
	apiServer := &http.Server{Handler: http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		w.WriteHeader(http.StatusNoContent)
	})}
	go func() { _ = apiServer.Serve(listener) }()
	t.Cleanup(func() { _ = apiServer.Close() })

	var httpShutdownCalls int
	d := &Daemon{
		streamRegistry:     registry,
		apiServer:          apiServer,
		apiShutdownTimeout: 30 * time.Millisecond,
		shutdownAPIServer: func(ctx context.Context) error {
			httpShutdownCalls++
			// Registry wait consumed its own budget. HTTP shutdown must receive
			// a fresh context rather than an already-expired shared context.
			require.NoError(t, ctx.Err())
			return apiServer.Shutdown(ctx)
		},
	}
	d.shutdown()

	require.Equal(t, 1, httpShutdownCalls,
		"HTTP Shutdown must still run after registry drain times out")
	require.ErrorIs(t, context.Cause(entry.Context()), streamregistry.ReasonDaemonShutdown)
	_, registerErr := registry.Register(context.Background(), "late", 1, "server-a")
	require.ErrorIs(t, registerErr, streamregistry.ErrRegistryClosed,
		"CancelAll must close admission before bounded wait")
	require.Equal(t, profile.ServerLifecyclePreparing, servers.LifecycleMode(),
		"registry timeout must keep server mutations fail-closed")
	_, err = servers.Patch("missing", profile.ServerUpdate{})
	require.True(t,
		errors.Is(err, profile.ErrServerRuntimePreparing),
		"server mutation error = %v, want ErrServerRuntimePreparing", err)
}