package tools

import (
	"context"
	"encoding/json"
	"errors"
	"net/http"
	"net/http/httptest"
	"net/url"
	"strings"
	"sync"
	"testing"

	"security-center-mcp/internal/client"
)

// vulnRecordedRequest captures one request seen by the mock backend.
type vulnRecordedRequest struct {
	Path  string
	Query url.Values
}

// newVulnMockBackend starts an httptest server wrapping handler, records
// every request, and returns a client pointed at it plus an accessor for
// the recorded requests.
func newVulnMockBackend(t *testing.T, handler http.HandlerFunc) (*client.Client, func() []vulnRecordedRequest) {
	t.Helper()
	var mu sync.Mutex
	var reqs []vulnRecordedRequest
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		mu.Lock()
		reqs = append(reqs, vulnRecordedRequest{Path: r.URL.Path, Query: r.URL.Query()})
		mu.Unlock()
		handler(w, r)
	}))
	t.Cleanup(srv.Close)
	c := &client.Client{BaseURL: srv.URL, APIKey: "test-key", HTTP: srv.Client()}
	return c, func() []vulnRecordedRequest {
		mu.Lock()
		defer mu.Unlock()
		return append([]vulnRecordedRequest(nil), reqs...)
	}
}

// writeVulnEnvelope writes a backend-style {code, message, data} envelope.
func writeVulnEnvelope(t *testing.T, w http.ResponseWriter, code int, message string, data any) {
	t.Helper()
	w.Header().Set("Content-Type", "application/json")
	if err := json.NewEncoder(w).Encode(map[string]any{
		"code":      code,
		"message":   message,
		"data":      data,
		"timestamp": 1700000000,
	}); err != nil {
		t.Errorf("write envelope: %v", err)
	}
}

// vulnErrPayload asserts data is a structured {"error": {...}} payload and
// returns the inner error object.
func vulnErrPayload(t *testing.T, data any) map[string]any {
	t.Helper()
	m, ok := data.(map[string]any)
	if !ok {
		t.Fatalf("expected error payload map, got %T (%v)", data, data)
	}
	e, ok := m["error"].(map[string]any)
	if !ok {
		t.Fatalf("expected data.error object, got %v", m)
	}
	return e
}

func TestDispatchVulnHandledSemantics(t *testing.T) {
	c, reqs := newVulnMockBackend(t, func(w http.ResponseWriter, r *http.Request) {
		writeVulnEnvelope(t, w, 0, "ok", map[string]any{})
	})

	// Vuln tool names are handled (args cover the required params so each
	// call runs to completion).
	for _, name := range []string{
		"search_vulnerabilities",
		"get_vulnerability",
		"list_vuln_lookup_options",
		"get_vuln_tags_stats",
	} {
		args := map[string]any{"public_id": "VULN-1", "kind": "tags"}
		_, handled, err := dispatchVuln(context.Background(), c, name, args)
		if err != nil {
			t.Fatalf("%s: unexpected err: %v", name, err)
		}
		if !handled {
			t.Fatalf("%s: expected handled=true", name)
		}
	}

	// Any other name is not handled and must not hit the backend.
	for _, name := range []string{"search_assets", "get_asset", "no_such_tool"} {
		data, handled, err := dispatchVuln(context.Background(), c, name, nil)
		if handled || data != nil || err != nil {
			t.Fatalf("%s: expected (nil,false,nil), got (%v,%v,%v)", name, data, handled, err)
		}
	}

	// search(1) + get_vulnerability(3) + lookup(1) + stats(1) = 6 requests.
	if got := len(reqs()); got != 6 {
		t.Fatalf("expected 6 backend requests, got %d", got)
	}
}

func TestSearchVulnerabilitiesSuccess(t *testing.T) {
	c, reqs := newVulnMockBackend(t, func(w http.ResponseWriter, r *http.Request) {
		writeVulnEnvelope(t, w, 0, "ok", map[string]any{
			"total": 1,
			"items": []any{map[string]any{"public_id": "VULN-1"}},
		})
	})

	data, handled, err := dispatchVuln(context.Background(), c, "search_vulnerabilities", map[string]any{
		"status":   []any{"open", "investigating"},
		"severity": []any{"critical", "high"},
		"tags":     "web,sql",
		"search":   "rce",
		"page":     2,
		"has_poc":  true,
	})
	if err != nil || !handled {
		t.Fatalf("expected (data,true,nil), got handled=%v err=%v", handled, err)
	}

	got := reqs()
	if len(got) != 1 {
		t.Fatalf("expected 1 request, got %d", len(got))
	}
	if got[0].Path != "/api/v1/vulnerabilities" {
		t.Fatalf("unexpected path %q", got[0].Path)
	}
	q := got[0].Query
	want := map[string]string{
		"status":   "open,investigating", // 数组 → 逗号连接
		"severity": "critical,high",
		"tags":     "web,sql", // 字符串原样透传
		"search":   "rce",
		"page":     "2",
		"has_poc":  "true",
	}
	for k, v := range want {
		if q.Get(k) != v {
			t.Errorf("query %s: want %q, got %q", k, v, q.Get(k))
		}
	}
	// Args not supplied must not appear in the query.
	for _, k := range []string{"application", "sort_by", "page_size", "is_overdue"} {
		if _, ok := q[k]; ok {
			t.Errorf("unexpected query param %s", k)
		}
	}

	m, ok := data.(map[string]any)
	if !ok || m["total"] != float64(1) {
		t.Fatalf("unexpected data %v", data)
	}
}

func TestGetVulnerabilitySuccess(t *testing.T) {
	c, reqs := newVulnMockBackend(t, func(w http.ResponseWriter, r *http.Request) {
		switch {
		case strings.HasSuffix(r.URL.Path, "/events"):
			writeVulnEnvelope(t, w, 0, "ok", []any{map[string]any{"type": "status_change"}})
		case strings.HasSuffix(r.URL.Path, "/comments"):
			writeVulnEnvelope(t, w, 0, "ok", []any{map[string]any{"body": "确认中"}})
		default:
			writeVulnEnvelope(t, w, 0, "ok", map[string]any{"public_id": "VULN-2024-0001", "severity": "high"})
		}
	})

	data, handled, err := dispatchVuln(context.Background(), c, "get_vulnerability", map[string]any{
		"public_id": "VULN-2024-0001",
	})
	if err != nil || !handled {
		t.Fatalf("expected (data,true,nil), got handled=%v err=%v", handled, err)
	}

	got := reqs()
	if len(got) != 3 {
		t.Fatalf("expected 3 requests, got %d", len(got))
	}
	wantPaths := []string{
		"/api/v1/vulnerabilities/VULN-2024-0001",
		"/api/v1/vulnerabilities/VULN-2024-0001/events",
		"/api/v1/vulnerabilities/VULN-2024-0001/comments",
	}
	for i, p := range wantPaths {
		if got[i].Path != p {
			t.Errorf("request #%d: want path %q, got %q", i, p, got[i].Path)
		}
	}

	m, ok := data.(map[string]any)
	if !ok {
		t.Fatalf("expected merged map, got %T", data)
	}
	vuln, ok := m["vulnerability"].(map[string]any)
	if !ok || vuln["severity"] != "high" {
		t.Errorf("unexpected vulnerability field: %v", m["vulnerability"])
	}
	if _, ok := m["events"].([]any); !ok {
		t.Errorf("unexpected events field: %v", m["events"])
	}
	if _, ok := m["comments"].([]any); !ok {
		t.Errorf("unexpected comments field: %v", m["comments"])
	}
}

func TestGetVulnerabilityMissingPublicID(t *testing.T) {
	c, reqs := newVulnMockBackend(t, func(w http.ResponseWriter, r *http.Request) {
		t.Errorf("backend must not be called without public_id: %s", r.URL.Path)
	})

	for _, args := range []map[string]any{{}, {"public_id": ""}} {
		data, handled, err := dispatchVuln(context.Background(), c, "get_vulnerability", args)
		if data != nil || !handled {
			t.Fatalf("args %v: expected (nil,true,ArgError), got data=%v handled=%v", args, data, handled)
		}
		var argErr *ArgError
		if !errors.As(err, &argErr) {
			t.Fatalf("args %v: expected *ArgError, got %T (%v)", args, err, err)
		}
		if !strings.Contains(argErr.Msg, "public_id") {
			t.Errorf("args %v: message should mention public_id, got %q", args, argErr.Msg)
		}
	}
	if got := len(reqs()); got != 0 {
		t.Fatalf("expected no backend requests, got %d", got)
	}
}

func TestGetVulnerabilityDetailNotFound(t *testing.T) {
	c, reqs := newVulnMockBackend(t, func(w http.ResponseWriter, r *http.Request) {
		if strings.HasSuffix(r.URL.Path, "/events") || strings.HasSuffix(r.URL.Path, "/comments") {
			t.Errorf("events/comments must not be called after detail 404: %s", r.URL.Path)
		}
		writeVulnEnvelope(t, w, 404, "vulnerability not found", nil)
	})

	data, handled, err := dispatchVuln(context.Background(), c, "get_vulnerability", map[string]any{
		"public_id": "NOPE",
	})
	if err != nil || !handled {
		t.Fatalf("expected (error-payload,true,nil), got handled=%v err=%v", handled, err)
	}
	e := vulnErrPayload(t, data)
	if e["code"] != 404 {
		t.Errorf("want code 404, got %v", e["code"])
	}
	if e["message"] != "vulnerability not found" {
		t.Errorf("want passthrough message, got %v", e["message"])
	}
	if got := len(reqs()); got != 1 {
		t.Fatalf("expected exactly 1 request (detail only), got %d", got)
	}
}

func TestListVulnLookupOptionsKinds(t *testing.T) {
	cases := []struct {
		kind string
		path string
	}{
		{"tags", "/api/v1/vulnerability-tags"},
		{"types", "/api/v1/vulnerability-types"},
		{"sources", "/api/v1/vulnerability-sources"},
		{"channels", "/api/v1/vulnerability-discovery-channels"},
	}
	for _, tc := range cases {
		t.Run(tc.kind, func(t *testing.T) {
			c, reqs := newVulnMockBackend(t, func(w http.ResponseWriter, r *http.Request) {
				writeVulnEnvelope(t, w, 0, "ok", map[string]any{"items": []any{}})
			})
			data, handled, err := dispatchVuln(context.Background(), c, "list_vuln_lookup_options", map[string]any{
				"kind": tc.kind,
				"q":    "x",
			})
			if err != nil || !handled {
				t.Fatalf("expected (data,true,nil), got handled=%v err=%v", handled, err)
			}
			got := reqs()
			if len(got) != 1 {
				t.Fatalf("expected 1 request, got %d", len(got))
			}
			if got[0].Path != tc.path {
				t.Errorf("want path %q, got %q", tc.path, got[0].Path)
			}
			if got[0].Query.Get("q") != "x" {
				t.Errorf("want q=x, got %q", got[0].Query.Get("q"))
			}
			if _, ok := data.(map[string]any); !ok {
				t.Errorf("unexpected data %v", data)
			}
		})
	}
}

func TestListVulnLookupOptionsTypeFilters(t *testing.T) {
	c, reqs := newVulnMockBackend(t, func(w http.ResponseWriter, r *http.Request) {
		writeVulnEnvelope(t, w, 0, "ok", map[string]any{})
	})

	// kind=types declares severity/category/system_only.
	if _, _, err := dispatchVuln(context.Background(), c, "list_vuln_lookup_options", map[string]any{
		"kind": "types", "severity": "high", "category": "web", "system_only": true,
	}); err != nil {
		t.Fatalf("unexpected err: %v", err)
	}
	q := reqs()[0].Query
	for k, v := range map[string]string{"severity": "high", "category": "web", "system_only": "true"} {
		if q.Get(k) != v {
			t.Errorf("query %s: want %q, got %q", k, v, q.Get(k))
		}
	}
}

func TestListVulnLookupOptionsDispatchArgsNotMapped(t *testing.T) {
	c, reqs := newVulnMockBackend(t, func(w http.ResponseWriter, r *http.Request) {
		writeVulnEnvelope(t, w, 0, "ok", map[string]any{})
	})

	// kind is dispatch-only and the tags endpoint declares no severity:
	// neither may reach the backend query.
	if _, _, err := dispatchVuln(context.Background(), c, "list_vuln_lookup_options", map[string]any{
		"kind": "tags", "severity": "high",
	}); err != nil {
		t.Fatalf("unexpected err: %v", err)
	}
	q := reqs()[0].Query
	for _, k := range []string{"kind", "severity"} {
		if _, ok := q[k]; ok {
			t.Errorf("param %s must not be sent to the backend", k)
		}
	}
}

func TestListVulnLookupOptionsInvalidKind(t *testing.T) {
	c, reqs := newVulnMockBackend(t, func(w http.ResponseWriter, r *http.Request) {
		t.Errorf("backend must not be called with invalid kind: %s", r.URL.Path)
	})

	for _, args := range []map[string]any{{}, {"kind": "foo"}} {
		data, handled, err := dispatchVuln(context.Background(), c, "list_vuln_lookup_options", args)
		if data != nil || !handled {
			t.Fatalf("args %v: expected (nil,true,ArgError), got data=%v handled=%v", args, data, handled)
		}
		var argErr *ArgError
		if !errors.As(err, &argErr) {
			t.Fatalf("args %v: expected *ArgError, got %T (%v)", args, err, err)
		}
		for _, legal := range []string{"tags", "types", "sources", "channels"} {
			if !strings.Contains(argErr.Msg, legal) {
				t.Errorf("args %v: message should list legal value %q, got %q", args, legal, argErr.Msg)
			}
		}
		if args["kind"] == "foo" && !strings.Contains(argErr.Msg, "foo") {
			t.Errorf("message should echo the illegal value, got %q", argErr.Msg)
		}
	}
	if got := len(reqs()); got != 0 {
		t.Fatalf("expected no backend requests, got %d", got)
	}
}

func TestGetVulnTagsStatsSuccess(t *testing.T) {
	c, reqs := newVulnMockBackend(t, func(w http.ResponseWriter, r *http.Request) {
		writeVulnEnvelope(t, w, 0, "ok", []any{map[string]any{"tag": "web", "count": 3}})
	})

	data, handled, err := dispatchVuln(context.Background(), c, "get_vuln_tags_stats", nil)
	if err != nil || !handled {
		t.Fatalf("expected (data,true,nil), got handled=%v err=%v", handled, err)
	}
	got := reqs()
	if len(got) != 1 {
		t.Fatalf("expected 1 request, got %d", len(got))
	}
	if got[0].Path != "/api/v1/vulnerability-tags/stats" {
		t.Fatalf("unexpected path %q", got[0].Path)
	}
	if len(got[0].Query) != 0 {
		t.Fatalf("expected no query params, got %v", got[0].Query)
	}
	items, ok := data.([]any)
	if !ok || len(items) != 1 {
		t.Fatalf("unexpected data %v", data)
	}
}

func TestVulnAPIErrorMapping(t *testing.T) {
	c, _ := newVulnMockBackend(t, func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Retry-After", "30")
		writeVulnEnvelope(t, w, 429, "rate limit exceeded", nil)
	})

	data, handled, err := dispatchVuln(context.Background(), c, "get_vuln_tags_stats", nil)
	if err != nil || !handled {
		t.Fatalf("expected (error-payload,true,nil), got handled=%v err=%v", handled, err)
	}
	e := vulnErrPayload(t, data)
	if e["code"] != 429 {
		t.Errorf("want code 429, got %v", e["code"])
	}
	if e["message"] != "rate limit exceeded" {
		t.Errorf("want passthrough message, got %v", e["message"])
	}
	if e["retry_after"] != int64(30) {
		t.Errorf("want retry_after 30, got %v", e["retry_after"])
	}
}

// TestDispatchRoutesVulnTools verifies the end-to-end wiring through
// call.go's public Dispatch: a vuln tool name falls through the CAASM
// routing to callVulnTool, which vuln.go's init replaced with
// dispatchVuln.
func TestDispatchRoutesVulnTools(t *testing.T) {
	c, reqs := newVulnMockBackend(t, func(w http.ResponseWriter, r *http.Request) {
		writeVulnEnvelope(t, w, 0, "ok", map[string]any{"total": 0})
	})

	res, err := Dispatch(context.Background(), c, "search_vulnerabilities", map[string]any{
		"severity": []any{"high", "critical"},
	})
	if err != nil {
		t.Fatalf("Dispatch returned error: %v", err)
	}
	m, ok := res.(map[string]any)
	if !ok || m["total"] != float64(0) {
		t.Fatalf("unexpected result %v", res)
	}
	got := reqs()
	if len(got) != 1 || got[0].Path != "/api/v1/vulnerabilities" {
		t.Fatalf("unexpected requests %v", got)
	}
	if got[0].Query.Get("severity") != "high,critical" {
		t.Fatalf("unexpected query %v", got[0].Query)
	}
}

func TestVulnTransportErrorMapping(t *testing.T) {
	c, _ := newVulnMockBackend(t, func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusInternalServerError)
		_, _ = w.Write([]byte("boom"))
	})

	data, handled, err := dispatchVuln(context.Background(), c, "get_vuln_tags_stats", nil)
	if err != nil || !handled {
		t.Fatalf("expected (error-payload,true,nil), got handled=%v err=%v", handled, err)
	}
	e := vulnErrPayload(t, data)
	if e["code"] != "upstream" {
		t.Errorf("want code upstream, got %v", e["code"])
	}
	msg, _ := e["message"].(string)
	if !strings.Contains(msg, "500") {
		t.Errorf("message should carry the transport status, got %q", msg)
	}
}