package control

import (
	"bytes"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"sync"
	"testing"
	"time"

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

	"github.com/tidbcloud/AgentTeam/desktop/internal/daemoncontrol"
	"github.com/tidbcloud/AgentTeam/desktop/internal/token"
	"github.com/tidbcloud/AgentTeam/x/proc"
)

type handlerReceiptStore struct {
	receipt daemoncontrol.Receipt
}

func (s *handlerReceiptStore) RewriteMatching(_, replacement daemoncontrol.Receipt) error {
	s.receipt = replacement
	return nil
}

func TestIdentityRequiresHeaderOnlyControlToken(t *testing.T) {
	router, receipt, _ := newControlRouter(t)

	request := httptest.NewRequest(http.MethodGet, RootPath+"/identity?controlToken="+receipt.ControlToken, nil)
	recorder := httptest.NewRecorder()
	router.ServeHTTP(recorder, request)
	require.Equal(t, http.StatusBadRequest, recorder.Code)

	request = httptest.NewRequest(http.MethodGet, RootPath+"/identity", nil)
	request.Header.Set(ControlTokenHeader, receipt.ControlToken)
	recorder = httptest.NewRecorder()
	router.ServeHTTP(recorder, request)
	require.Equal(t, http.StatusOK, recorder.Code)
	require.NotContains(t, recorder.Body.String(), receipt.ControlToken)
	require.Contains(t, recorder.Body.String(), receipt.InstanceID)
}

func TestShutdownCommitsResponseBeforeSchedulingDrain(t *testing.T) {
	router, receipt, shutdown := newControlRouter(t)
	payload := OwnerProof{
		InstanceID:     receipt.InstanceID,
		OwnerID:        receipt.Owner.ID,
		OwnerPID:       receipt.Owner.PID,
		OwnerStartedAt: receipt.Owner.StartedAt,
		OwnerStartID:   receipt.Owner.StartID,
	}
	request := newJSONRequest(t, http.MethodPost, RootPath+"/shutdown", payload)
	request.Header.Set(ControlTokenHeader, receipt.ControlToken)
	recorder := httptest.NewRecorder()
	router.ServeHTTP(recorder, request)

	require.Equal(t, http.StatusAccepted, recorder.Code)
	require.JSONEq(t, `{"status":"accepted"}`, recorder.Body.String())
	select {
	case <-shutdown:
	case <-time.After(time.Second):
		t.Fatal("shutdown was not scheduled after response")
	}
}

func TestShutdownRejectsBodyTokenWrongOwnerAndServiceOwnership(t *testing.T) {
	router, receipt, _ := newControlRouter(t)
	payload := OwnerProof{
		InstanceID:       receipt.InstanceID,
		OwnerID:          receipt.Owner.ID,
		OwnerPID:         receipt.Owner.PID,
		OwnerStartedAt:   receipt.Owner.StartedAt,
		OwnerStartID:     receipt.Owner.StartID,
		UnexpectedSecret: receipt.ControlToken,
	}
	request := newJSONRequest(t, http.MethodPost, RootPath+"/shutdown", payload)
	request.Header.Set(ControlTokenHeader, receipt.ControlToken)
	recorder := httptest.NewRecorder()
	router.ServeHTTP(recorder, request)
	require.Equal(t, http.StatusBadRequest, recorder.Code)

	payload.UnexpectedSecret = ""
	payload.OwnerID = "wrong-owner"
	request = newJSONRequest(t, http.MethodPost, RootPath+"/shutdown", payload)
	request.Header.Set(ControlTokenHeader, receipt.ControlToken)
	recorder = httptest.NewRecorder()
	router.ServeHTTP(recorder, request)
	require.Equal(t, http.StatusUnauthorized, recorder.Code)

	payload.OwnerID = receipt.Owner.ID
	payload.OwnerStartID = ""
	request = newJSONRequest(t, http.MethodPost, RootPath+"/shutdown", payload)
	request.Header.Set(ControlTokenHeader, receipt.ControlToken)
	recorder = httptest.NewRecorder()
	router.ServeHTTP(recorder, request)
	require.Equal(t, http.StatusUnauthorized, recorder.Code)

	serviceReceipt := receipt
	serviceReceipt.Owner.Kind = daemoncontrol.OwnerKindService
	router, _, _ = newControlRouterWithReceipt(t, serviceReceipt)
	payload.OwnerID = serviceReceipt.Owner.ID
	payload.OwnerStartID = serviceReceipt.Owner.StartID
	request = newJSONRequest(t, http.MethodPost, RootPath+"/shutdown", payload)
	request.Header.Set(ControlTokenHeader, serviceReceipt.ControlToken)
	recorder = httptest.NewRecorder()
	router.ServeHTTP(recorder, request)
	require.Equal(t, http.StatusConflict, recorder.Code)
}

func TestControlRoutesBypassBusinessAuthButUnknownControlRouteDoesNotFallThrough(t *testing.T) {
	router, receipt, _ := newControlRouter(t)
	request := httptest.NewRequest(http.MethodGet, RootPath+"/identity", nil)
	request.Header.Set(ControlTokenHeader, receipt.ControlToken)
	recorder := httptest.NewRecorder()
	router.ServeHTTP(recorder, request)
	require.Equal(t, http.StatusOK, recorder.Code)

	request = httptest.NewRequest(http.MethodGet, RootPath+"/unknown", nil)
	request.Header.Set(ControlTokenHeader, receipt.ControlToken)
	recorder = httptest.NewRecorder()
	router.ServeHTTP(recorder, request)
	require.Equal(t, http.StatusNotFound, recorder.Code)
}

func TestUISessionControlErrorsExposeStableRecoveryCodes(t *testing.T) {
	gin.SetMode(gin.TestMode)
	for _, test := range []struct {
		name   string
		err    error
		status int
		body   string
	}{
		{
			name:   "invalid session",
			err:    token.ErrUISessionInvalid,
			status: http.StatusUnauthorized,
			body:   `{"code":"UI_SESSION_INVALID","error":"UI session is invalid"}`,
		},
		{
			name:   "scope mismatch",
			err:    token.ErrUISessionMismatch,
			status: http.StatusConflict,
			body:   `{"code":"UI_SESSION_SCOPE_MISMATCH","error":"UI session scope mismatch"}`,
		},
	} {
		t.Run(test.name, func(t *testing.T) {
			recorder := httptest.NewRecorder()
			ctx, _ := gin.CreateTestContext(recorder)
			writeControlError(ctx, test.err)
			require.Equal(t, test.status, recorder.Code)
			require.JSONEq(t, test.body, recorder.Body.String())
		})
	}
}

func newControlRouter(t *testing.T) (*gin.Engine, daemoncontrol.Receipt, <-chan struct{}) {
	t.Helper()
	return newControlRouterWithReceipt(t, validControlReceipt())
}

func newControlRouterWithReceipt(
	t *testing.T,
	receipt daemoncontrol.Receipt,
) (*gin.Engine, daemoncontrol.Receipt, <-chan struct{}) {
	t.Helper()
	store := &handlerReceiptStore{receipt: receipt}
	controller, err := daemoncontrol.NewController(daemoncontrol.ControllerOptions{
		Receipt: receipt,
		Store:   store,
		LookupProcess: func(pid int) (proc.ProcessInfo, bool) {
			if pid != receipt.Owner.PID {
				return proc.ProcessInfo{}, false
			}
			return proc.ProcessInfo{
				PID:       pid,
				StartedAt: receipt.Owner.StartedAt,
				StartID:   receipt.Owner.StartID,
			}, true
		},
		PIDAlive: func(int) bool { return false },
	})
	require.NoError(t, err)
	shutdown := make(chan struct{})
	router := gin.New()
	var shutdownOnce sync.Once
	RegisterRoutes(router, Dependencies{
		Controller: controller,
		RequestShutdown: func() {
			shutdownOnce.Do(func() {
				close(shutdown)
			})
		},
	})
	return router, receipt, shutdown
}

func validControlReceipt() daemoncontrol.Receipt {
	startedAt := time.Unix(1710000000, 0).UTC()
	return daemoncontrol.Receipt{
		SchemaVersion:     daemoncontrol.SchemaVersion,
		DesktopAPIVersion: daemoncontrol.DesktopAPIVersion,
		InstanceID:        "instance-control-test",
		PID:               9900,
		StartedAt:         startedAt,
		BaseURL:           "http://127.0.0.1:43123",
		Owner: daemoncontrol.Owner{
			Kind: daemoncontrol.OwnerKindElectron,
			ID:   "owner-control-test",
			ProcessIdentity: daemoncontrol.ProcessIdentity{
				PID:       9901,
				StartedAt: startedAt.Add(-time.Minute),
				StartID:   "linux:boot-test:9901",
			},
		},
		ControlToken: "control-test-token",
	}
}

func newJSONRequest(t *testing.T, method, path string, payload any) *http.Request {
	t.Helper()
	data, err := json.Marshal(payload)
	require.NoError(t, err)
	request := httptest.NewRequest(method, path, bytes.NewReader(data))
	request.Header.Set("Content-Type", "application/json")
	return request
}