package wire

import (
	"context"
	"crypto/tls"
	"crypto/x509"
	"errors"
	"net"
	"net/http"
	"net/http/httptest"
	"net/url"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"github.com/gorilla/websocket"
)

// gatewayWebSocketDialCapability is the immutable, per-call P3 capability.
// The API consumer supplies only sanitized request data; wire owns every
// transport/security field on the gorilla Dialer it creates internally.
type gatewayWebSocketDialCapability interface {
	GatewayWebSocketDial(
		ctx context.Context,
		targetURL string,
		header http.Header,
		subprotocols []string,
	) (*websocket.Conn, *http.Response, error)
}

// gatewayWebSocketHandshakePolicy exposes only the pure, immutable timeout
// decision to same-package tests. It prevents elapsed-time tests from being
// satisfied by an accidentally tiny timeout while keeping the public
// ServerGatewayConnection surface limited to the dial capability.
type gatewayWebSocketHandshakePolicy interface {
	gatewayWebSocketHandshakeTimeout() (time.Duration, error)
}

func requireGatewayWebSocketDialCapability(t *testing.T, protocol *Protocol) gatewayWebSocketDialCapability {
	t.Helper()
	capability, ok := any(protocol).(gatewayWebSocketDialCapability)
	if !ok {
		t.Fatalf("Protocol does not implement GatewayWebSocketDial")
	}
	return capability
}

func websocketURL(raw string) string {
	u, _ := url.Parse(raw)
	switch u.Scheme {
	case "https":
		u.Scheme = "wss"
	default:
		u.Scheme = "ws"
	}
	return u.String()
}

func TestGatewayWebSocketDialDefaultTransportProductionPath(t *testing.T) {
	upgrader := websocket.Upgrader{
		CheckOrigin:       func(*http.Request) bool { return true },
		EnableCompression: false,
	}
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		conn, err := upgrader.Upgrade(w, r, nil)
		if err != nil {
			return
		}
		_ = conn.Close()
	}))
	defer server.Close()

	// Production wire.New constructs the control-plane http.Client with
	// Transport=nil (see BuildGatewayHTTPClient). A previous revision of this
	// test masked the production shape by manually injecting
	// http.DefaultTransport.Clone(); the gateway then rejected nil Transport
	// with ErrUnsupportedGatewayTransport, breaking every real WS dial.
	// Keep this test honest: the Protocol must look exactly like production.
	protocol := &Protocol{httpClient: &http.Client{}}
	dial := requireGatewayWebSocketDialCapability(t, protocol)
	conn, response, err := dial.GatewayWebSocketDial(
		context.Background(),
		websocketURL(server.URL),
		nil,
		nil,
	)
	if response != nil && response.Body != nil {
		defer response.Body.Close()
	}
	if err != nil {
		t.Fatalf("GatewayWebSocketDial(default transport): %v", err)
	}
	if conn == nil {
		t.Fatal("GatewayWebSocketDial returned nil connection")
	}
	_ = conn.Close()
}

// TestGatewayWebSocketDialResolvesNilTransportViaPrivateClone verifies the
// codex5 blocker-1 fix: when the production http.Client carries Transport=nil
// (the wire.New shape), GatewayWebSocketDial must (a) resolve nil to
// http.DefaultTransport, (b) Clone it per call so concurrent callers never
// observe mutations to the shared global, and (c) leave the global
// http.DefaultTransport pointer stable across calls.
func TestGatewayWebSocketDialResolvesNilTransportViaPrivateClone(t *testing.T) {
	upgrader := websocket.Upgrader{
		CheckOrigin:       func(*http.Request) bool { return true },
		EnableCompression: false,
	}
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		conn, err := upgrader.Upgrade(w, r, nil)
		if err == nil {
			_ = conn.Close()
		}
	}))
	defer server.Close()

	before := http.DefaultTransport
	protocol := &Protocol{httpClient: &http.Client{}}
	dial := requireGatewayWebSocketDialCapability(t, protocol)

	const callers = 4
	var wg sync.WaitGroup
	start := make(chan struct{})
	errs := make(chan error, callers)
	for i := 0; i < callers; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			<-start
			conn, response, err := dial.GatewayWebSocketDial(
				context.Background(),
				websocketURL(server.URL),
				nil,
				nil,
			)
			if response != nil && response.Body != nil {
				_ = response.Body.Close()
			}
			if err != nil {
				errs <- err
				return
			}
			_ = conn.Close()
		}()
	}
	close(start)
	wg.Wait()
	close(errs)
	for err := range errs {
		if err != nil {
			t.Fatalf("concurrent nil-Transport dial failed: %v", err)
		}
	}
	if http.DefaultTransport != before {
		t.Fatalf("http.DefaultTransport was replaced: got %#v, want %#v", http.DefaultTransport, before)
	}
}

func TestGatewayWebSocketDialInheritsProxyDialAndTLS(t *testing.T) {
	upgrader := websocket.Upgrader{CheckOrigin: func(*http.Request) bool { return true }}
	server := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if got := r.TLS.NegotiatedProtocol; got != "" && got != "http/1.1" {
			t.Errorf("negotiated protocol = %q, want HTTP/1.1", got)
		}
		conn, err := upgrader.Upgrade(w, r, nil)
		if err == nil {
			_ = conn.Close()
		}
	}))
	defer server.Close()

	pool := x509.NewCertPool()
	pool.AddCert(server.Certificate())
	var proxyCalls atomic.Int32
	var dialCalls atomic.Int32
	baseDialer := &net.Dialer{}
	transport := http.DefaultTransport.(*http.Transport).Clone()
	transport.TLSClientConfig = &tls.Config{RootCAs: pool, MinVersion: tls.VersionTLS12}
	transport.Proxy = func(*http.Request) (*url.URL, error) {
		proxyCalls.Add(1)
		return nil, nil
	}
	transport.DialContext = func(ctx context.Context, network, address string) (net.Conn, error) {
		dialCalls.Add(1)
		return baseDialer.DialContext(ctx, network, address)
	}
	protocol := &Protocol{httpClient: &http.Client{Transport: transport}}
	dial := requireGatewayWebSocketDialCapability(t, protocol)

	conn, response, err := dial.GatewayWebSocketDial(
		context.Background(),
		websocketURL(server.URL),
		nil,
		nil,
	)
	if response != nil && response.Body != nil {
		defer response.Body.Close()
	}
	if err != nil {
		t.Fatalf("GatewayWebSocketDial(TLS): %v", err)
	}
	if proxyCalls.Load() != 1 {
		t.Fatalf("Proxy calls = %d, want 1", proxyCalls.Load())
	}
	if dialCalls.Load() == 0 {
		t.Fatal("inherited DialContext was not called")
	}
	_ = conn.Close()
}

func TestGatewayWebSocketDialIsPerCallImmutableAndCompressionDisabled(t *testing.T) {
	type observed struct {
		protocols  string
		extensions string
		selected   string
	}
	seen := make(chan observed, 2)
	upgrader := websocket.Upgrader{
		CheckOrigin:       func(*http.Request) bool { return true },
		EnableCompression: false,
	}
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		requested := websocket.Subprotocols(r)
		selected := ""
		if len(requested) > 0 {
			selected = requested[0]
		}
		responseHeader := make(http.Header)
		if selected != "" {
			responseHeader.Set("Sec-WebSocket-Protocol", selected)
		}
		seen <- observed{
			protocols:  r.Header.Get("Sec-WebSocket-Protocol"),
			extensions: r.Header.Get("Sec-WebSocket-Extensions"),
			selected:   selected,
		}
		conn, err := upgrader.Upgrade(w, r, responseHeader)
		if err == nil {
			_ = conn.Close()
		}
	}))
	defer server.Close()

	protocol := &Protocol{
		httpClient: &http.Client{Transport: http.DefaultTransport.(*http.Transport).Clone()},
	}
	dial := requireGatewayWebSocketDialCapability(t, protocol)

	inputs := [][]string{{"proto-a"}, {"proto-b"}}
	var wg sync.WaitGroup
	errs := make(chan error, len(inputs))
	for _, protocols := range inputs {
		protocols := append([]string(nil), protocols...)
		wg.Add(1)
		go func() {
			defer wg.Done()
			conn, response, dialErr := dial.GatewayWebSocketDial(
				context.Background(),
				websocketURL(server.URL),
				nil,
				protocols,
			)
			if response != nil && response.Body != nil {
				_ = response.Body.Close()
			}
			if dialErr != nil {
				errs <- dialErr
				return
			}
			if conn.Subprotocol() != protocols[0] {
				errs <- errors.New("subprotocol leaked across calls")
			}
			_ = conn.Close()
		}()
	}
	wg.Wait()
	close(errs)
	for err := range errs {
		if err != nil {
			t.Fatal(err)
		}
	}

	got := map[string]observed{}
	for range inputs {
		value := <-seen
		got[value.selected] = value
	}
	for _, protocol := range []string{"proto-a", "proto-b"} {
		value, ok := got[protocol]
		if !ok {
			t.Fatalf("missing isolated request for %q: %#v", protocol, got)
		}
		if value.protocols != protocol {
			t.Fatalf("%s request protocols = %q", protocol, value.protocols)
		}
		if value.extensions != "" {
			t.Fatalf("%s negotiated extension request = %q, want empty", protocol, value.extensions)
		}
	}
}

func TestGatewayWebSocketDialFailsClosedForUnsupportedTransportPolicies(t *testing.T) {
	var networkDials atomic.Int32
	dial := func(context.Context, string, string) (net.Conn, error) {
		networkDials.Add(1)
		return nil, errors.New("network dial must not run")
	}
	nonEmptyHeader := http.Header{"Proxy-Authorization": {"secret"}}

	cases := []struct {
		name      string
		transport http.RoundTripper
	}{
		{name: "custom round tripper", transport: roundTripperFunc(func(*http.Request) (*http.Response, error) {
			return nil, errors.New("custom RT must not run")
		})},
		{name: "legacy dial", transport: func() *http.Transport {
			tr := http.DefaultTransport.(*http.Transport).Clone()
			tr.DialContext = nil
			tr.Dial = func(string, string) (net.Conn, error) { return nil, errors.New("must not run") }
			return tr
		}()},
		{name: "dial TLS context", transport: func() *http.Transport {
			tr := http.DefaultTransport.(*http.Transport).Clone()
			tr.DialContext = dial
			tr.DialTLSContext = func(context.Context, string, string) (net.Conn, error) {
				return nil, errors.New("must not run")
			}
			return tr
		}()},
		{name: "legacy dial TLS", transport: func() *http.Transport {
			tr := http.DefaultTransport.(*http.Transport).Clone()
			tr.DialContext = dial
			tr.DialTLS = func(string, string) (net.Conn, error) { return nil, errors.New("must not run") }
			return tr
		}()},
		{name: "proxy connect header", transport: func() *http.Transport {
			tr := http.DefaultTransport.(*http.Transport).Clone()
			tr.DialContext = dial
			tr.ProxyConnectHeader = nonEmptyHeader
			return tr
		}()},
		{name: "dynamic proxy connect header", transport: func() *http.Transport {
			tr := http.DefaultTransport.(*http.Transport).Clone()
			tr.DialContext = dial
			tr.GetProxyConnectHeader = func(context.Context, *url.URL, string) (http.Header, error) {
				return nil, nil
			}
			return tr
		}()},
		{name: "proxy connect response hook", transport: func() *http.Transport {
			tr := http.DefaultTransport.(*http.Transport).Clone()
			tr.DialContext = dial
			tr.OnProxyConnectResponse = func(context.Context, *url.URL, *http.Request, *http.Response) error {
				return nil
			}
			return tr
		}()},
		{name: "max conns per host", transport: func() *http.Transport {
			tr := http.DefaultTransport.(*http.Transport).Clone()
			tr.DialContext = dial
			tr.MaxConnsPerHost = 1
			return tr
		}()},
		{name: "max response header bytes", transport: func() *http.Transport {
			tr := http.DefaultTransport.(*http.Transport).Clone()
			tr.DialContext = dial
			tr.MaxResponseHeaderBytes = 1024
			return tr
		}()},
		{name: "http2 only protocols", transport: func() *http.Transport {
			tr := http.DefaultTransport.(*http.Transport).Clone()
			tr.DialContext = dial
			protocols := new(http.Protocols)
			protocols.SetHTTP2(true)
			protocols.SetHTTP1(false)
			tr.Protocols = protocols
			return tr
		}()},
	}

	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			var proxyCalls atomic.Int32
			if transport, ok := tc.transport.(*http.Transport); ok {
				transport.Proxy = func(*http.Request) (*url.URL, error) {
					proxyCalls.Add(1)
					return nil, nil
				}
			}
			protocol := &Protocol{httpClient: &http.Client{Transport: tc.transport}}
			capability := requireGatewayWebSocketDialCapability(t, protocol)
			before := networkDials.Load()
			conn, response, err := capability.GatewayWebSocketDial(
				context.Background(),
				"ws://127.0.0.1:1/socket",
				nil,
				nil,
			)
			if conn != nil {
				_ = conn.Close()
				t.Fatal("unsupported policy returned a connection")
			}
			if response != nil && response.Body != nil {
				_ = response.Body.Close()
			}
			if err == nil {
				t.Fatal("unsupported policy returned nil error")
			}
			if networkDials.Load() != before {
				t.Fatal("unsupported policy attempted a network dial")
			}
			if proxyCalls.Load() != 0 {
				t.Fatalf("unsupported policy called Proxy %d times before fail-closed validation", proxyCalls.Load())
			}
		})
	}
}

func TestGatewayWebSocketDialEffectiveHandshakeTimeoutExactPolicy(t *testing.T) {
	cases := []struct {
		name                  string
		tlsHandshakeTimeout   time.Duration
		responseHeaderTimeout time.Duration
		want                  time.Duration
	}{
		{
			name:                  "production-default",
			tlsHandshakeTimeout:   10 * time.Second,
			responseHeaderTimeout: 30 * time.Second,
			want:                  10 * time.Second,
		},
		{
			name:                  "zero-tls-uses-header-cap",
			responseHeaderTimeout: 30 * time.Second,
			want:                  30 * time.Second,
		},
		{
			name: "both-zero-use-absolute-cap",
			want: 30 * time.Second,
		},
		{
			name:                  "larger-values-are-capped",
			tlsHandshakeTimeout:   45 * time.Second,
			responseHeaderTimeout: time.Minute,
			want:                  30 * time.Second,
		},
		{
			name:                  "smaller-response-header-wins",
			tlsHandshakeTimeout:   10 * time.Second,
			responseHeaderTimeout: 3 * time.Second,
			want:                  3 * time.Second,
		},
	}
	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			transport := http.DefaultTransport.(*http.Transport).Clone()
			transport.TLSHandshakeTimeout = tc.tlsHandshakeTimeout
			transport.ResponseHeaderTimeout = tc.responseHeaderTimeout
			protocol := &Protocol{httpClient: &http.Client{Transport: transport}}
			policy, ok := any(protocol).(gatewayWebSocketHandshakePolicy)
			if !ok {
				t.Fatal("Protocol does not expose the same-package immutable WebSocket handshake timeout policy")
			}
			got, err := policy.gatewayWebSocketHandshakeTimeout()
			if err != nil {
				t.Fatalf("gatewayWebSocketHandshakeTimeout: %v", err)
			}
			if got != tc.want {
				t.Fatalf("effective handshake timeout = %v, want exactly %v", got, tc.want)
			}
		})
	}
}

func TestGatewayWebSocketDialEffectiveHandshakeTimeoutEnforcedWithinBounds(t *testing.T) {
	listener, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("Listen: %v", err)
	}
	defer listener.Close()
	accepted := make(chan net.Conn, 1)
	go func() {
		conn, acceptErr := listener.Accept()
		if acceptErr == nil {
			accepted <- conn
		}
	}()

	transport := http.DefaultTransport.(*http.Transport).Clone()
	transport.TLSHandshakeTimeout = 40 * time.Millisecond
	transport.ResponseHeaderTimeout = 200 * time.Millisecond
	protocol := &Protocol{httpClient: &http.Client{Transport: transport}}
	dial := requireGatewayWebSocketDialCapability(t, protocol)

	start := time.Now()
	conn, response, err := dial.GatewayWebSocketDial(
		context.Background(),
		"wss://"+listener.Addr().String()+"/socket",
		nil,
		nil,
	)
	if conn != nil {
		_ = conn.Close()
	}
	if response != nil && response.Body != nil {
		_ = response.Body.Close()
	}
	if err == nil {
		t.Fatal("stalled TLS handshake returned nil error")
	}
	if elapsed := time.Since(start); elapsed < 20*time.Millisecond || elapsed > 500*time.Millisecond {
		t.Fatalf("stalled handshake took %v, want inherited 40ms budget within [20ms,500ms]", elapsed)
	}
	select {
	case acceptedConn := <-accepted:
		_ = acceptedConn.Close()
	default:
	}
}

// roundTripperFunc avoids adding a production fake merely to prove custom
// RoundTrippers are rejected before invocation.
type roundTripperFunc func(*http.Request) (*http.Response, error)

func (f roundTripperFunc) RoundTrip(request *http.Request) (*http.Response, error) {
	return f(request)
}