package egress

import (
	"context"
	"errors"
	"io"
	"net"
	"net/http"
	"net/http/httptest"
	"net/url"
	"strings"
	"sync/atomic"
	"testing"

	"github.com/stretchr/testify/require"

	"github.com/tidbcloud/AgentTeam/internal/extension/binding"
)

func TestClientInjectsAuthAndReturnsBoundedJSON(t *testing.T) {
	var receivedAuth string
	server := httptest.NewTLSServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
		receivedAuth = request.Header.Get("Authorization")
		writer.Header().Set("Content-Type", "application/json")
		writer.WriteHeader(http.StatusCreated)
		_, _ = io.WriteString(writer, `{"ok":true}`)
	}))
	defer server.Close()

	client := newTestClient(t, server)
	result, err := client.Execute(context.Background(), boundRequest(t, client.originKey, http.MethodPost),
		Credential{Type: CredentialOAuth2Bearer, Value: "provider-secret"})
	require.NoError(t, err)
	require.Equal(t, "Bearer provider-secret", receivedAuth)
	require.Equal(t, http.StatusCreated, result.UpstreamStatus)
	require.Equal(t, int64(len(`{"ok":true}`)), result.ResponseBytes)
	require.JSONEq(t, `{"ok":true}`, string(result.Data))
	require.LessOrEqual(t, result.FinalBytes, MaxFinalResponseBytes)
}

func TestClientRejectsCrossOriginRedirect(t *testing.T) {
	other := httptest.NewTLSServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) {
		t.Fatal("cross-origin redirect target must not be reached")
	}))
	defer other.Close()
	origin := httptest.NewTLSServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
		http.Redirect(writer, request, other.URL+"/stolen", http.StatusFound)
	}))
	defer origin.Close()

	client := newTestClient(t, origin)
	_, err := client.Execute(context.Background(), boundRequest(t, client.originKey, http.MethodGet), Credential{})
	require.ErrorIs(t, err, ErrCrossOriginRedirect)
}

func TestClientRejectsSameOriginRedirectOutsideFixedBasePath(t *testing.T) {
	server := httptest.NewTLSServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
		if request.URL.Path == "/admin" {
			t.Fatal("redirect outside the fixed endpoint base path must not be reached")
		}
		http.Redirect(writer, request, "/admin", http.StatusFound)
	}))
	defer server.Close()
	config := publicTestConfig(t, server)
	config.Origin.Path = "/api"
	client, err := NewClient(config)
	require.NoError(t, err)
	_, err = client.Execute(
		context.Background(),
		boundRequest(t, config.Origin.String(), http.MethodGet),
		Credential{},
	)
	require.ErrorIs(t, err, ErrDestinationDenied)
}

func TestClientRejectsRedirectThatChangesBoundMethod(t *testing.T) {
	server := httptest.NewTLSServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
		if request.URL.Path == "/v1/redirected" {
			t.Fatal("redirect that changes POST to GET must not be reached")
		}
		http.Redirect(writer, request, "/v1/redirected", http.StatusFound)
	}))
	defer server.Close()
	client := newTestClient(t, server)
	_, err := client.Execute(
		context.Background(),
		boundRequest(t, client.originKey, http.MethodPost),
		Credential{},
	)
	require.ErrorIs(t, err, ErrDestinationDenied)
}

func TestClientRevalidatesDNSAtConnectTime(t *testing.T) {
	var dialCalls atomic.Int32
	resolver := staticResolver{ips: []net.IP{net.ParseIP("127.0.0.1")}}
	origin, err := url.Parse("https://rebind.example")
	require.NoError(t, err)
	client, err := NewClient(Config{
		Origin:   origin,
		Resolver: resolver,
		DialContext: func(context.Context, string, string) (net.Conn, error) {
			dialCalls.Add(1)
			return nil, errors.New("must not dial rejected address")
		},
	})
	require.NoError(t, err)

	_, err = client.Execute(context.Background(), boundRequest(t, origin.String(), http.MethodGet), Credential{})
	require.ErrorIs(t, err, ErrDestinationDenied)
	require.Zero(t, dialCalls.Load())
}

func TestClientRejectsEveryNonPublicAddressClass(t *testing.T) {
	for _, rawIP := range []string{
		"10.0.0.1",
		"100.64.0.1",
		"169.254.10.20",
		"192.0.2.1",
		"198.18.0.1",
		"203.0.113.1",
		"240.0.0.1",
		"224.0.0.1",
		"0.0.0.0",
		"::1",
		"::ffff:100.64.0.1",
		"2001:db8::1",
	} {
		t.Run(rawIP, func(t *testing.T) {
			var dialCalls atomic.Int32
			origin, err := url.Parse("https://blocked.example")
			require.NoError(t, err)
			client, err := NewClient(Config{
				Origin: origin,
				Resolver: staticResolver{
					ips: []net.IP{net.ParseIP(rawIP)},
				},
				DialContext: func(context.Context, string, string) (net.Conn, error) {
					dialCalls.Add(1)
					return nil, errors.New("must not dial rejected address")
				},
			})
			require.NoError(t, err)
			_, err = client.Execute(
				context.Background(),
				boundRequest(t, origin.String(), http.MethodGet),
				Credential{},
			)
			require.ErrorIs(t, err, ErrDestinationDenied)
			require.Zero(t, dialCalls.Load())
		})
	}
}

func TestClientAllowsPublicDestination(t *testing.T) {
	var dialCalls atomic.Int32
	origin, err := url.Parse("https://public.example")
	require.NoError(t, err)
	client, err := NewClient(Config{
		Origin: origin,
		Resolver: staticResolver{
			ips: []net.IP{net.ParseIP("8.8.8.8")},
		},
		DialContext: func(context.Context, string, string) (net.Conn, error) {
			dialCalls.Add(1)
			return nil, errors.New("dial attempted")
		},
	})
	require.NoError(t, err)
	_, err = client.Execute(
		context.Background(),
		boundRequest(t, origin.String(), http.MethodGet),
		Credential{},
	)
	require.ErrorIs(t, err, ErrRequestFailed)
	require.NotContains(t, err.Error(), "dial attempted",
		"transport errors are sanitized because net/http may render a credential-bearing URL")
	require.NotErrorIs(t, err, ErrDestinationDenied)
	require.Equal(t, int32(1), dialCalls.Load())
}

func TestClientTransportErrorNeverLeaksQueryAPIKey(t *testing.T) {
	origin, err := url.Parse("https://public.example")
	require.NoError(t, err)
	client, err := NewClient(Config{
		Origin: origin,
		Resolver: staticResolver{
			ips: []net.IP{net.ParseIP("8.8.8.8")},
		},
		DialContext: func(context.Context, string, string) (net.Conn, error) {
			return nil, errors.New("dial failure containing internal diagnostics")
		},
	})
	require.NoError(t, err)
	bound := boundRequest(t, origin.String()+"/items?visible=safe", http.MethodGet)
	_, err = client.Execute(context.Background(), bound, Credential{
		Type: CredentialAPIKey, In: "query", Name: "api_key", Value: "query-secret",
	})
	require.ErrorIs(t, err, ErrRequestFailed)
	require.NotContains(t, err.Error(), "query-secret")
	require.NotContains(t, err.Error(), "api_key")
	require.NotContains(t, err.Error(), origin.String())
}

func TestClientUnexpectedStatusReturnsOnlySanitizedStatusMetadata(t *testing.T) {
	server := httptest.NewTLSServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) {
		writer.Header().Set("Content-Type", "application/json")
		writer.WriteHeader(http.StatusUnauthorized)
		_, _ = io.WriteString(writer, `{"error":"provider-secret-detail"}`)
	}))
	defer server.Close()
	client := newTestClient(t, server)
	result, err := client.Execute(
		context.Background(),
		boundRequest(t, client.originKey, http.MethodGet),
		Credential{Type: CredentialOAuth2Bearer, Value: "provider-secret"},
	)
	require.ErrorIs(t, err, ErrUnexpectedStatus)
	status, ok := UpstreamStatus(err)
	require.True(t, ok)
	require.Equal(t, http.StatusUnauthorized, status)
	require.Equal(t, "4xx", result.StatusClass)
	require.Equal(t, http.StatusUnauthorized, result.UpstreamStatus)
	require.NotContains(t, err.Error(), "provider-secret")
	require.NotContains(t, err.Error(), "provider-secret-detail")
	require.Empty(t, result.Data, "unexpected response plaintext must not escape normalization")
}

func TestClientRejectsAPIKeyInReservedHeader(t *testing.T) {
	server := httptest.NewTLSServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) {
		t.Fatal("reserved credential header must fail before network")
	}))
	defer server.Close()
	client := newTestClient(t, server)
	_, err := client.Execute(
		context.Background(),
		boundRequest(t, client.originKey, http.MethodGet),
		Credential{
			Type: CredentialAPIKey, In: "header", Name: "Cookie", Value: "secret",
		},
	)
	require.Error(t, err)
	require.NotContains(t, err.Error(), "secret")
}

func TestClientRejectsAPIKeyCollisionWithEmptyBoundHeader(t *testing.T) {
	target := boundRequest(t, "https://public.example", http.MethodGet)
	target.Headers.Set("X-Provider-Key", "")
	err := injectCredential(target.URL, target.Headers, Credential{
		Type:  CredentialAPIKey,
		In:    "header",
		Name:  "x-provider-key",
		Value: "provider-secret",
	})
	require.ErrorContains(t, err, "collides")
	require.Equal(t, []string{""}, target.Headers.Values("X-Provider-Key"))
}

func TestClientEnforcesRawAndFinalBudgets(t *testing.T) {
	t.Run("raw", func(t *testing.T) {
		server := httptest.NewTLSServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) {
			writer.Header().Set("Content-Type", "application/json")
			_, _ = io.WriteString(writer, `"`+strings.Repeat("x", MaxRawResponseBytes)+`"`)
		}))
		defer server.Close()
		client := newTestClient(t, server)
		_, err := client.Execute(context.Background(), boundRequest(t, client.originKey, http.MethodGet), Credential{})
		require.ErrorIs(t, err, ErrRawResponseTooLarge)
	})
	t.Run("final", func(t *testing.T) {
		server := httptest.NewTLSServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) {
			writer.Header().Set("Content-Type", "application/json")
			_, _ = io.WriteString(writer, `"`+strings.Repeat("x", MaxFinalResponseBytes)+`"`)
		}))
		defer server.Close()
		client := newTestClient(t, server)
		_, err := client.Execute(context.Background(), boundRequest(t, client.originKey, http.MethodGet), Credential{})
		require.ErrorIs(t, err, ErrFinalResponseTooLarge)
	})
}

func TestClientRejectsBinaryContentType(t *testing.T) {
	server := httptest.NewTLSServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) {
		writer.Header().Set("Content-Type", "application/octet-stream")
		_, _ = writer.Write([]byte{0, 1, 2})
	}))
	defer server.Close()
	client := newTestClient(t, server)
	_, err := client.Execute(context.Background(), boundRequest(t, client.originKey, http.MethodGet), Credential{})
	require.ErrorIs(t, err, ErrUnsupportedContentType)
}

type staticResolver struct {
	ips []net.IP
	err error
}

func (resolver staticResolver) LookupIP(context.Context, string, string) ([]net.IP, error) {
	return append([]net.IP(nil), resolver.ips...), resolver.err
}

func newTestClient(t *testing.T, server *httptest.Server) *Client {
	t.Helper()
	client, err := NewClient(publicTestConfig(t, server))
	require.NoError(t, err)
	return client
}

func publicTestConfig(t *testing.T, server *httptest.Server) Config {
	t.Helper()
	listenerAddress := server.Listener.Addr().String()
	_, port, err := net.SplitHostPort(listenerAddress)
	require.NoError(t, err)
	origin := &url.URL{
		Scheme: "https",
		Host:   net.JoinHostPort("public.example", port),
	}
	tlsConfig := server.Client().Transport.(*http.Transport).
		TLSClientConfig.Clone()
	tlsConfig.ServerName = "example.com"
	return Config{
		Origin:   origin,
		Resolver: staticResolver{ips: []net.IP{net.ParseIP("8.8.8.8")}},
		DialContext: func(ctx context.Context, network, _ string) (net.Conn, error) {
			return (&net.Dialer{}).DialContext(ctx, network, listenerAddress)
		},
		TLSClientConfig: tlsConfig,
	}
}

func boundRequest(t *testing.T, origin, method string) binding.BoundRequest {
	t.Helper()
	target, err := url.Parse(origin + "/v1/action")
	require.NoError(t, err)
	return binding.BoundRequest{
		Method: method,
		URL:    target,
		Headers: http.Header{
			"Accept": []string{"application/json"},
		},
		Body: []byte(`{"input":true}`),
		ResponsePolicy: binding.ResponsePolicy{
			Mode:            "passthrough",
			SuccessStatuses: []int{http.StatusOK, http.StatusCreated},
			ContentTypes:    []string{"application/json", "text/plain"},
		},
		SecretFreeHash: "sha256:request",
	}
}