package onboarding

import (
	"bytes"
	"context"
	"crypto/ed25519"
	"crypto/rand"
	"encoding/json"
	"fmt"
	"net/http"
	"net/http/httptest"
	"sync/atomic"
	"testing"
	"time"

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

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

func TestRoutesReturnUnavailableWithoutManager(t *testing.T) {
	gin.SetMode(gin.TestMode)
	router := gin.New()
	internal := router.Group("/api/internal", auth.Middleware())
	RegisterRoutes(internal, Dependencies{})

	recorder := performRequest(router, token.CreateSession(), http.MethodGet, "/api/internal/onboarding/servers", nil)

	require.Equal(t, http.StatusServiceUnavailable, recorder.Code)
	require.JSONEq(t, `{"error":"Desktop authentication is unavailable."}`, recorder.Body.String())
}

func TestProbeCancellationDoesNotPreventReconnect(t *testing.T) {
	requestStarted := make(chan struct{})
	upstreamCanceled := make(chan struct{})
	var requestCount atomic.Int32
	upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
		if requestCount.Add(1) == 1 {
			close(requestStarted)
			<-request.Context().Done()
			close(upstreamCanceled)
			return
		}
		writer.Header().Set("Content-Type", "application/json")
		_, _ = writer.Write([]byte(`{"pingcapAuthEnabled":true,"passwordAuthEnabled":true}`))
	}))
	t.Cleanup(upstream.Close)

	manager, err := userauth.New(t.TempDir())
	require.NoError(t, err)
	gin.SetMode(gin.TestMode)
	router := gin.New()
	internal := router.Group("/api/internal", auth.Middleware())
	RegisterRoutes(internal, Dependencies{Auth: manager})

	requestContext, cancel := context.WithCancel(context.Background())
	defer cancel()
	request := httptest.NewRequest(http.MethodPost, "/api/internal/onboarding/probe", bytes.NewReader([]byte(`{"serverUrl":"`+upstream.URL+`"}`))).WithContext(requestContext)
	request.Header.Set("Authorization", "Bearer "+token.CreateSession())
	request.Header.Set("Content-Type", "application/json")
	recorder := httptest.NewRecorder()
	handlerDone := make(chan struct{})

	go func() {
		router.ServeHTTP(recorder, request)
		close(handlerDone)
	}()

	select {
	case <-requestStarted:
	case <-time.After(time.Second):
		t.Fatal("probe did not reach the upstream server")
	}
	cancel()
	select {
	case <-upstreamCanceled:
	case <-time.After(time.Second):
		t.Fatal("canceling the renderer request did not cancel the upstream probe")
	}
	select {
	case <-handlerDone:
	case <-time.After(time.Second):
		t.Fatal("probe handler did not return after cancellation")
	}

	require.NotEqual(t, http.StatusBadGateway, recorder.Code)
	require.Empty(t, recorder.Body.String())

	reconnect := performRequest(
		router,
		token.CreateSession(),
		http.MethodPost,
		"/api/internal/onboarding/probe",
		map[string]string{"serverUrl": upstream.URL},
	)
	require.Equal(t, http.StatusOK, reconnect.Code)
	require.JSONEq(t, `{
		"name":"127.0.0.1",
		"serverUrl":"`+upstream.URL+`",
		"pingcapAuthEnabled":true,
		"passwordAuthEnabled":true,
		"registrationMode":"invite_only",
		"openRegistration":false
	}`, reconnect.Body.String())
}

func TestResumeClassifiesSavedLoginFailures(t *testing.T) {
	for _, test := range []struct {
		name       string
		status     int
		wantStatus int
		wantBody   string
	}{
		{
			name:       "upstream rejects refresh token",
			status:     http.StatusUnauthorized,
			wantStatus: http.StatusUnauthorized,
			wantBody:   `{"code":"reauthentication_required","error":"Saved login requires reauthentication."}`,
		},
		{
			name:       "upstream is temporarily unavailable",
			status:     http.StatusServiceUnavailable,
			wantStatus: http.StatusServiceUnavailable,
			wantBody:   `{"code":"temporarily_unavailable","error":"Saved login is temporarily unavailable."}`,
		},
	} {
		t.Run(test.name, func(t *testing.T) {
			setTestEncryptionKeys(t)
			upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
				writer.Header().Set("Content-Type", "application/json")
				switch request.URL.Path {
				case "/api/v1/auth/login":
					_, _ = writer.Write([]byte(`{"accessToken":"access","refreshToken":"refresh","user":{"id":"usr_1","email":"user@example.com","name":"User"}}`))
				case "/api/v1/workspaces":
					_, _ = writer.Write([]byte(`[{"id":"wsp_1","name":"Workspace","slug":"workspace"}]`))
				case "/api/v1/auth/refresh":
					writer.WriteHeader(test.status)
					_, _ = writer.Write([]byte(`{"error":"upstream failure"}`))
				default:
					http.NotFound(writer, request)
				}
			}))
			t.Cleanup(upstream.Close)

			directory := t.TempDir()
			manager, err := userauth.New(directory)
			require.NoError(t, err)
			login, err := manager.Login(context.Background(), upstream.URL, "user@example.com", "password")
			require.NoError(t, err)
			manager, err = userauth.New(directory)
			require.NoError(t, err)

			router := gin.New()
			internal := router.Group("/api/internal", auth.Middleware())
			RegisterRoutes(internal, Dependencies{Auth: manager})
			recorder := performRequest(router, token.CreateSession(), http.MethodPost, "/api/internal/onboarding/resume", map[string]string{
				"authId": login.Server.ID,
			})

			require.Equal(t, test.wantStatus, recorder.Code)
			require.JSONEq(t, test.wantBody, recorder.Body.String())
		})
	}
}

func TestActivateLeavesOnboardingBootstrapToRenderer(t *testing.T) {
	setTestEncryptionKeys(t)
	bootstrapCalls := 0
	upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
		switch request.URL.Path {
		case "/api/v1/auth/login":
			writer.Header().Set("Content-Type", "application/json")
			_, _ = writer.Write([]byte(`{"accessToken":"access-secret","refreshToken":"refresh-secret","user":{"id":"usr_1","email":"user@example.com","name":"User"}}`))
		case "/api/v1/workspaces":
			require.Equal(t, "Bearer access-secret", request.Header.Get("Authorization"))
			writer.Header().Set("Content-Type", "application/json")
			_, _ = writer.Write([]byte(`[{"id":"wsp_1","name":"Workspace","slug":"workspace"}]`))
		case "/api/v1/workspaces/wsp_1/onboarding/bootstrap":
			bootstrapCalls++
			http.Error(writer, "unexpected daemon bootstrap", http.StatusInternalServerError)
		default:
			http.NotFound(writer, request)
		}
	}))
	t.Cleanup(upstream.Close)

	manager, err := userauth.New(t.TempDir())
	require.NoError(t, err)
	login, err := manager.Login(context.Background(), upstream.URL, "user@example.com", "password-secret")
	require.NoError(t, err)
	require.NotEmpty(t, login.Server.ID)

	sessions := uisession.NewStore(uisession.Options{})
	gin.SetMode(gin.TestMode)
	router := gin.New()
	internal := router.Group("/api/internal", auth.Middleware())
	RegisterRoutes(internal, Dependencies{Auth: manager, Sessions: sessions})
	sessionToken := token.CreateSession()
	recorder := performRequest(router, sessionToken, http.MethodPost, "/api/internal/onboarding/activate", map[string]string{
		"authId":      login.Server.ID,
		"workspaceId": "wsp_1",
	})

	require.Equal(t, http.StatusOK, recorder.Code)
	require.JSONEq(t, `{
		"server":{"id":"`+userauth.BuildTargetID(login.Server.ID, "wsp_1")+`","name":"127.0.0.1"},
		"binding":{"workspaceId":"wsp_1","machineId":""},
		"generation":1
	}`, recorder.Body.String())
	require.Zero(t, bootstrapCalls)
	require.NotContains(t, recorder.Body.String(), upstream.URL)
	require.NotContains(t, recorder.Body.String(), "access-secret")
	require.NotContains(t, recorder.Body.String(), "refresh-secret")
	require.NotContains(t, recorder.Body.String(), "password-secret")
	snapshot, ok := sessions.Get(sessionToken)
	require.True(t, ok)
	require.Equal(t, userauth.BuildTargetID(login.Server.ID, "wsp_1"), snapshot.ActiveServerID)
	lastTarget, ok := manager.LastActiveTarget()
	require.True(t, ok)
	require.Equal(t, snapshot.ActiveServerID, lastTarget)
}

func TestReauthenticatePreservesActiveWorkspaceGeneration(t *testing.T) {
	setTestEncryptionKeys(t)
	loginCount := 0
	upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
		writer.Header().Set("Content-Type", "application/json")
		switch request.URL.Path {
		case "/api/v1/public/config":
			_, _ = writer.Write([]byte(`{"pingcapAuthEnabled":true,"passwordAuthEnabled":true}`))
		case "/api/v1/auth/login":
			loginCount++
			_, _ = writer.Write([]byte(`{"accessToken":"access-` + fmt.Sprint(loginCount) + `","refreshToken":"refresh-` + fmt.Sprint(loginCount) + `","user":{"id":"usr_1","email":"user@example.com","name":"User"}}`))
		case "/api/v1/workspaces":
			_, _ = writer.Write([]byte(`[{"id":"wsp_1","name":"Workspace","slug":"workspace"}]`))
		default:
			http.NotFound(writer, request)
		}
	}))
	t.Cleanup(upstream.Close)

	manager, err := userauth.New(t.TempDir())
	require.NoError(t, err)
	login, err := manager.Login(context.Background(), upstream.URL, "user@example.com", "old-password")
	require.NoError(t, err)
	sessions := uisession.NewStore(uisession.Options{})
	router := gin.New()
	internal := router.Group("/api/internal", auth.Middleware())
	RegisterRoutes(internal, Dependencies{Auth: manager, Sessions: sessions})
	sessionToken := token.CreateSession()
	activate := performRequest(router, sessionToken, http.MethodPost, "/api/internal/onboarding/activate", map[string]string{
		"authId": login.Server.ID, "workspaceId": "wsp_1",
	})
	require.Equal(t, http.StatusOK, activate.Code)
	targetID := userauth.BuildTargetID(login.Server.ID, "wsp_1")

	contextRecorder := performRequest(router, sessionToken, http.MethodGet, "/api/internal/onboarding/reauthenticate", nil)
	require.Equal(t, http.StatusOK, contextRecorder.Code, contextRecorder.Body.String())
	require.JSONEq(t, `{
		"serverId":"`+targetID+`",
		"serverUrl":"`+upstream.URL+`",
		"serverName":"127.0.0.1",
		"pingcapAuthEnabled":true,
		"passwordAuthEnabled":true
	}`, contextRecorder.Body.String())

	recorder := performRequest(router, sessionToken, http.MethodPost, "/api/internal/onboarding/reauthenticate", map[string]string{
		"serverId": targetID, "identifier": "user@example.com", "password": "new-password",
	})

	require.Equal(t, http.StatusNoContent, recorder.Code, recorder.Body.String())
	snapshot, ok := sessions.Get(sessionToken)
	require.True(t, ok)
	require.Equal(t, targetID, snapshot.ActiveServerID)
	require.Equal(t, uint64(1), snapshot.Generation)
	_, err = manager.Resume(context.Background(), login.Server.ID)
	require.NoError(t, err)
	require.Equal(t, 2, loginCount)
}

func TestActivateCreatesAndConnectsFirstDesktopMachine(t *testing.T) {
	setTestEncryptionKeys(t)
	var createdMachines int
	upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
		writer.Header().Set("Content-Type", "application/json")
		switch {
		case request.URL.Path == "/api/v1/auth/login":
			_, _ = writer.Write([]byte(`{"accessToken":"access-secret","refreshToken":"refresh-secret","user":{"id":"usr_1","email":"user@example.com","name":"User"}}`))
		case request.URL.Path == "/api/v1/workspaces":
			_, _ = writer.Write([]byte(`[{"id":"wsp_1","name":"Workspace","slug":"workspace"}]`))
		case request.Method == http.MethodGet && request.URL.Path == "/api/v1/workspaces/wsp_1/machines":
			require.Equal(t, "own", request.URL.Query().Get("scope"))
			_, _ = writer.Write([]byte(`{"machines":[],"latestDaemonVersion":"","hasMore":false}`))
		case request.Method == http.MethodPost && request.URL.Path == "/api/v1/workspaces/wsp_1/machines":
			createdMachines++
			_, _ = writer.Write([]byte(`{"machine":{"id":"mch_local"},"apiKey":"machine-secret"}`))
		default:
			http.NotFound(writer, request)
		}
	}))
	t.Cleanup(upstream.Close)

	authManager, err := userauth.New(t.TempDir())
	require.NoError(t, err)
	login, err := authManager.Login(context.Background(), upstream.URL, "user@example.com", "password")
	require.NoError(t, err)
	profileManager := profile.New(t.TempDir())
	machineBindings := machine.NewService(machine.Dependencies{
		Auth:    authManager,
		Servers: profileManager.Servers(),
	})
	sessions := uisession.NewStore(uisession.Options{})
	router := gin.New()
	internal := router.Group("/api/internal", auth.Middleware())
	RegisterRoutes(internal, Dependencies{
		Auth:            authManager,
		MachineBindings: machineBindings,
		Servers:         profileManager.Servers(),
		Sessions:        sessions,
	})

	sessionToken := token.CreateSession()
	recorder := performRequest(router, sessionToken, http.MethodPost, "/api/internal/onboarding/activate", map[string]string{
		"authId": login.Server.ID, "workspaceId": "wsp_1",
	})

	require.Equal(t, http.StatusOK, recorder.Code)
	require.Contains(t, recorder.Body.String(), `"machineId":"mch_local"`)
	require.NotContains(t, recorder.Body.String(), "machine-secret")
	require.Equal(t, 1, createdMachines)
	servers := profileManager.Servers().All()
	require.Len(t, servers, 1)
	require.Equal(t, "mch_local", servers[0].MachineID)
	require.Equal(t, "machine-secret", servers[0].APIKey)
}

func TestLogoutActiveRemovesSavedAccountAndClearsSelection(t *testing.T) {
	setTestEncryptionKeys(t)
	upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
		writer.Header().Set("Content-Type", "application/json")
		switch request.URL.Path {
		case "/api/v1/auth/login":
			_, _ = writer.Write([]byte(`{"accessToken":"access","refreshToken":"refresh","user":{"id":"usr_1","email":"user@example.com","name":"User"}}`))
		case "/api/v1/workspaces":
			_, _ = writer.Write([]byte(`[{"id":"wsp_1","name":"Workspace","slug":"workspace"}]`))
		case "/api/v1/auth/logout":
			writer.WriteHeader(http.StatusNoContent)
		default:
			http.NotFound(writer, request)
		}
	}))
	t.Cleanup(upstream.Close)
	manager, err := userauth.New(t.TempDir())
	require.NoError(t, err)
	login, err := manager.Login(context.Background(), upstream.URL, "user@example.com", "password")
	require.NoError(t, err)

	sessions := uisession.NewStore(uisession.Options{})
	router := gin.New()
	internal := router.Group("/api/internal", auth.Middleware())
	RegisterRoutes(internal, Dependencies{Auth: manager, Sessions: sessions})
	sessionToken := token.CreateSession()
	activate := performRequest(router, sessionToken, http.MethodPost, "/api/internal/onboarding/activate", map[string]string{
		"authId": login.Server.ID, "workspaceId": "wsp_1",
	})
	require.Equal(t, http.StatusOK, activate.Code)

	logout := performRequest(router, sessionToken, http.MethodDelete, "/api/internal/onboarding/session", nil)
	require.Equal(t, http.StatusNoContent, logout.Code)
	snapshot, ok := sessions.Get(sessionToken)
	require.True(t, ok)
	require.Empty(t, snapshot.ActiveServerID)
	_, err = manager.Resume(context.Background(), login.Server.ID)
	require.EqualError(t, err, "saved server login not found")
	_, ok = manager.LastActiveTarget()
	require.False(t, ok)
}

func TestRemoveAccountEndpointInvalidatesEverySelectedWorkspaceTarget(t *testing.T) {
	setTestEncryptionKeys(t)
	upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
		writer.Header().Set("Content-Type", "application/json")
		switch request.URL.Path {
		case "/api/v1/auth/login":
			_, _ = writer.Write([]byte(`{"accessToken":"access","refreshToken":"refresh","user":{"id":"usr_1","email":"user@example.com","name":"User"}}`))
		case "/api/v1/workspaces":
			_, _ = writer.Write([]byte(`[{"id":"wsp_1","name":"One","slug":"one"},{"id":"wsp_2","name":"Two","slug":"two"}]`))
		case "/api/v1/auth/logout":
			writer.WriteHeader(http.StatusNoContent)
		default:
			http.NotFound(writer, request)
		}
	}))
	t.Cleanup(upstream.Close)
	manager, err := userauth.New(t.TempDir())
	require.NoError(t, err)
	login, err := manager.Login(context.Background(), upstream.URL, "user@example.com", "password")
	require.NoError(t, err)
	sessions := uisession.NewStore(uisession.Options{})
	firstSession := token.CreateSession()
	secondSession := token.CreateSession()
	expiresAt := time.Now().Add(time.Hour)
	_, err = sessions.Select(firstSession, expiresAt, userauth.BuildTargetID(login.Server.ID, "wsp_1"))
	require.NoError(t, err)
	_, err = sessions.Select(secondSession, expiresAt, userauth.BuildTargetID(login.Server.ID, "wsp_2"))
	require.NoError(t, err)

	router := gin.New()
	internal := router.Group("/api/internal", auth.Middleware())
	RegisterRoutes(internal, Dependencies{Auth: manager, Sessions: sessions})
	logout := performRequest(
		router,
		firstSession,
		http.MethodDelete,
		"/api/internal/onboarding/accounts/"+login.Server.ID,
		nil,
	)
	require.Equal(t, http.StatusNoContent, logout.Code)
	first, _ := sessions.Get(firstSession)
	second, _ := sessions.Get(secondSession)
	require.Empty(t, first.ActiveServerID)
	require.Empty(t, second.ActiveServerID)
}

func performRequest(router http.Handler, sessionToken, method, path string, body any) *httptest.ResponseRecorder {
	var data []byte
	if body != nil {
		data, _ = json.Marshal(body)
	}
	request := httptest.NewRequest(method, path, bytes.NewReader(data))
	if sessionToken != "" {
		request.Header.Set("Authorization", "Bearer "+sessionToken)
	}
	if body != nil {
		request.Header.Set("Content-Type", "application/json")
	}
	recorder := httptest.NewRecorder()
	router.ServeHTTP(recorder, request)
	return recorder
}

func setTestEncryptionKeys(t *testing.T) {
	t.Helper()
	publicKey, privateKey, err := ed25519.GenerateKey(rand.Reader)
	require.NoError(t, err)
	require.NoError(t, crypto.SetKeys(publicKey, privateKey))
	t.Cleanup(crypto.ResetKeys)
}