package machine

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"testing"

	"github.com/gin-gonic/gin"
	"github.com/stretchr/testify/require"

	"github.com/tidbcloud/AgentTeam/desktop/internal/api/auth"
	"github.com/tidbcloud/AgentTeam/desktop/internal/profile"
	"github.com/tidbcloud/AgentTeam/desktop/internal/token"
	"github.com/tidbcloud/AgentTeam/desktop/internal/uisession"
)

func TestEnsureCurrentBindingResolvesLegacyBindingBeforeCreating(t *testing.T) {
	setMachineTestEncryptionKeys(t)
	for _, onlyIfUserHasNoMachines := range []bool{false, true} {
		t.Run(fmt.Sprintf("only_if_user_has_no_machines_%t", onlyIfUserHasNoMachines), func(t *testing.T) {
			var listRequests, createRequests, resolveRequests int
			upstream := newMachineTestServer(t, func(writer http.ResponseWriter, request *http.Request) {
				switch {
				case request.URL.Path == "/api/v1/machine-auth-resolutions":
					resolveRequests++
					_, _ = writer.Write([]byte(`{"machineId":"mch_existing","workspaceId":"wsp_1"}`))
				case request.Method == http.MethodGet && request.URL.Path == "/api/v1/workspaces/wsp_1/machines":
					listRequests++
				case request.Method == http.MethodPost && request.URL.Path == "/api/v1/workspaces/wsp_1/machines":
					createRequests++
				default:
					http.NotFound(writer, request)
				}
			})
			authManager, target := loginMachineTestUser(t, upstream.URL)
			profileManager := profile.New(t.TempDir())
			legacy := addMachineTestBinding(t, profileManager, upstream.URL, "", "legacy-secret")
			readiness := &machineTestReadiness{}
			service := NewService(Dependencies{
				Auth:                authManager,
				ConnectionReadiness: readiness,
				Servers:             profileManager.Servers(),
			})

			binding, created, status, err := service.EnsureCurrentBinding(context.Background(), target, onlyIfUserHasNoMachines)

			require.NoError(t, err)
			require.False(t, created)
			require.Equal(t, http.StatusOK, status)
			require.Equal(t, "mch_existing", binding.MachineID)
			require.Equal(t, 1, resolveRequests)
			require.Zero(t, listRequests)
			require.Zero(t, createRequests)
			servers := profileManager.Servers().All()
			require.Len(t, servers, 1)
			require.Equal(t, legacy.ID, servers[0].ID)
			require.Equal(t, "mch_existing", servers[0].MachineID)
			require.Equal(t, "legacy-secret", servers[0].APIKey)
			require.Equal(t, legacy.ID, readiness.serverID)
		})
	}
}

func TestEnsureCurrentBindingDoesNotReplaceUnresolvedLegacyBinding(t *testing.T) {
	setMachineTestEncryptionKeys(t)
	var createRequests int
	upstream := newMachineTestServer(t, func(writer http.ResponseWriter, request *http.Request) {
		switch {
		case request.URL.Path == "/api/v1/machine-auth-resolutions":
			writer.WriteHeader(http.StatusUnauthorized)
			_, _ = writer.Write([]byte(`{"error":"invalid machine API key"}`))
		case request.Method == http.MethodPost && request.URL.Path == "/api/v1/workspaces/wsp_1/machines":
			createRequests++
		default:
			http.NotFound(writer, request)
		}
	})
	authManager, target := loginMachineTestUser(t, upstream.URL)
	profileManager := profile.New(t.TempDir())
	legacy := addMachineTestBinding(t, profileManager, upstream.URL, "", "legacy-secret")
	service := NewService(Dependencies{Auth: authManager, Servers: profileManager.Servers()})

	_, created, status, err := service.EnsureCurrentBinding(context.Background(), target, false)

	require.ErrorContains(t, err, "restore an existing Machine")
	require.False(t, created)
	require.Equal(t, http.StatusConflict, status)
	require.Zero(t, createRequests)
	servers := profileManager.Servers().All()
	require.Len(t, servers, 1)
	require.Equal(t, legacy.ID, servers[0].ID)
	require.Empty(t, servers[0].MachineID)
	require.Equal(t, "legacy-secret", servers[0].APIKey)
}

func TestCurrentBindingRequiresConfirmationBeforeReplacingUnresolvedLegacyBinding(t *testing.T) {
	setMachineTestEncryptionKeys(t)
	var createRequests int
	upstream := newMachineTestServer(t, func(writer http.ResponseWriter, request *http.Request) {
		switch {
		case request.URL.Path == "/api/v1/machine-auth-resolutions":
			writer.WriteHeader(http.StatusUnauthorized)
			_, _ = writer.Write([]byte(`{"error":"invalid machine API key"}`))
		case request.Method == http.MethodPost && request.URL.Path == "/api/v1/workspaces/wsp_1/machines":
			createRequests++
			_, _ = writer.Write([]byte(`{"machine":{"id":"mch_replacement"},"apiKey":"replacement-secret"}`))
		default:
			http.NotFound(writer, request)
		}
	})
	authManager, target := loginMachineTestUser(t, upstream.URL)
	profileManager := profile.New(t.TempDir())
	legacy := addMachineTestBinding(t, profileManager, upstream.URL, "", "legacy-secret")
	sessions := uisession.NewStore(uisession.Options{})
	sessionID := token.CreateSession()
	expiresAt, ok := token.SessionExpiresAt(sessionID)
	require.True(t, ok)
	_, err := sessions.Select(sessionID, expiresAt, target.ID)
	require.NoError(t, err)

	gin.SetMode(gin.TestMode)
	router := gin.New()
	internal := router.Group("/api/internal", auth.Middleware())
	RegisterRoutes(internal, Dependencies{
		Auth:     authManager,
		Servers:  profileManager.Servers(),
		Sessions: sessions,
	})

	current := performMachineRequest(router, sessionID, http.MethodGet, "/api/internal/machines/current-binding?workspaceId=wsp_1", nil)
	require.Equal(t, http.StatusOK, current.Code)
	require.JSONEq(t, `{"status":"recovery_required"}`, current.Body.String())

	unconfirmed := performMachineRequest(router, sessionID, http.MethodPost, "/api/internal/machines/current-bindings", map[string]any{
		"workspaceId": "wsp_1",
	})
	require.Equal(t, http.StatusConflict, unconfirmed.Code)
	require.JSONEq(t, `{"code":"DESKTOP_MACHINE_BINDING_RECOVERY_REQUIRED","error":"the saved Desktop connection could not be verified; restore an existing Machine or confirm replacement"}`, unconfirmed.Body.String())
	require.Zero(t, createRequests)

	confirmed := performMachineRequest(router, sessionID, http.MethodPost, "/api/internal/machines/current-bindings", map[string]any{
		"workspaceId":       "wsp_1",
		"replaceUnresolved": true,
	})
	require.Equal(t, http.StatusOK, confirmed.Code)
	require.JSONEq(t, `{"machineId":"mch_replacement","status":"created"}`, confirmed.Body.String())
	require.Equal(t, 1, createRequests)
	servers := profileManager.Servers().All()
	require.Len(t, servers, 1)
	require.Equal(t, legacy.ID, servers[0].ID)
	require.Equal(t, "mch_replacement", servers[0].MachineID)
	require.Equal(t, "replacement-secret", servers[0].APIKey)
}

func TestEnsureCurrentBindingRejectsLegacyBindingFromAnotherWorkspace(t *testing.T) {
	setMachineTestEncryptionKeys(t)
	var createRequests int
	upstream := newMachineTestServer(t, func(writer http.ResponseWriter, request *http.Request) {
		switch {
		case request.URL.Path == "/api/v1/machine-auth-resolutions":
			_, _ = writer.Write([]byte(`{"machineId":"mch_foreign","workspaceId":"wsp_foreign"}`))
		case request.Method == http.MethodPost && request.URL.Path == "/api/v1/workspaces/wsp_1/machines":
			createRequests++
		default:
			http.NotFound(writer, request)
		}
	})
	authManager, target := loginMachineTestUser(t, upstream.URL)
	profileManager := profile.New(t.TempDir())
	legacy := addMachineTestBinding(t, profileManager, upstream.URL, "", "legacy-secret")
	service := NewService(Dependencies{Auth: authManager, Servers: profileManager.Servers()})

	_, created, status, err := service.EnsureCurrentBinding(context.Background(), target, false)

	require.ErrorContains(t, err, "restore an existing Machine")
	require.False(t, created)
	require.Equal(t, http.StatusConflict, status)
	require.Zero(t, createRequests)
	servers := profileManager.Servers().All()
	require.Len(t, servers, 1)
	require.Equal(t, legacy.ID, servers[0].ID)
	require.Empty(t, servers[0].MachineID)
}

func TestCurrentBindingPrefersCompleteBindingOverLegacyCandidate(t *testing.T) {
	setMachineTestEncryptionKeys(t)
	var legacyResolveRequests int
	upstream := newMachineTestServer(t, func(writer http.ResponseWriter, request *http.Request) {
		var body map[string]string
		require.NoError(t, json.NewDecoder(request.Body).Decode(&body))
		if body["apiKey"] == "legacy-secret" {
			legacyResolveRequests++
		}
		_, _ = writer.Write([]byte(`{"machineId":"mch_current","workspaceId":"wsp_1"}`))
	})
	authManager, target := loginMachineTestUser(t, upstream.URL)
	profileManager := profile.New(t.TempDir())
	complete := addMachineTestBinding(t, profileManager, upstream.URL, "mch_current", "current-secret")
	legacy := addMachineTestBinding(t, profileManager, upstream.URL, "", "legacy-secret")
	service := NewService(Dependencies{Auth: authManager, Servers: profileManager.Servers()})

	binding, exists, status, err := service.ReconcileCurrentBinding(context.Background(), target)

	require.NoError(t, err)
	require.True(t, exists)
	require.Equal(t, http.StatusOK, status)
	require.Equal(t, "mch_current", binding.MachineID)
	require.Zero(t, legacyResolveRequests)
	servers := profileManager.Servers().All()
	require.Len(t, servers, 2)
	require.Equal(t, complete.ID, servers[0].ID)
	require.Equal(t, legacy.ID, servers[1].ID)
	require.Empty(t, servers[1].MachineID)
}

func TestReconcileCurrentBindingKeepsLegacyBindingOnServerFailure(t *testing.T) {
	setMachineTestEncryptionKeys(t)
	upstream := newMachineTestServer(t, func(writer http.ResponseWriter, request *http.Request) {
		writer.WriteHeader(http.StatusServiceUnavailable)
		_, _ = writer.Write([]byte(`{"error":"temporarily unavailable"}`))
	})
	authManager, target := loginMachineTestUser(t, upstream.URL)
	profileManager := profile.New(t.TempDir())
	legacy := addMachineTestBinding(t, profileManager, upstream.URL, "", "legacy-secret")
	service := NewService(Dependencies{Auth: authManager, Servers: profileManager.Servers()})

	_, exists, status, err := service.ReconcileCurrentBinding(context.Background(), target)

	require.ErrorContains(t, err, "temporarily unavailable")
	require.False(t, exists)
	require.Equal(t, http.StatusServiceUnavailable, status)
	servers := profileManager.Servers().All()
	require.Len(t, servers, 1)
	require.Equal(t, legacy.ID, servers[0].ID)
	require.Empty(t, servers[0].MachineID)
	require.Equal(t, "legacy-secret", servers[0].APIKey)
}