package tools

import (
	"context"
	"encoding/json"
	"errors"
	"io"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"testing"

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

// uploadRecordedCall captures one request of the three-step upload
// flow: step is "presign" (POST /uploads), "put" (raw PUT to the
// presigned URL) or "confirm" (POST /uploads/confirm).
type uploadRecordedCall struct {
	Step        string
	Method      string
	Path        string
	ContentType string
	APIKey      string
	JSONBody    map[string]any
	RawBody     []byte
}

// newUploadMockBackend starts one httptest server doubling as the
// backend API and the object store (the presigned upload_url it hands
// out points back at itself under /mock-s3/). handler is invoked per
// request with the classified step and the server URL (for building
// presigned URLs); every request is recorded in call order.
func newUploadMockBackend(t *testing.T, handler func(step, srvURL string, w http.ResponseWriter, r *http.Request)) (*client.Client, func() []uploadRecordedCall) {
	t.Helper()
	var mu sync.Mutex
	var calls []uploadRecordedCall
	var srv *httptest.Server
	srv = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		raw, _ := io.ReadAll(r.Body)
		call := uploadRecordedCall{
			Method:      r.Method,
			Path:        r.URL.Path,
			ContentType: r.Header.Get("Content-Type"),
			APIKey:      r.Header.Get("X-API-Key"),
		}
		switch {
		case r.Method == http.MethodPut:
			call.Step = "put"
			call.RawBody = raw
		case strings.HasSuffix(r.URL.Path, "/uploads/confirm"):
			call.Step = "confirm"
			_ = json.Unmarshal(raw, &call.JSONBody)
		default:
			call.Step = "presign"
			_ = json.Unmarshal(raw, &call.JSONBody)
		}
		mu.Lock()
		calls = append(calls, call)
		mu.Unlock()
		handler(call.Step, srv.URL, w, r)
	}))
	t.Cleanup(srv.Close)
	c := &client.Client{BaseURL: srv.URL, APIKey: "test-key", HTTP: srv.Client()}
	return c, func() []uploadRecordedCall {
		mu.Lock()
		defer mu.Unlock()
		return append([]uploadRecordedCall(nil), calls...)
	}
}

// happyUploadHandler answers all three steps successfully: presign
// returns attachment 123 with a presigned URL pointing at the mock
// store, the store accepts the PUT, confirm marks it ready.
func happyUploadHandler(t *testing.T) func(step, srvURL string, w http.ResponseWriter, r *http.Request) {
	return func(step, srvURL string, w http.ResponseWriter, r *http.Request) {
		switch step {
		case "presign":
			writeVulnEnvelope(t, w, 0, "ok", map[string]any{
				"attachment_id": 123,
				"upload_url":    srvURL + "/mock-s3/vuln-attachments/key-1.png?X-Amz-Signature=redacted&X-Amz-Expires=900",
				"key":           "vuln-attachments/key-1.png",
				"expires_in":    900,
			})
		case "put":
			w.WriteHeader(http.StatusOK)
		case "confirm":
			writeVulnEnvelope(t, w, 0, "ok", map[string]any{"id": 123, "status": "ready"})
		}
	}
}

// writeUploadTempFile writes content to a t.TempDir file and returns
// its path.
func writeUploadTempFile(t *testing.T, name string, content []byte) string {
	t.Helper()
	path := filepath.Join(t.TempDir(), name)
	if err := os.WriteFile(path, content, 0o600); err != nil {
		t.Fatalf("write temp file: %v", err)
	}
	return path
}

// TestUploadToolRegistered asserts the tool's schema shape, the
// contract endpoints, and that the description carries the known
// limitations (同机/stdio 形态、仅图片 ≤10MB、无 SVG).
func TestUploadToolRegistered(t *testing.T) {
	tool, ok := findTool("upload_vulnerability_attachment")
	if !ok {
		t.Fatal("tool upload_vulnerability_attachment not registered")
	}

	if len(tool.InputSchema.Required) != 1 || tool.InputSchema.Required[0] != "file_path" {
		t.Errorf("required = %v, want [file_path]", tool.InputSchema.Required)
	}
	ct, ok := tool.InputSchema.Properties["content_type"]
	if !ok {
		t.Fatal("content_type property missing")
	}
	wantEnum := []string{"image/png", "image/jpeg", "image/webp", "image/gif"}
	if len(ct.Enum) != len(wantEnum) {
		t.Fatalf("content_type enum = %v, want %v", ct.Enum, wantEnum)
	}
	for i, v := range wantEnum {
		if ct.Enum[i] != v {
			t.Errorf("content_type enum[%d] = %q, want %q", i, ct.Enum[i], v)
		}
	}

	for _, sub := range []string{"同机", "stdio", "10MB", "SVG", "attachment_id", "data-sc-attachment-id"} {
		if !strings.Contains(tool.Description, sub) {
			t.Errorf("description missing limitation wording %q", sub)
		}
	}

	// 编排工具：登记契约两端点（post /uploads + post /uploads/confirm）
	// 供注册表级 swagger 校验覆盖，但无 ParamMap/BodyMap。
	if len(tool.Endpoints) != 2 {
		t.Fatalf("want 2 endpoints, got %d", len(tool.Endpoints))
	}
	wantPaths := []string{"/api/v1/vulnerabilities/uploads", "/api/v1/vulnerabilities/uploads/confirm"}
	for i, ep := range tool.Endpoints {
		if ep.Method != "post" || ep.Path != wantPaths[i] {
			t.Errorf("endpoint #%d = %s %s, want post %s", i, ep.Method, ep.Path, wantPaths[i])
		}
		if len(ep.ParamMap) > 0 || len(ep.BodyMap) > 0 {
			t.Errorf("endpoint #%d: orchestrated tool must not declare ParamMap/BodyMap", i)
		}
	}
}

// TestUploadAttachmentThreeStepSequence asserts the happy path: presign
// → raw PUT → confirm, in order, with the presign body declaring
// filename+content_type, the PUT carrying the exact file bytes with the
// declared Content-Type (and no API key), and the confirm body carrying
// the attachment_id; the result surfaces attachment_id (+key).
func TestUploadAttachmentThreeStepSequence(t *testing.T) {
	fileContent := []byte("\x89PNG\r\n\x1a\n-fake-screenshot-bytes")
	path := writeUploadTempFile(t, "evidence.png", fileContent)

	c, calls := newUploadMockBackend(t, happyUploadHandler(t))
	data, handled, err := dispatchVuln(context.Background(), c, "upload_vulnerability_attachment", map[string]any{
		"file_path": path,
	})
	if err != nil || !handled {
		t.Fatalf("expected (data,true,nil), got handled=%v err=%v", handled, err)
	}

	got := calls()
	if len(got) != 3 {
		t.Fatalf("expected 3 requests, got %d", len(got))
	}
	for i, step := range []string{"presign", "put", "confirm"} {
		if got[i].Step != step {
			t.Fatalf("request #%d: want step %q, got %q (%s %s)", i, step, got[i].Step, got[i].Method, got[i].Path)
		}
	}

	// ① 预签名：POST /uploads，body 声明 filename + content_type。
	presign := got[0]
	if presign.Method != http.MethodPost || presign.Path != "/api/v1/vulnerabilities/uploads" {
		t.Errorf("presign: want POST /api/v1/vulnerabilities/uploads, got %s %s", presign.Method, presign.Path)
	}
	if presign.JSONBody["filename"] != "evidence.png" {
		t.Errorf("presign body filename = %v, want evidence.png", presign.JSONBody["filename"])
	}
	if presign.JSONBody["content_type"] != "image/png" {
		t.Errorf("presign body content_type = %v, want image/png", presign.JSONBody["content_type"])
	}

	// ② 直传：字节原样 PUT 到预签名地址，Content-Type 用申请时声明
	// 值（S3 预签名绑定），不带 X-API-Key（预签名本身即凭证）。
	put := got[1]
	if put.Method != http.MethodPut || put.Path != "/mock-s3/vuln-attachments/key-1.png" {
		t.Errorf("put: want PUT /mock-s3/vuln-attachments/key-1.png, got %s %s", put.Method, put.Path)
	}
	if string(put.RawBody) != string(fileContent) {
		t.Errorf("put body = %q, want exact file bytes", put.RawBody)
	}
	if put.ContentType != "image/png" {
		t.Errorf("put Content-Type = %q, want image/png (预签名申请时的声明值)", put.ContentType)
	}
	if put.APIKey != "" {
		t.Errorf("put must not carry X-API-Key, got %q", put.APIKey)
	}

	// ③ 确认：POST /uploads/confirm，body 只有 attachment_id。
	confirm := got[2]
	if confirm.Method != http.MethodPost || confirm.Path != "/api/v1/vulnerabilities/uploads/confirm" {
		t.Errorf("confirm: want POST /api/v1/vulnerabilities/uploads/confirm, got %s %s", confirm.Method, confirm.Path)
	}
	if confirm.JSONBody["attachment_id"] != float64(123) {
		t.Errorf("confirm body attachment_id = %v, want 123", confirm.JSONBody["attachment_id"])
	}
	if len(confirm.JSONBody) != 1 {
		t.Errorf("confirm body must carry only attachment_id, got %v", confirm.JSONBody)
	}

	// 返回 attachment_id（+ key）透出给 LLM 引用。
	m, ok := data.(map[string]any)
	if !ok {
		t.Fatalf("expected result map, got %T (%v)", data, data)
	}
	if m["attachment_id"] != uint64(123) {
		t.Errorf("result attachment_id = %v (%T), want 123", m["attachment_id"], m["attachment_id"])
	}
	if m["key"] != "vuln-attachments/key-1.png" {
		t.Errorf("result key = %v, want vuln-attachments/key-1.png", m["key"])
	}
}

// TestUploadAttachmentContentType asserts content_type inference from
// the extension (case-insensitive) and that an explicit content_type
// overrides the inferred one on both the presign body and the PUT.
func TestUploadAttachmentContentType(t *testing.T) {
	cases := []struct {
		name        string
		fileName    string
		contentType string // explicit argument, "" = infer
		wantMIME    string
	}{
		{name: "jpg inferred", fileName: "a.jpg", wantMIME: "image/jpeg"},
		{name: "jpeg inferred", fileName: "b.jpeg", wantMIME: "image/jpeg"},
		{name: "webp inferred", fileName: "c.webp", wantMIME: "image/webp"},
		{name: "gif inferred", fileName: "d.gif", wantMIME: "image/gif"},
		{name: "uppercase ext", fileName: "E.PNG", wantMIME: "image/png"},
		{name: "explicit override", fileName: "f.png", contentType: "image/webp", wantMIME: "image/webp"},
	}
	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			path := writeUploadTempFile(t, tc.fileName, []byte("img"))
			c, calls := newUploadMockBackend(t, happyUploadHandler(t))

			args := map[string]any{"file_path": path}
			if tc.contentType != "" {
				args["content_type"] = tc.contentType
			}
			_, handled, err := dispatchVuln(context.Background(), c, "upload_vulnerability_attachment", args)
			if err != nil || !handled {
				t.Fatalf("expected (data,true,nil), got handled=%v err=%v", handled, err)
			}

			got := calls()
			if len(got) != 3 {
				t.Fatalf("expected 3 requests, got %d", len(got))
			}
			if got[0].JSONBody["content_type"] != tc.wantMIME {
				t.Errorf("presign content_type = %v, want %s", got[0].JSONBody["content_type"], tc.wantMIME)
			}
			if got[1].ContentType != tc.wantMIME {
				t.Errorf("PUT Content-Type = %q, want %s (必须与预签名声明一致)", got[1].ContentType, tc.wantMIME)
			}
		})
	}
}

// TestUploadAttachmentLocalInterception asserts local problems (missing
// file, unsupported extension, >10MB, unsupported explicit
// content_type, missing file_path) fail fast as *ArgError before any
// HTTP call.
func TestUploadAttachmentLocalInterception(t *testing.T) {
	dir := t.TempDir()
	svgPath := writeUploadTempFile(t, "evil.svg", []byte("<svg/>"))
	noExtPath := writeUploadTempFile(t, "evidence", []byte("img"))
	pngPath := writeUploadTempFile(t, "ok.png", []byte("img"))
	bigPath := filepath.Join(dir, "big.png")
	if err := os.WriteFile(bigPath, make([]byte, maxAttachmentBytes+1), 0o600); err != nil {
		t.Fatalf("write big temp file: %v", err)
	}

	cases := []struct {
		name    string
		args    map[string]any
		wantMsg []string
	}{
		{
			name:    "missing file_path",
			args:    map[string]any{},
			wantMsg: []string{"file_path"},
		},
		{
			name:    "file not found",
			args:    map[string]any{"file_path": filepath.Join(dir, "nope.png")},
			wantMsg: []string{"不存在"},
		},
		{
			name:    "svg rejected",
			args:    map[string]any{"file_path": svgPath},
			wantMsg: []string{".svg", "SVG"},
		},
		{
			name:    "no extension",
			args:    map[string]any{"file_path": noExtPath},
			wantMsg: []string{"扩展名"},
		},
		{
			name:    "over 10MB",
			args:    map[string]any{"file_path": bigPath},
			wantMsg: []string{"10MB"},
		},
		{
			name:    "unsupported explicit content_type",
			args:    map[string]any{"file_path": pngPath, "content_type": "image/svg+xml"},
			wantMsg: []string{"content_type", "image/svg+xml"},
		},
	}
	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			c, calls := newUploadMockBackend(t, func(step, srvURL string, w http.ResponseWriter, r *http.Request) {
				t.Errorf("backend must not be called on local interception: %s %s", r.Method, r.URL.Path)
			})

			data, handled, err := dispatchVuln(context.Background(), c, "upload_vulnerability_attachment", tc.args)
			if data != nil || !handled {
				t.Fatalf("expected (nil,true,ArgError), got data=%v handled=%v", data, handled)
			}
			var argErr *ArgError
			if !errors.As(err, &argErr) {
				t.Fatalf("expected *ArgError, got %T (%v)", err, err)
			}
			for _, sub := range tc.wantMsg {
				if !strings.Contains(argErr.Msg, sub) {
					t.Errorf("message should mention %q, got %q", sub, argErr.Msg)
				}
			}
			if got := len(calls()); got != 0 {
				t.Fatalf("expected zero HTTP calls, got %d", got)
			}
		})
	}
}

// TestUploadAttachmentStepFailures asserts a failure at any of the
// three steps (presign 403, PUT 403, confirm 400) aborts the flow and
// returns the structured error payload whose message names the failed
// step, with no later step attempted.
func TestUploadAttachmentStepFailures(t *testing.T) {
	path := writeUploadTempFile(t, "evidence.png", []byte("\x89PNG-fake"))

	t.Run("presign 403", func(t *testing.T) {
		c, calls := newUploadMockBackend(t, func(step, srvURL string, w http.ResponseWriter, r *http.Request) {
			writeVulnEnvelope(t, w, 403, "not authorized: missing scope vuln:submit", nil)
		})
		data, handled, err := dispatchVuln(context.Background(), c, "upload_vulnerability_attachment", map[string]any{"file_path": path})
		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"] != 403 {
			t.Errorf("want code 403, got %v", e["code"])
		}
		msg, _ := e["message"].(string)
		for _, sub := range []string{"第 1 步", "预签名", "missing scope"} {
			if !strings.Contains(msg, sub) {
				t.Errorf("message should mention %q, got %q", sub, msg)
			}
		}
		if got := len(calls()); got != 1 {
			t.Fatalf("flow must stop at step 1, got %d requests", got)
		}
	})

	t.Run("put 403", func(t *testing.T) {
		c, calls := newUploadMockBackend(t, func(step, srvURL string, w http.ResponseWriter, r *http.Request) {
			if step == "put" {
				w.WriteHeader(http.StatusForbidden)
				_, _ = w.Write([]byte("SignatureDoesNotMatch"))
				return
			}
			happyUploadHandler(t)(step, srvURL, w, r)
		})
		data, handled, err := dispatchVuln(context.Background(), c, "upload_vulnerability_attachment", map[string]any{"file_path": path})
		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)
		for _, sub := range []string{"第 2 步", "直传", "403"} {
			if !strings.Contains(msg, sub) {
				t.Errorf("message should mention %q, got %q", sub, msg)
			}
		}
		if got := len(calls()); got != 2 {
			t.Fatalf("flow must stop at step 2, got %d requests", got)
		}
	})

	t.Run("confirm 400", func(t *testing.T) {
		c, calls := newUploadMockBackend(t, func(step, srvURL string, w http.ResponseWriter, r *http.Request) {
			if step == "confirm" {
				writeVulnEnvelope(t, w, 400, "attachment rejected: size exceeds 10MB", nil)
				return
			}
			happyUploadHandler(t)(step, srvURL, w, r)
		})
		data, handled, err := dispatchVuln(context.Background(), c, "upload_vulnerability_attachment", map[string]any{"file_path": path})
		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"] != 400 {
			t.Errorf("want code 400, got %v", e["code"])
		}
		msg, _ := e["message"].(string)
		for _, sub := range []string{"第 3 步", "确认", "10MB"} {
			if !strings.Contains(msg, sub) {
				t.Errorf("message should mention %q, got %q", sub, msg)
			}
		}
		if got := len(calls()); got != 3 {
			t.Fatalf("flow must run all 3 steps before failing, got %d requests", got)
		}
	})
}

// TestUploadAttachmentThroughPublicAPI verifies the end-to-end wiring
// of the upload tool through call.go's public Dispatch (routing falls
// through CAASM to callVulnTool → dispatchVuln → vulnUploadAttachment).
func TestUploadAttachmentThroughPublicAPI(t *testing.T) {
	path := writeUploadTempFile(t, "evidence.png", []byte("\x89PNG-fake"))
	c, calls := newUploadMockBackend(t, happyUploadHandler(t))

	res, err := Dispatch(context.Background(), c, "upload_vulnerability_attachment", map[string]any{
		"file_path": path,
	})
	if err != nil {
		t.Fatalf("Dispatch returned error: %v", err)
	}
	m, ok := res.(map[string]any)
	if !ok || m["attachment_id"] != uint64(123) {
		t.Fatalf("unexpected result %v", res)
	}
	if got := len(calls()); got != 3 {
		t.Fatalf("expected 3 backend requests, got %d", got)
	}
}