package daemon

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

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

	"github.com/tidbcloud/AgentTeam/internal/core/secrets"
	"github.com/tidbcloud/AgentTeam/internal/store/agent"
	"github.com/tidbcloud/AgentTeam/internal/versioning"
	"github.com/tidbcloud/AgentTeam/internal/workspace/owner"
)

type daemonRouteAuth struct {
	apiKey                   string
	identity                 agent.MachineIdentity
	issuedToken              string
	authErr                  error
	issueErr                 error
	issueMachine             string
	issueAgent               string
	issueAllowWorkspaceWrite bool
}

func (a *daemonRouteAuth) AuthenticateMachine(_ context.Context, apiKey string) (agent.MachineIdentity, error) {
	if apiKey != a.apiKey {
		return agent.MachineIdentity{}, agent.ErrUnauthorized
	}
	if a.authErr != nil {
		return agent.MachineIdentity{}, a.authErr
	}
	return a.identity, nil
}

func (a *daemonRouteAuth) IssueAgentAuthTokenForMachine(_ context.Context, machineID string, agentID string, allowWorkspaceWrite bool) (string, error) {
	a.issueMachine = machineID
	a.issueAgent = agentID
	a.issueAllowWorkspaceWrite = allowWorkspaceWrite
	if a.issueErr != nil {
		return "", a.issueErr
	}
	return a.issuedToken, nil
}

func TestDaemonAgentAuthTokenRouteIssuesTokenForMachine(t *testing.T) {
	gin.SetMode(gin.TestMode)

	auth := &daemonRouteAuth{
		apiKey:      "machine-key",
		identity:    agent.MachineIdentity{MachineID: "mac_token_refresh", WorkspaceID: "wsp_token_refresh"},
		issuedToken: "fresh-agent-token",
	}
	router := gin.New()
	RegisterRoutes(router, nil, auth, nil, RegisterRoutesOptions{})

	req := httptest.NewRequest(http.MethodPost, "/api/v1/workspaces/wsp_token_refresh/daemon/agents/agt_token_refresh/auth-token", nil)
	req.Header.Set("Authorization", "Bearer machine-key")
	rec := httptest.NewRecorder()
	router.ServeHTTP(rec, req)

	require.Equal(t, http.StatusOK, rec.Code)
	require.Equal(t, "no-store", rec.Header().Get("Cache-Control"))
	require.JSONEq(t, `{"authToken":"fresh-agent-token"}`, rec.Body.String())
	require.Equal(t, "mac_token_refresh", auth.issueMachine)
	require.Equal(t, "agt_token_refresh", auth.issueAgent)
	require.False(t, auth.issueAllowWorkspaceWrite)
}

func TestDaemonAgentAuthTokenWorkspaceScopedRouteIssuesTokenForMachine(t *testing.T) {
	gin.SetMode(gin.TestMode)

	auth := &daemonRouteAuth{
		apiKey:      "machine-key",
		identity:    agent.MachineIdentity{MachineID: "mac_token_refresh", WorkspaceID: "wsp_token_refresh"},
		issuedToken: "fresh-agent-token",
	}
	router := gin.New()
	RegisterRoutes(router, nil, auth, nil, RegisterRoutesOptions{})

	req := httptest.NewRequest(http.MethodPost, "/api/v1/workspaces/wsp_token_refresh/daemon/agents/agt_token_refresh/auth-token", nil)
	req.Header.Set("Authorization", "Bearer machine-key")
	rec := httptest.NewRecorder()
	router.ServeHTTP(rec, req)

	require.Equal(t, http.StatusOK, rec.Code)
	require.Equal(t, "no-store", rec.Header().Get("Cache-Control"))
	require.JSONEq(t, `{"authToken":"fresh-agent-token"}`, rec.Body.String())
	require.Equal(t, "mac_token_refresh", auth.issueMachine)
	require.Equal(t, "agt_token_refresh", auth.issueAgent)
}

func TestDaemonAgentAuthTokenWorkspaceScopedRouteRejectsWorkspaceMismatch(t *testing.T) {
	gin.SetMode(gin.TestMode)

	auth := &daemonRouteAuth{
		apiKey:      "machine-key",
		identity:    agent.MachineIdentity{MachineID: "mac_token_refresh", WorkspaceID: "wsp_token_refresh"},
		issuedToken: "fresh-agent-token",
	}
	router := gin.New()
	RegisterRoutes(router, nil, auth, nil, RegisterRoutesOptions{})

	req := httptest.NewRequest(http.MethodPost, "/api/v1/workspaces/wsp_other/daemon/agents/agt_token_refresh/auth-token", nil)
	req.Header.Set("Authorization", "Bearer machine-key")
	rec := httptest.NewRecorder()
	router.ServeHTTP(rec, req)

	require.Equal(t, http.StatusMisdirectedRequest, rec.Code)
	require.JSONEq(t, `{"code":"workspace_owner_mismatch","error":"Workspace owner mismatch","workspaceId":"wsp_token_refresh"}`, rec.Body.String())
	require.Empty(t, auth.issueAgent)
}

func TestDaemonAgentAuthTokenRoutePassesConnectionProofForCapableDaemon(t *testing.T) {
	gin.SetMode(gin.TestMode)

	auth := &daemonRouteAuth{
		apiKey:      "machine-key",
		identity:    agent.MachineIdentity{MachineID: "mac_token_refresh", WorkspaceID: "wsp_token_refresh"},
		issuedToken: "fresh-agent-token",
	}
	connectionToken := "connection-token"
	hub := &Hub{
		clients: map[string]*clientState{
			"mac_token_refresh": {
				capabilities:          map[string]struct{}{agentPermissionModeCapability: {}},
				tokenRefreshProofHash: secrets.HashToken(connectionToken),
			},
		},
	}
	router := gin.New()
	RegisterRoutes(router, hub, auth, nil, RegisterRoutesOptions{})

	req := httptest.NewRequest(http.MethodPost, "/api/v1/workspaces/wsp_token_refresh/daemon/agents/agt_token_refresh/auth-token", nil)
	req.Header.Set("Authorization", "Bearer machine-key")
	req.Header.Set(daemonTokenRefreshProofHeader, connectionToken)
	rec := httptest.NewRecorder()
	router.ServeHTTP(rec, req)

	require.Equal(t, http.StatusOK, rec.Code)
	require.JSONEq(t, `{"authToken":"fresh-agent-token"}`, rec.Body.String())
	require.Equal(t, "mac_token_refresh", auth.issueMachine)
	require.Equal(t, "agt_token_refresh", auth.issueAgent)
	require.True(t, auth.issueAllowWorkspaceWrite)
}

func TestDaemonAgentAuthTokenRouteIssuesLegacyWorkspaceWriteTokenWithoutConnectionProof(t *testing.T) {
	gin.SetMode(gin.TestMode)

	auth := &daemonRouteAuth{
		apiKey:      "machine-key",
		identity:    agent.MachineIdentity{MachineID: "mac_token_refresh", WorkspaceID: "wsp_token_refresh"},
		issuedToken: "fresh-agent-token",
	}
	// This fake mirrors the post-normalization store path: legacy
	// workspace-write rows return a token without requiring the refresh proof.
	hub := &Hub{
		clients: map[string]*clientState{
			"mac_token_refresh": {
				capabilities:          map[string]struct{}{agentPermissionModeCapability: {}},
				tokenRefreshProofHash: secrets.HashToken("connection-token"),
			},
		},
	}
	router := gin.New()
	RegisterRoutes(router, hub, auth, nil, RegisterRoutesOptions{})

	req := httptest.NewRequest(http.MethodPost, "/api/v1/workspaces/wsp_token_refresh/daemon/agents/agt_token_refresh/auth-token", nil)
	req.Header.Set("Authorization", "Bearer machine-key")
	rec := httptest.NewRecorder()
	router.ServeHTTP(rec, req)

	require.Equal(t, http.StatusOK, rec.Code)
	require.JSONEq(t, `{"authToken":"fresh-agent-token"}`, rec.Body.String())
	require.Equal(t, "mac_token_refresh", auth.issueMachine)
	require.Equal(t, "agt_token_refresh", auth.issueAgent)
	require.False(t, auth.issueAllowWorkspaceWrite)
}

func TestAgentTokenRefreshWorkspaceWriteAllowedRequiresCapabilityAndProof(t *testing.T) {
	hub := &Hub{
		clients: map[string]*clientState{
			"mac_token_refresh": {
				capabilities:          map[string]struct{}{agentPermissionModeCapability: {}},
				tokenRefreshProofHash: secrets.HashToken("connection-token"),
			},
			"mac_no_capability": {
				capabilities:          map[string]struct{}{},
				tokenRefreshProofHash: secrets.HashToken("connection-token"),
			},
		},
	}

	require.True(t, hub.agentTokenRefreshWorkspaceWriteAllowed("mac_token_refresh", "connection-token"))
	require.False(t, hub.agentTokenRefreshWorkspaceWriteAllowed("mac_token_refresh", "wrong-token"))
	require.False(t, hub.agentTokenRefreshWorkspaceWriteAllowed("mac_token_refresh", ""))
	require.False(t, hub.agentTokenRefreshWorkspaceWriteAllowed("mac_no_capability", "connection-token"))
	require.False(t, hub.agentTokenRefreshWorkspaceWriteAllowed("mac_missing", "connection-token"))
}

func TestDaemonAgentAuthTokenRouteRejectsOwnerMismatch(t *testing.T) {
	gin.SetMode(gin.TestMode)

	auth := &daemonRouteAuth{
		apiKey:  "machine-key",
		authErr: owner.ErrOwnerMismatch,
	}
	router := gin.New()
	RegisterRoutes(router, nil, auth, nil, RegisterRoutesOptions{})

	req := httptest.NewRequest(http.MethodPost, "/api/v1/workspaces/wsp_token_refresh/daemon/agents/agt_token_refresh/auth-token", nil)
	req.Header.Set("Authorization", "Bearer machine-key")
	rec := httptest.NewRecorder()
	router.ServeHTTP(rec, req)

	require.Equal(t, http.StatusMisdirectedRequest, rec.Code)
	require.Contains(t, rec.Body.String(), "workspace_owner_mismatch")
	require.Empty(t, auth.issueAgent)
}

func TestDaemonAgentAuthTokenRouteRejectsRouteWorkspaceMismatch(t *testing.T) {
	gin.SetMode(gin.TestMode)

	auth := &daemonRouteAuth{
		apiKey:      "machine-key",
		identity:    agent.MachineIdentity{MachineID: "mac_token_refresh", WorkspaceID: "wsp_token_refresh"},
		issuedToken: "fresh-agent-token",
	}
	router := gin.New()
	RegisterRoutes(router, nil, auth, nil, RegisterRoutesOptions{})

	req := httptest.NewRequest(http.MethodPost, "/api/v1/workspaces/wsp_other/daemon/agents/agt_token_refresh/auth-token", nil)
	req.Header.Set("Authorization", "Bearer machine-key")
	req.Header.Set(headerRouteWorkspaceID, "wsp_other")
	rec := httptest.NewRecorder()
	router.ServeHTTP(rec, req)

	require.Equal(t, http.StatusMisdirectedRequest, rec.Code)
	require.JSONEq(t, `{"code":"workspace_owner_mismatch","error":"Workspace owner mismatch","workspaceId":"wsp_token_refresh"}`, rec.Body.String())
	require.Empty(t, auth.issueAgent)
}

func TestDaemonAgentAuthTokenRouteIncludesWorkspaceIdOnIssueOwnerMismatch(t *testing.T) {
	gin.SetMode(gin.TestMode)

	auth := &daemonRouteAuth{
		apiKey:   "machine-key",
		identity: agent.MachineIdentity{MachineID: "mac_token_refresh", WorkspaceID: "wsp_resolved"},
		issueErr: owner.ErrOwnerMismatch,
	}
	router := gin.New()
	RegisterRoutes(router, nil, auth, nil, RegisterRoutesOptions{})

	req := httptest.NewRequest(http.MethodPost, "/api/v1/workspaces/wsp_resolved/daemon/agents/agt_token_refresh/auth-token", nil)
	req.Header.Set("Authorization", "Bearer machine-key")
	rec := httptest.NewRecorder()
	router.ServeHTTP(rec, req)

	require.Equal(t, http.StatusMisdirectedRequest, rec.Code)
	require.JSONEq(t, `{"code":"workspace_owner_mismatch","error":"Workspace owner mismatch","workspaceId":"wsp_resolved"}`, rec.Body.String())
}

func TestDaemonAgentAuthTokenRouteIncludesWorkspaceIdOnAuthOwnerMismatch(t *testing.T) {
	gin.SetMode(gin.TestMode)

	auth := &daemonRouteAuth{
		apiKey:  "machine-key",
		authErr: owner.OwnerMismatchError{WorkspaceID: "wsp_resolved"},
	}
	router := gin.New()
	RegisterRoutes(router, nil, auth, nil, RegisterRoutesOptions{})

	req := httptest.NewRequest(http.MethodPost, "/api/v1/workspaces/wsp_resolved/daemon/agents/agt_token_refresh/auth-token", nil)
	req.Header.Set("Authorization", "Bearer machine-key")
	rec := httptest.NewRecorder()
	router.ServeHTTP(rec, req)

	require.Equal(t, http.StatusMisdirectedRequest, rec.Code)
	require.JSONEq(t, `{"code":"workspace_owner_mismatch","error":"Workspace owner mismatch","workspaceId":"wsp_resolved"}`, rec.Body.String())
	require.Empty(t, auth.issueAgent)
}

func TestDaemonAgentAuthTokenRouteRejectsInvalidMachineOrAgent(t *testing.T) {
	gin.SetMode(gin.TestMode)

	t.Run("invalid machine key", func(t *testing.T) {
		auth := &daemonRouteAuth{apiKey: "machine-key"}
		router := gin.New()
		RegisterRoutes(router, nil, auth, nil, RegisterRoutesOptions{})

		req := httptest.NewRequest(http.MethodPost, "/api/v1/workspaces/wsp_token_refresh/daemon/agents/agt_token_refresh/auth-token", nil)
		req.Header.Set("Authorization", "Bearer wrong-key")
		rec := httptest.NewRecorder()
		router.ServeHTTP(rec, req)

		require.Equal(t, http.StatusUnauthorized, rec.Code)
		require.Empty(t, auth.issueAgent)
	})

	t.Run("agent not owned by machine", func(t *testing.T) {
		auth := &daemonRouteAuth{
			apiKey:   "machine-key",
			identity: agent.MachineIdentity{MachineID: "mac_token_refresh", WorkspaceID: "wsp_token_refresh"},
			issueErr: agent.ErrUnauthorized,
		}
		router := gin.New()
		RegisterRoutes(router, nil, auth, nil, RegisterRoutesOptions{})

		req := httptest.NewRequest(http.MethodPost, "/api/v1/workspaces/wsp_token_refresh/daemon/agents/agt_other/auth-token?key=machine-key", nil)
		rec := httptest.NewRecorder()
		router.ServeHTTP(rec, req)

		require.Equal(t, http.StatusUnauthorized, rec.Code)
		require.Equal(t, "mac_token_refresh", auth.issueMachine)
		require.Equal(t, "agt_other", auth.issueAgent)
	})

	t.Run("store failure", func(t *testing.T) {
		auth := &daemonRouteAuth{
			apiKey:   "machine-key",
			identity: agent.MachineIdentity{MachineID: "mac_token_refresh", WorkspaceID: "wsp_token_refresh"},
			issueErr: errors.New("database unavailable"),
		}
		router := gin.New()
		RegisterRoutes(router, nil, auth, nil, RegisterRoutesOptions{})

		req := httptest.NewRequest(http.MethodPost, "/api/v1/workspaces/wsp_token_refresh/daemon/agents/agt_other/auth-token?key=machine-key", nil)
		rec := httptest.NewRecorder()
		router.ServeHTTP(rec, req)

		require.Equal(t, http.StatusInternalServerError, rec.Code)
	})
}

func TestDaemonConnectRouteOwnerMismatchIncludesWorkspaceId(t *testing.T) {
	gin.SetMode(gin.TestMode)

	auth := &daemonRouteAuth{
		apiKey:  "machine-key",
		authErr: owner.OwnerMismatchError{WorkspaceID: "wsp_viyi2993wddldo7my6d1"},
	}
	router := gin.New()
	RegisterRoutes(router, nil, auth, nil, RegisterRoutesOptions{})

	req := httptest.NewRequest(http.MethodGet, "/api/v1/workspaces/wsp_viyi2993wddldo7my6d1/daemon/connect?key=machine-key", nil)
	rec := httptest.NewRecorder()
	router.ServeHTTP(rec, req)

	require.Equal(t, http.StatusMisdirectedRequest, rec.Code)
	require.JSONEq(
		t,
		`{"code":"workspace_owner_mismatch","error":"Workspace owner mismatch","workspaceId":"wsp_viyi2993wddldo7my6d1"}`,
		rec.Body.String(),
	)
}

func TestDaemonConnectRouteOwnerMismatchOmitsWorkspaceIdWhenUnknown(t *testing.T) {
	gin.SetMode(gin.TestMode)

	auth := &daemonRouteAuth{
		apiKey:  "machine-key",
		authErr: owner.ErrOwnerMismatch,
	}
	router := gin.New()
	RegisterRoutes(router, nil, auth, nil, RegisterRoutesOptions{})

	req := httptest.NewRequest(http.MethodGet, "/api/v1/workspaces/wsp_unknown/daemon/connect?key=machine-key", nil)
	rec := httptest.NewRecorder()
	router.ServeHTTP(rec, req)

	require.Equal(t, http.StatusMisdirectedRequest, rec.Code)
	require.JSONEq(
		t,
		`{"code":"workspace_owner_mismatch","error":"Workspace owner mismatch"}`,
		rec.Body.String(),
	)
}

func TestDaemonConnectRouteClassifiesMachineAuthFailures(t *testing.T) {
	gin.SetMode(gin.TestMode)

	tests := []struct {
		name       string
		authErr    error
		statusCode int
		body       string
	}{
		{
			name:       "invalid key",
			authErr:    agent.ErrUnauthorized,
			statusCode: http.StatusUnauthorized,
			body:       `{"error":"invalid machine api key"}`,
		},
		{
			name:       "authentication store unavailable",
			authErr:    errors.New("database unavailable"),
			statusCode: http.StatusServiceUnavailable,
			body:       `{"error":"machine authentication unavailable"}`,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			auth := &daemonRouteAuth{apiKey: "machine-key", authErr: tt.authErr}
			router := gin.New()
			RegisterRoutes(router, nil, auth, nil, RegisterRoutesOptions{})

			req := httptest.NewRequest(http.MethodGet, "/api/v1/workspaces/wsp_machine/daemon/connect?key=machine-key", nil)
			rec := httptest.NewRecorder()
			router.ServeHTTP(rec, req)

			require.Equal(t, tt.statusCode, rec.Code)
			require.JSONEq(t, tt.body, rec.Body.String())
		})
	}
}

func TestDaemonConnectRouteRouteWorkspaceMismatchIncludesAuthenticatedWorkspaceId(t *testing.T) {
	gin.SetMode(gin.TestMode)

	auth := &daemonRouteAuth{
		apiKey:   "machine-key",
		identity: agent.MachineIdentity{MachineID: "mac_token_refresh", WorkspaceID: "wsp_viyi2993wddldo7my6d1"},
	}
	router := gin.New()
	RegisterRoutes(router, nil, auth, nil, RegisterRoutesOptions{})

	req := httptest.NewRequest(http.MethodGet, "/api/v1/workspaces/wsp_other/daemon/connect?key=machine-key", nil)
	req.Header.Set(headerRouteWorkspaceID, "wsp_other")
	rec := httptest.NewRecorder()
	router.ServeHTTP(rec, req)

	require.Equal(t, http.StatusMisdirectedRequest, rec.Code)
	require.JSONEq(
		t,
		`{"code":"workspace_owner_mismatch","error":"Workspace owner mismatch","workspaceId":"wsp_viyi2993wddldo7my6d1"}`,
		rec.Body.String(),
	)
}

func TestDaemonUpgradeInfoRouteReturnsServerVersionAndRecommendation(t *testing.T) {
	gin.SetMode(gin.TestMode)

	auth := &daemonRouteAuth{
		apiKey:   "machine-key",
		identity: agent.MachineIdentity{MachineID: "mac_upgrade", WorkspaceID: "wsp_upgrade"},
	}
	router := gin.New()
	RegisterRoutes(router, nil, auth, nil, RegisterRoutesOptions{
		LatestDesktopVersion: versioning.NewStaticProvider("v0.2.5"),
		MinDesktopVersion:    "v0.2.0",
	})

	req := httptest.NewRequest(http.MethodGet, "/api/v1/workspaces/wsp_upgrade/daemon/upgrade-info", nil)
	req.Header.Set("Authorization", "Bearer machine-key")
	rec := httptest.NewRecorder()
	router.ServeHTTP(rec, req)

	require.Equal(t, http.StatusOK, rec.Code)
	body := rec.Body.String()
	require.Contains(t, body, `"serverVersion"`)
	require.Contains(t, body, `"recommendedDesktopVersion":"0.2.5"`)
	require.Contains(t, body, `"minDesktopVersion":"v0.2.0"`)
}

func TestDaemonUpgradeInfoRouteOmitsEmptyFields(t *testing.T) {
	gin.SetMode(gin.TestMode)

	auth := &daemonRouteAuth{
		apiKey:   "machine-key",
		identity: agent.MachineIdentity{MachineID: "mac_upgrade", WorkspaceID: "wsp_upgrade"},
	}
	router := gin.New()
	RegisterRoutes(router, nil, auth, nil, RegisterRoutesOptions{})

	req := httptest.NewRequest(http.MethodGet, "/api/v1/workspaces/wsp_upgrade/daemon/upgrade-info", nil)
	req.Header.Set("Authorization", "Bearer machine-key")
	rec := httptest.NewRecorder()
	router.ServeHTTP(rec, req)

	require.Equal(t, http.StatusOK, rec.Code)
	require.Contains(t, rec.Body.String(), `"serverVersion"`)
	require.NotContains(t, rec.Body.String(), `"recommendedDesktopVersion"`)
	require.NotContains(t, rec.Body.String(), `"minDesktopVersion"`)
}

func TestDaemonUpgradeInfoRouteRejectsMissingAPIKey(t *testing.T) {
	gin.SetMode(gin.TestMode)

	auth := &daemonRouteAuth{apiKey: "machine-key"}
	router := gin.New()
	RegisterRoutes(router, nil, auth, nil, RegisterRoutesOptions{})

	req := httptest.NewRequest(http.MethodGet, "/api/v1/workspaces/wsp_upgrade/daemon/upgrade-info", nil)
	rec := httptest.NewRecorder()
	router.ServeHTTP(rec, req)

	require.Equal(t, http.StatusUnauthorized, rec.Code)
}

func TestDaemonUpgradeInfoRouteRejectsWorkspaceMismatch(t *testing.T) {
	gin.SetMode(gin.TestMode)

	auth := &daemonRouteAuth{
		apiKey:   "machine-key",
		identity: agent.MachineIdentity{MachineID: "mac_upgrade", WorkspaceID: "wsp_upgrade"},
	}
	router := gin.New()
	RegisterRoutes(router, nil, auth, nil, RegisterRoutesOptions{})

	req := httptest.NewRequest(http.MethodGet, "/api/v1/workspaces/wsp_other/daemon/upgrade-info", nil)
	req.Header.Set("Authorization", "Bearer machine-key")
	rec := httptest.NewRecorder()
	router.ServeHTTP(rec, req)

	require.Equal(t, http.StatusMisdirectedRequest, rec.Code)
}