package main

import (
	"fmt"
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"
)

func TestFillHostBedrock(t *testing.T) {
	rt := &route{Upstream: "https://bedrock-runtime.{region}.amazonaws.com", HostParams: []hostParam{{"region", awsRegion}}}
	cases := []struct {
		path         string
		wantUpstream string
		wantRest     string
		wantOK       bool
	}{
		{"/us-west-2/openai/v1/chat/completions", "https://bedrock-runtime.us-west-2.amazonaws.com", "/openai/v1/chat/completions", true},
		{"/ap-southeast-1/model/x/invoke", "https://bedrock-runtime.ap-southeast-1.amazonaws.com", "/model/x/invoke", true},
		{"/us-gov-east-1/openai/v1/chat/completions", "https://bedrock-runtime.us-gov-east-1.amazonaws.com", "/openai/v1/chat/completions", true},
		{"/eu-central-1", "https://bedrock-runtime.eu-central-1.amazonaws.com", "/", true},
		{"/us-west-2/openai/v1/models?limit=1", "https://bedrock-runtime.us-west-2.amazonaws.com", "/openai/v1/models?limit=1", true},

		// Anything that could widen the host past bedrock-runtime.<region>.amazonaws.com.
		{"/s3/bucket/key", "", "", false},              // bedrock-runtime.s3.amazonaws.com
		{"/us-west-2.evil.com/x", "", "", false},       // extra labels (dot)
		{"/../../x", "", "", false},                    // traversal
		{"/US-WEST-2/x", "", "", false},                // case
		{"/us_west_2/x", "", "", false},                // underscore
		{"/us-west-2:80/x", "", "", false},             // port
		{"/", "", "", false},                           // empty segment
		{"/openai/v1/chat/completions", "", "", false}, // region omitted entirely
	}
	for _, c := range cases {
		upstream, rest, _, ok := rt.fillHost(c.path)
		if ok != c.wantOK || upstream != c.wantUpstream || rest != c.wantRest {
			t.Errorf("fillHost(%q) = (%q, %q, %v), want (%q, %q, %v)",
				c.path, upstream, rest, ok, c.wantUpstream, c.wantRest, c.wantOK)
		}
	}
}

func TestFillHostO11y(t *testing.T) {
	rt := &route{Upstream: "https://www.ds.{region}.{vendor}.observability.tidbcloud.com", HostParams: []hostParam{{"region", o11yLabel}, {"vendor", o11yLabel}}}
	cases := []struct {
		path         string
		wantUpstream string
		wantRest     string
		wantOK       bool
	}{
		{"/us-west-2/aws/internal/metrics/x/api/v1/query", "https://www.ds.us-west-2.aws.observability.tidbcloud.com", "/internal/metrics/x/api/v1/query", true},
		{"/us-central1/gcp/internal/metrics/x/api/v1/query_range?step=60", "https://www.ds.us-central1.gcp.observability.tidbcloud.com", "/internal/metrics/x/api/v1/query_range?step=60", true},
		{"/ap-southeast-1/alicloud", "https://www.ds.ap-southeast-1.alicloud.observability.tidbcloud.com", "/", true},

		{"/us-west-2", "", "", false},            // vendor missing
		{"/us-west-2.evil/aws/x", "", "", false}, // dot in region
		{"/us-west-2/aws.evil/x", "", "", false}, // dot in vendor
		{"/US-WEST-2/aws/x", "", "", false},      // case
		{"/", "", "", false},                     // empty
	}
	for _, c := range cases {
		upstream, rest, _, ok := rt.fillHost(c.path)
		if ok != c.wantOK || upstream != c.wantUpstream || rest != c.wantRest {
			t.Errorf("fillHost(%q) = (%q, %q, %v), want (%q, %q, %v)",
				c.path, upstream, rest, ok, c.wantUpstream, c.wantRest, c.wantOK)
		}
	}
}

func TestServeCredentials(t *testing.T) {
	do := func(p *proxy, method string) *httptest.ResponseRecorder {
		w := httptest.NewRecorder()
		p.ServeHTTP(w, httptest.NewRequest(method, "/o11y/credentials", nil))
		return w
	}
	enabled := func(o11y *o11yAuth) *proxy {
		return &proxy{secrets: &secretsStore{data: secrets{O11yExposeCreds: true}}, o11y: o11y}
	}
	full := func() *o11yAuth {
		return &o11yAuth{token: "eyTOKEN", auth0: "eyAUTH0", awsAK: "ASIA1", awsSK: "secret1", awsST: "session1"}
	}

	// Toggle off (default) → 404, even with creds present: reads as absent.
	off := &proxy{secrets: &secretsStore{}, o11y: full()}
	if w := do(off, http.MethodGet); w.Code != http.StatusNotFound {
		t.Errorf("toggle off: got %d, want 404", w.Code)
	}
	if b := do(off, http.MethodGet).Body.String(); strings.Contains(b, "ASIA1") || strings.Contains(b, "eyTOKEN") {
		t.Errorf("toggle off leaked creds: %s", b)
	}

	// Enabled but o11y rotation disabled → 503.
	if w := do(enabled(nil), http.MethodGet); w.Code != http.StatusServiceUnavailable {
		t.Errorf("o11y disabled: got %d, want 503", w.Code)
	}

	// Enabled but no creds fetched yet → 503.
	if w := do(enabled(&o11yAuth{}), http.MethodGet); w.Code != http.StatusServiceUnavailable {
		t.Errorf("no creds yet: got %d, want 503", w.Code)
	}

	// Enabled with creds → 200 with o11y_token + STS, but never the auth0 token.
	p := enabled(full())
	w := do(p, http.MethodGet)
	if w.Code != http.StatusOK {
		t.Fatalf("with creds: got %d, want 200", w.Code)
	}
	b := w.Body.String()
	if !strings.Contains(b, "eyTOKEN") || !strings.Contains(b, "ASIA1") || !strings.Contains(b, "session1") {
		t.Errorf("body missing creds: %s", b)
	}
	if strings.Contains(b, "eyAUTH0") {
		t.Errorf("body leaked the proxy-injected auth0 token: %s", b)
	}

	// Non-GET → 405, and never leaks creds.
	if w := do(p, http.MethodPost); w.Code != http.StatusMethodNotAllowed {
		t.Errorf("POST: got %d, want 405", w.Code)
	}
}

// slowBody yields 3 chunks, 150ms apart, like an SSE stream.
type slowBody struct{ n int }

func (s *slowBody) Read(p []byte) (int, error) {
	if s.n >= 3 {
		return 0, io.EOF
	}
	time.Sleep(150 * time.Millisecond)
	s.n++
	return copy(p, "data: chunk\n\n"), nil
}

// firstByteLatency measures how long after the request starts the client can
// read the first body byte.
func firstByteLatency(t *testing.T, copyFn func(http.ResponseWriter, io.Reader)) time.Duration {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "text/event-stream")
		w.WriteHeader(200)
		copyFn(w, &slowBody{})
	}))
	defer srv.Close()

	start := time.Now()
	resp, err := http.Get(srv.URL)
	if err != nil {
		t.Fatal(err)
	}
	defer resp.Body.Close()
	buf := make([]byte, 8)
	if _, err := io.ReadFull(resp.Body, buf); err != nil {
		t.Fatal(err)
	}
	return time.Since(start)
}

func TestCopyStreamFlushesEachChunk(t *testing.T) {
	streamed := firstByteLatency(t, copyStream)
	buffered := firstByteLatency(t, func(w http.ResponseWriter, r io.Reader) { io.Copy(w, r) })
	t.Logf("copyStream first byte: %v | io.Copy first byte: %v",
		streamed.Round(10*time.Millisecond), buffered.Round(10*time.Millisecond))

	// 3 chunks x 150ms: streaming delivers at ~150ms, buffering at ~450ms.
	if streamed > 300*time.Millisecond {
		t.Fatalf("copyStream first byte took %v — not streaming", streamed)
	}
	if buffered < 400*time.Millisecond {
		t.Fatalf("control did not buffer as expected (%v); test is not proving anything", buffered)
	}
}

func TestPermissionMode(t *testing.T) {
	const cbPath = "/api/v1/alertstore/alerts/ALT-1/agent-runs/42"
	cases := []struct {
		mode         permissionMode
		method, path string
		want         bool
	}{
		// reads always pass
		{permReadOnly, "GET", "/api/v1/x", true},
		{permAgent, "GET", "/api/v1/x", true},
		// readonly rejects the callback PATCH; agent allows it
		{permReadOnly, "PATCH", cbPath, false},
		{permAgent, "PATCH", cbPath, true},
		// agent allows ONLY the callback path, not other writes
		{permAgent, "PATCH", "/api/v1/devops/clusters/1", false},
		{permAgent, "POST", cbPath, false},
		{permAgent, "DELETE", cbPath, false},
		// git push: blocked in readonly, allowed at method level in agent
		// (forward then restricts the refs); upload-pack stays read-only
		{permReadOnly, "POST", "/o/r.git/git-receive-pack", false},
		{permAgent, "POST", "/o/r.git/git-receive-pack", true},
		{permReadOnly, "POST", "/o/r.git/git-upload-pack", true},
		// full allows anything
		{permFull, "DELETE", "/whatever", true},
	}
	for _, c := range cases {
		p := &proxy{perm: c.mode}
		if got := p.methodAllowed(c.method, c.path); got != c.want {
			t.Errorf("mode=%s %s %s: got %v want %v", c.mode, c.method, c.path, got, c.want)
		}
	}
	if _, ok := parsePermissionMode("bogus"); ok {
		t.Error("parsePermissionMode accepted an invalid mode")
	}
}

// pkt builds a git pkt-line (4-hex length prefix + payload).
func pkt(payload string) string {
	return fmt.Sprintf("%04x%s", len(payload)+4, payload)
}

const flushPkt = "0000"

func TestValidateAndBufferReceivePack(t *testing.T) {
	const zero = "0000000000000000000000000000000000000000"
	const oid = "1234567890123456789012345678901234567890"
	// first command carries a capability list after a NUL
	cmd := func(ref string) string {
		return pkt(zero+" "+oid+" "+ref+"\x00 report-status side-band-64k") + flushPkt
	}
	packfile := "PACKtrailingbytes"

	cases := []struct {
		name    string
		section string // pkt-line command section (up to+incl flush-pkt)
		wantOK  bool
	}{
		{"agent branch", cmd("refs/heads/agent/foo"), true},
		{"agent nested", cmd("refs/heads/agent/x/y/z"), true},
		{"human branch", cmd("refs/heads/main"), false},
		{"agent delete allowed", pkt(oid+" "+zero+" refs/heads/agent/foo\x00report-status") + flushPkt, true},
		{"tag rejected", cmd("refs/tags/v1"), false},
		{"prefix lookalike rejected", cmd("refs/heads/agentX"), false},
		{"malformed rejected", pkt("garbage") + flushPkt, false},
		{"truncated rejected", pkt(zero + " " + oid + " refs/heads/agent/foo"), false}, // no flush-pkt
	}

	for _, c := range cases {
		t.Run(c.name, func(t *testing.T) {
			full := c.section + packfile
			consumed, ok := validateAndBufferReceivePack(strings.NewReader(full))
			if ok != c.wantOK {
				t.Fatalf("ok=%v want %v", ok, c.wantOK)
			}
			if !ok {
				return
			}
			// consumed must be exactly the command section, leaving the packfile
			// for the caller to stream through unread.
			if string(consumed) != c.section {
				t.Errorf("consumed=%q want %q", consumed, c.section)
			}
		})
	}
}

func TestReceivePackMultiRefAllOrNothing(t *testing.T) {
	const oid = "1234567890123456789012345678901234567890"
	const zero = "0000000000000000000000000000000000000000"
	// two commands: one allowed, one human branch -> whole push rejected
	section := pkt(zero+" "+oid+" refs/heads/agent/ok\x00report-status") +
		pkt(zero+" "+oid+" refs/heads/main") + flushPkt
	if _, ok := validateAndBufferReceivePack(strings.NewReader(section + "PACK")); ok {
		t.Error("push with a human-branch command should be rejected")
	}
}