package middleware

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

	"github.com/gin-gonic/gin"
)

// setupAgentAllowlistRouter 注册白名单内代表性路由 + 白名单外路由,
// 用请求头 X-Test-Auth-Method 模拟认证中间件写入的 auth_method。
func setupAgentAllowlistRouter(reached *map[string]bool) *gin.Engine {
	gin.SetMode(gin.TestMode)
	r := gin.New()
	r.Use(func(c *gin.Context) {
		if m := c.GetHeader("X-Test-Auth-Method"); m != "" {
			c.Set("auth_method", m)
		}
		c.Next()
	})
	r.Use(AgentRouteAllowlist())

	mark := func(key string) gin.HandlerFunc {
		return func(c *gin.Context) {
			(*reached)[key] = true
			c.Status(http.StatusOK)
		}
	}

	// 白名单内 GET
	r.GET("/api/v1/agent/whoami", mark("agent-whoami"))
	r.GET("/api/v1/caasm/assets", mark("caasm-assets"))
	r.GET("/api/v1/caasm/assets/:id", mark("caasm-asset-detail"))
	r.GET("/api/v1/vulnerabilities", mark("vuln-list"))
	r.GET("/api/v1/vulnerabilities/:public_id", mark("vuln-detail"))
	r.GET("/api/v1/vulnerabilities/:public_id/comments", mark("vuln-comments-list"))
	r.GET("/api/v1/openapi.json", mark("openapi-json"))
	// 白名单内写路由(iteration-002 P1-T4 新增 6 条)
	r.POST("/api/v1/vulnerabilities", mark("vuln-create"))
	r.PUT("/api/v1/vulnerabilities/:public_id", mark("vuln-update"))
	r.POST("/api/v1/vulnerabilities/:public_id/transition", mark("vuln-transition"))
	r.POST("/api/v1/vulnerabilities/:public_id/comments", mark("vuln-comment-add"))
	r.POST("/api/v1/vulnerabilities/uploads", mark("vuln-upload"))
	r.POST("/api/v1/vulnerabilities/uploads/confirm", mark("vuln-upload-confirm"))
	// 路径在 GET 白名单内但 method 未放行
	r.POST("/api/v1/caasm/assets", mark("caasm-assets-post"))
	// 白名单外
	r.POST("/api/v1/caasm/ingest", mark("caasm-ingest"))
	r.GET("/api/v1/users", mark("users-list"))
	r.DELETE("/api/v1/caasm/assets/:id", mark("caasm-asset-delete"))
	r.DELETE("/api/v1/vulnerabilities/:public_id", mark("vuln-delete"))
	return r
}

func doRequest(r *gin.Engine, method, path, authMethod string) *httptest.ResponseRecorder {
	req := httptest.NewRequest(method, path, nil)
	if authMethod != "" {
		req.Header.Set("X-Test-Auth-Method", authMethod)
	}
	w := httptest.NewRecorder()
	r.ServeHTTP(w, req)
	return w
}

func TestAgentRouteAllowlist(t *testing.T) {
	t.Run("agent + whitelisted GET passes", func(t *testing.T) {
		cases := []struct {
			path    string
			handler string
		}{
			{"/api/v1/agent/whoami", "agent-whoami"},
			{"/api/v1/caasm/assets", "caasm-assets"},
			{"/api/v1/caasm/assets/123", "caasm-asset-detail"},
			{"/api/v1/vulnerabilities", "vuln-list"},
			{"/api/v1/vulnerabilities/VULN-2026-0001", "vuln-detail"},
			{"/api/v1/vulnerabilities/VULN-2026-0001/comments", "vuln-comments-list"},
			{"/api/v1/openapi.json", "openapi-json"},
		}
		for _, tc := range cases {
			reached := map[string]bool{}
			r := setupAgentAllowlistRouter(&reached)
			w := doRequest(r, http.MethodGet, tc.path, "agent")
			if !reached[tc.handler] {
				t.Fatalf("GET %s: expected handler %q to be reached, body=%s", tc.path, tc.handler, w.Body.String())
			}
		}
	})

	t.Run("agent + whitelisted write routes pass", func(t *testing.T) {
		cases := []struct {
			method  string
			path    string
			handler string
		}{
			{http.MethodPost, "/api/v1/vulnerabilities", "vuln-create"},
			{http.MethodPut, "/api/v1/vulnerabilities/VULN-2026-0001", "vuln-update"},
			{http.MethodPost, "/api/v1/vulnerabilities/VULN-2026-0001/transition", "vuln-transition"},
			{http.MethodPost, "/api/v1/vulnerabilities/VULN-2026-0001/comments", "vuln-comment-add"},
			{http.MethodPost, "/api/v1/vulnerabilities/uploads", "vuln-upload"},
			{http.MethodPost, "/api/v1/vulnerabilities/uploads/confirm", "vuln-upload-confirm"},
		}
		for _, tc := range cases {
			reached := map[string]bool{}
			r := setupAgentAllowlistRouter(&reached)
			w := doRequest(r, tc.method, tc.path, "agent")
			if !reached[tc.handler] {
				t.Fatalf("%s %s: expected handler %q to be reached, body=%s", tc.method, tc.path, tc.handler, w.Body.String())
			}
		}
	})

	t.Run("agent + POST on GET-only whitelisted path forbidden", func(t *testing.T) {
		reached := map[string]bool{}
		r := setupAgentAllowlistRouter(&reached)
		w := doRequest(r, http.MethodPost, "/api/v1/caasm/assets", "agent")
		if reached["caasm-assets-post"] {
			t.Fatal("POST /api/v1/caasm/assets: handler must not be reached for agent (GET-only entry)")
		}
		if !strings.Contains(w.Body.String(), `"code":403`) {
			t.Fatalf("expected body code 403, got %s", w.Body.String())
		}
	})

	t.Run("agent + POST ingest forbidden", func(t *testing.T) {
		reached := map[string]bool{}
		r := setupAgentAllowlistRouter(&reached)
		w := doRequest(r, http.MethodPost, "/api/v1/caasm/ingest", "agent")
		if reached["caasm-ingest"] {
			t.Fatal("POST /api/v1/caasm/ingest: handler must not be reached for agent")
		}
		if !strings.Contains(w.Body.String(), `"code":403`) {
			t.Fatalf("expected body code 403, got %s", w.Body.String())
		}
	})

	t.Run("agent + non-whitelisted GET forbidden", func(t *testing.T) {
		reached := map[string]bool{}
		r := setupAgentAllowlistRouter(&reached)
		w := doRequest(r, http.MethodGet, "/api/v1/users", "agent")
		if reached["users-list"] {
			t.Fatal("GET /api/v1/users: handler must not be reached for agent")
		}
		if !strings.Contains(w.Body.String(), `"code":403`) {
			t.Fatalf("expected body code 403, got %s", w.Body.String())
		}
	})

	t.Run("agent + DELETE forbidden", func(t *testing.T) {
		cases := []struct {
			path    string
			handler string
		}{
			{"/api/v1/caasm/assets/123", "caasm-asset-delete"},
			{"/api/v1/vulnerabilities/VULN-2026-0001", "vuln-delete"},
		}
		for _, tc := range cases {
			reached := map[string]bool{}
			r := setupAgentAllowlistRouter(&reached)
			w := doRequest(r, http.MethodDelete, tc.path, "agent")
			if reached[tc.handler] {
				t.Fatalf("DELETE %s: handler must not be reached for agent", tc.path)
			}
			if !strings.Contains(w.Body.String(), `"code":403`) {
				t.Fatalf("DELETE %s: expected body code 403, got %s", tc.path, w.Body.String())
			}
		}
	})

	t.Run("human (no auth_method) passes through", func(t *testing.T) {
		cases := []struct {
			method  string
			path    string
			handler string
		}{
			{http.MethodGet, "/api/v1/caasm/assets", "caasm-assets"},
			{http.MethodPost, "/api/v1/caasm/assets", "caasm-assets-post"},
			{http.MethodPost, "/api/v1/caasm/ingest", "caasm-ingest"},
			{http.MethodGet, "/api/v1/users", "users-list"},
			{http.MethodPut, "/api/v1/vulnerabilities/VULN-2026-0001", "vuln-update"},
			{http.MethodDelete, "/api/v1/caasm/assets/123", "caasm-asset-delete"},
			{http.MethodDelete, "/api/v1/vulnerabilities/VULN-2026-0001", "vuln-delete"},
		}
		for _, tc := range cases {
			reached := map[string]bool{}
			r := setupAgentAllowlistRouter(&reached)
			w := doRequest(r, tc.method, tc.path, "")
			if !reached[tc.handler] {
				t.Fatalf("%s %s: expected handler %q to be reached for human caller, body=%s", tc.method, tc.path, tc.handler, w.Body.String())
			}
		}
	})

	t.Run("allowlist integrity: 37 GET + 6 write entries, no DELETE, no claim", func(t *testing.T) {
		if got := len(agentAllowedRoutes); got != 43 {
			t.Fatalf("expected exactly 43 allowlist entries (37 GET + 6 write), got %d", got)
		}
		writeRoutes := map[string]bool{
			"POST /api/v1/vulnerabilities":                       false,
			"PUT /api/v1/vulnerabilities/:public_id":             false,
			"POST /api/v1/vulnerabilities/:public_id/transition": false,
			"POST /api/v1/vulnerabilities/:public_id/comments":   false,
			"POST /api/v1/vulnerabilities/uploads":               false,
			"POST /api/v1/vulnerabilities/uploads/confirm":       false,
		}
		getCount := 0
		for key := range agentAllowedRoutes {
			parts := strings.SplitN(key, " ", 2)
			if len(parts) != 2 || !strings.HasPrefix(parts[1], "/api/v1/") {
				t.Fatalf("allowlist entry %q not in \"METHOD /api/v1/...\" form", key)
			}
			switch parts[0] {
			case http.MethodGet:
				getCount++
			case http.MethodPost, http.MethodPut:
				if _, ok := writeRoutes[key]; !ok {
					t.Fatalf("unexpected write entry %q in allowlist", key)
				}
				writeRoutes[key] = true
			default:
				t.Fatalf("unexpected method %q in allowlist entry %q (DELETE/PATCH must never appear)", parts[0], key)
			}
			if strings.Contains(key, "/claim") {
				t.Fatalf("claim route must not be opened to agent: %q", key)
			}
		}
		if getCount != 37 {
			t.Fatalf("expected 37 GET entries, got %d", getCount)
		}
		for route, seen := range writeRoutes {
			if !seen {
				t.Fatalf("expected write route %q missing from allowlist", route)
			}
		}
	})
}