package localsecurity

import (
	"net/http"
	"net/http/httptest"
	"net/url"
	"strings"
	"testing"

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

func TestPolicyHostOriginMatrix(t *testing.T) {
	policy, err := New(Options{
		Host:       "127.0.0.1:43123",
		DevOrigins: []string{"http://127.0.0.1:5173"},
	})
	require.NoError(t, err)
	tests := []struct {
		name   string
		host   string
		origin []string
		ok     bool
	}{
		{"originless CLI exact host", "127.0.0.1:43123", nil, true},
		{"Electron app origin", "127.0.0.1:43123", []string{"app://agentteam"}, true},
		{"configured Vite origin", "127.0.0.1:43123", []string{"http://127.0.0.1:5173"}, true},
		{"wrong port", "127.0.0.1:43124", nil, false},
		{"DNS rebinding host", "evil.example:43123", nil, false},
		{"localhost alias", "localhost:43123", nil, false},
		{"multiple origin", "127.0.0.1:43123", []string{"app://agentteam", "http://127.0.0.1:5173"}, false},
		{"reflected origin", "127.0.0.1:43123", []string{"https://evil.example"}, false},
		{"wildcard origin", "127.0.0.1:43123", []string{"*"}, false},
		{"malformed origin", "127.0.0.1:43123", []string{"app://agentteam/path"}, false},
		{"untrusted custom scheme", "127.0.0.1:43123", []string{"file://agentteam"}, false},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			request := &http.Request{
				Method: http.MethodGet,
				Host:   tt.host,
				URL:    &url.URL{Path: "/api/health"},
				Header: make(http.Header),
			}
			for _, origin := range tt.origin {
				request.Header.Add("Origin", origin)
			}
			err := policy.Validate(request)
			if tt.ok {
				require.NoError(t, err)
			} else {
				require.Error(t, err)
			}
		})
	}
}

func TestPolicyRendererCORSContract(t *testing.T) {
	policy, err := New(Options{
		Host:       "127.0.0.1:43123",
		DevOrigins: []string{"http://127.0.0.1:5173"},
	})
	require.NoError(t, err)

	for _, origin := range []string{AppOrigin, "http://127.0.0.1:5173"} {
		t.Run(origin, func(t *testing.T) {
			response := servePolicyRequest(t, policy, http.MethodGet, origin, "")
			require.Equal(t, http.StatusOK, response.Code)
			require.Equal(t, origin, response.Header().Get("Access-Control-Allow-Origin"))
			require.Equal(t, "Origin", response.Header().Get("Vary"))
			require.Equal(
				t,
				"ETag, X-Plan-Draft-ETag, X-Binding-Revision-ETag, X-AgentTeam-Workflow-Agent-Edit-Active-At-Archive, Content-Disposition, X-AgentTeam-Desktop-Auth-Error",
				response.Header().Get("Access-Control-Expose-Headers"),
			)
			require.Empty(t, response.Header().Get("Access-Control-Allow-Credentials"))

			preflight := servePolicyRequest(
				t,
				policy,
				http.MethodOptions,
				origin,
				"authorization,content-type,idempotency-key,if-match,x-agentteam-desktop-onboarding-registration-cutoff",
			)
			require.Equal(t, http.StatusNoContent, preflight.Code)
			require.Equal(t, origin, preflight.Header().Get("Access-Control-Allow-Origin"))
			require.Equal(
				t,
				"GET, HEAD, POST, PUT, PATCH, DELETE, OPTIONS",
				preflight.Header().Get("Access-Control-Allow-Methods"),
			)
			require.Equal(
				t,
				"Authorization, Content-Type, If-Match, Idempotency-Key, X-AgentTeam-Desktop-Onboarding-Registration-Cutoff",
				preflight.Header().Get("Access-Control-Allow-Headers"),
			)
			require.Equal(t, "600", preflight.Header().Get("Access-Control-Max-Age"))
			require.Empty(t, preflight.Header().Get("Access-Control-Allow-Credentials"))
		})
	}
}

func TestPolicyRendererCORSRejectsUntrustedAndNarrowsBrowserHeaders(t *testing.T) {
	policy, err := New(Options{Host: "127.0.0.1:43123"})
	require.NoError(t, err)

	originless := servePolicyRequest(t, policy, http.MethodGet, "", "")
	require.Equal(t, http.StatusOK, originless.Code)
	for name := range originless.Header() {
		require.False(t, strings.HasPrefix(name, "Access-Control-"), name)
	}

	originlessPreflight := servePolicyRequest(t, policy, http.MethodOptions, "", "")
	require.Equal(t, http.StatusForbidden, originlessPreflight.Code)
	for name := range originlessPreflight.Header() {
		require.False(t, strings.HasPrefix(name, "Access-Control-"), name)
	}

	untrusted := servePolicyRequest(t, policy, http.MethodGet, "https://evil.example", "")
	require.Equal(t, http.StatusForbidden, untrusted.Code)
	require.Empty(t, untrusted.Header().Get("Access-Control-Allow-Origin"))

	trusted := servePolicyRequest(t, policy, http.MethodOptions, AppOrigin, "authorization")
	require.Equal(t, http.StatusNoContent, trusted.Code)
	allowedHeaders := trusted.Header().Get("Access-Control-Allow-Headers")
	require.NotContains(t, allowedHeaders, "X-Session-Token")
	require.NotContains(t, allowedHeaders, "X-Api-Key")
	require.NotContains(t, allowedHeaders, "X-AgentTeam-Control-Token")
}

func servePolicyRequest(
	t *testing.T,
	policy *Policy,
	method string,
	origin string,
	requestHeaders string,
) *httptest.ResponseRecorder {
	t.Helper()
	router := gin.New()
	router.Use(policy.Middleware())
	router.Any("/api/test", func(c *gin.Context) {
		c.Status(http.StatusOK)
	})
	request := httptest.NewRequest(method, "/api/test", nil)
	request.Host = "127.0.0.1:43123"
	if origin != "" {
		request.Header.Set("Origin", origin)
	}
	if requestHeaders != "" {
		request.Header.Set("Access-Control-Request-Headers", requestHeaders)
	}
	response := httptest.NewRecorder()
	router.ServeHTTP(response, request)
	return response
}

func TestQuerySessionAllowlistAndDuplicateCredentials(t *testing.T) {
	tests := []struct {
		name   string
		method string
		path   string
		query  string
		header http.Header
		ok     bool
	}{
		{"WS", http.MethodGet, "/socket", "accessToken=one", nil, true},
		{"active stream", http.MethodGet, "/api/internal/servers/active/stream", "accessToken=one", nil, true},
		{"diagnostic archive", http.MethodGet, "/api/internal/machines/mch_1/diagnostics/archive", "accessToken=one", nil, true},
		{"diagnostic archive post", http.MethodPost, "/api/internal/machines/mch_1/diagnostics/archive", "accessToken=one", nil, false},
		{"diagnostic archive collection", http.MethodGet, "/api/internal/machines/mch_1/diagnostics/archive/extra", "accessToken=one", nil, false},
		{"attachment GET", http.MethodGet, "/api/v1/workspaces/w/attachments/a", "accessToken=one", nil, true},
		{"attachment HEAD", http.MethodHead, "/api/v1/workspaces/w/attachments/a", "accessToken=one", nil, true},
		{"managed avatar GET", http.MethodGet, "/api/v1/avatar-assets/a", "accessToken=one", nil, true},
		{"mutation", http.MethodPost, "/api/v1/workspaces/w/attachments/a", "accessToken=one", nil, false},
		{"ordinary GET", http.MethodGet, "/api/v1/workspaces/w/messages", "accessToken=one", nil, false},
		{"legacy token", http.MethodGet, "/socket", "token=one", nil, false},
		{"duplicate query", http.MethodGet, "/socket", "accessToken=one&accessToken=two", nil, false},
		{"duplicate header", http.MethodGet, "/socket", "accessToken=one", http.Header{"Authorization": []string{"Bearer one"}}, false},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			request := &http.Request{
				Method: tt.method,
				URL:    &url.URL{Path: tt.path, RawQuery: tt.query},
				Header: tt.header,
			}
			err := ValidateQueryCredentials(request)
			if tt.ok {
				require.NoError(t, err)
			} else {
				require.Error(t, err)
			}
		})
	}
}

func TestRedactURL(t *testing.T) {
	raw := "http://127.0.0.1:43123/socket?accessToken=secret&safe=value"
	redacted := RedactURL(raw)
	require.NotContains(t, redacted, "secret")
	require.Contains(t, redacted, "safe=value")
	require.Contains(t, redacted, "%5BREDACTED%5D")
}