package main

import (
	"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)
	}
}