// Copyright 2021-2022 Buf Technologies, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package connect

import (
	"context"
	"errors"
	"fmt"
	"io"
	"mime"
	"net/http"
	"net/url"
	"sort"
	"strings"
)

// The names of the Connect, gRPC, and gRPC-Web protocols (as exposed by
// [Peer.Protocol]). Additional protocols may be added in the future.
const (
	ProtocolConnect = "connect"
	ProtocolGRPC    = "grpc"
	ProtocolGRPCWeb = "grpcweb"
)

const (
	headerContentType = "Content-Type"
	headerUserAgent   = "User-Agent"

	discardLimit = 1024 * 1024 * 4 // 4MiB
)

var errNoTimeout = errors.New("no timeout")

// A Protocol defines the HTTP semantics to use when sending and receiving
// messages. It ties together codecs, compressors, and net/http to produce
// Senders and Receivers.
//
// For example, connect supports the gRPC protocol using this abstraction. Among
// many other things, the protocol implementation is responsible for
// translating timeouts from Go contexts to HTTP and vice versa. For gRPC, it
// converts timeouts to and from strings (for example, 10*time.Second <->
// "10S"), and puts those strings into the "Grpc-Timeout" HTTP header. Other
// protocols might encode durations differently, put them into a different HTTP
// header, or ignore them entirely.
//
// We don't have any short-term plans to export this interface; it's just here
// to separate the protocol-specific portions of connect from the
// protocol-agnostic plumbing.
type protocol interface {
	NewHandler(*protocolHandlerParams) protocolHandler
	NewClient(*protocolClientParams) (protocolClient, error)
}

// HandlerParams are the arguments provided to a Protocol's NewHandler
// method, bundled into a struct to allow backward-compatible argument
// additions. Protocol implementations should take care to use the supplied
// Spec rather than constructing their own, since new fields may have been
// added.
type protocolHandlerParams struct {
	Spec                         Spec
	Codecs                       readOnlyCodecs
	CompressionPools             readOnlyCompressionPools
	CompressMinBytes             int
	BufferPool                   *bufferPool
	ReadMaxBytes                 int
	SendMaxBytes                 int
	RequireConnectProtocolHeader bool
}

// Handler is the server side of a protocol. HTTP handlers typically support
// multiple protocols, codecs, and compressors.
type protocolHandler interface {
	// ContentTypes is the set of HTTP Content-Types that the protocol can
	// handle.
	ContentTypes() map[string]struct{}

	// SetTimeout runs before NewStream. Implementations may inspect the HTTP
	// request, parse any timeout set by the client, and return a modified
	// context and cancellation function.
	//
	// If the client didn't send a timeout, SetTimeout should return the
	// request's context, a nil cancellation function, and a nil error.
	SetTimeout(*http.Request) (context.Context, context.CancelFunc, error)

	// NewConn constructs a HandlerConn for the message exchange.
	NewConn(http.ResponseWriter, *http.Request) (handlerConnCloser, bool)
}

// ClientParams are the arguments provided to a Protocol's NewClient method,
// bundled into a struct to allow backward-compatible argument additions.
// Protocol implementations should take care to use the supplied Spec rather
// than constructing their own, since new fields may have been added.
type protocolClientParams struct {
	CompressionName  string
	CompressionPools readOnlyCompressionPools
	Codec            Codec
	CompressMinBytes int
	HTTPClient       HTTPClient
	URL              string
	BufferPool       *bufferPool
	ReadMaxBytes     int
	SendMaxBytes     int
	// The gRPC family of protocols always needs access to a Protobuf codec to
	// marshal and unmarshal errors.
	Protobuf Codec
}

// Client is the client side of a protocol. HTTP clients typically use a single
// protocol, codec, and compressor to send requests.
type protocolClient interface {
	// Peer describes the server for the RPC.
	Peer() Peer

	// WriteRequestHeader writes any protocol-specific request headers.
	WriteRequestHeader(StreamType, http.Header)

	// NewConn constructs a StreamingClientConn for the message exchange.
	//
	// Implementations should assume that the supplied HTTP headers have already
	// been populated by WriteRequestHeader. When constructing a stream for a
	// unary call, implementations may assume that the Sender's Send and Close
	// methods return before the Receiver's Receive or Close methods are called.
	NewConn(context.Context, Spec, http.Header) StreamingClientConn
}

// errorTranslatingHandlerConnCloser wraps a handlerConnCloser to ensure that
// we always return coded errors to users and write coded errors to the
// network.
//
// It's used in protocol implementations.
type errorTranslatingHandlerConnCloser struct {
	handlerConnCloser

	toWire   func(error) error
	fromWire func(error) error
}

func (hc *errorTranslatingHandlerConnCloser) Send(msg any) error {
	return hc.fromWire(hc.handlerConnCloser.Send(msg))
}

func (hc *errorTranslatingHandlerConnCloser) Receive(msg any) error {
	return hc.fromWire(hc.handlerConnCloser.Receive(msg))
}

func (hc *errorTranslatingHandlerConnCloser) Close(err error) error {
	closeErr := hc.handlerConnCloser.Close(hc.toWire(err))
	return hc.fromWire(closeErr)
}

// errorTranslatingClientConn wraps a StreamingClientConn to make sure that we always
// return coded errors from clients.
//
// It's used in protocol implementations.
type errorTranslatingClientConn struct {
	StreamingClientConn

	fromWire func(error) error
}

func (cc *errorTranslatingClientConn) Send(msg any) error {
	return cc.fromWire(cc.StreamingClientConn.Send(msg))
}

func (cc *errorTranslatingClientConn) Receive(msg any) error {
	return cc.fromWire(cc.StreamingClientConn.Receive(msg))
}

func (cc *errorTranslatingClientConn) CloseRequest() error {
	return cc.fromWire(cc.StreamingClientConn.CloseRequest())
}

func (cc *errorTranslatingClientConn) CloseResponse() error {
	return cc.fromWire(cc.StreamingClientConn.CloseResponse())
}

// wrapHandlerConnWithCodedErrors ensures that we (1) automatically code
// context-related errors correctly when writing them to the network, and (2)
// return *Errors from all exported APIs.
func wrapHandlerConnWithCodedErrors(conn handlerConnCloser) handlerConnCloser {
	return &errorTranslatingHandlerConnCloser{
		handlerConnCloser: conn,
		toWire:            wrapIfContextError,
		fromWire:          wrapIfUncoded,
	}
}

// wrapClientConnWithCodedErrors ensures that we always return *Errors from
// public APIs.
func wrapClientConnWithCodedErrors(conn StreamingClientConn) StreamingClientConn {
	return &errorTranslatingClientConn{
		StreamingClientConn: conn,
		fromWire:            wrapIfUncoded,
	}
}

func sortedAcceptPostValue(handlers []protocolHandler) string {
	contentTypes := make(map[string]struct{})
	for _, handler := range handlers {
		for contentType := range handler.ContentTypes() {
			contentTypes[contentType] = struct{}{}
		}
	}
	accept := make([]string, 0, len(contentTypes))
	for ct := range contentTypes {
		accept = append(accept, ct)
	}
	sort.Strings(accept)
	return strings.Join(accept, ", ")
}

func isCommaOrSpace(c rune) bool {
	return c == ',' || c == ' '
}

func discard(reader io.Reader) error {
	if lr, ok := reader.(*io.LimitedReader); ok {
		_, err := io.Copy(io.Discard, lr)
		return err
	}
	// We don't want to get stuck throwing data away forever, so limit how much
	// we're willing to do here.
	lr := &io.LimitedReader{R: reader, N: discardLimit}
	_, err := io.Copy(io.Discard, lr)
	return err
}

func validateRequestURL(uri string) *Error {
	_, err := url.ParseRequestURI(uri)
	if err == nil {
		return nil
	}
	if !strings.Contains(uri, "://") {
		// URL doesn't have a scheme, so the user is likely accustomed to
		// grpc-go's APIs.
		err = fmt.Errorf(
			"URL %q missing scheme: use http:// or https:// (unlike grpc-go)",
			uri,
		)
	}
	return NewError(CodeUnavailable, err)
}

// negotiateCompression determines and validates the request compression and
// response compression using the available compressors and protocol-specific
// Content-Encoding and Accept-Encoding headers.
func negotiateCompression( //nolint:nonamedreturns
	availableCompressors readOnlyCompressionPools,
	sent, accept string,
) (requestCompression, responseCompression string, clientVisibleErr *Error) {
	requestCompression = compressionIdentity
	if sent != "" && sent != compressionIdentity {
		// We default to identity, so we only care if the client sends something
		// other than the empty string or compressIdentity.
		if availableCompressors.Contains(sent) {
			requestCompression = sent
		} else {
			// To comply with
			// https://github.com/grpc/grpc/blob/master/doc/compression.md and the
			// Connect protocol, we should return CodeUnimplemented and specify
			// acceptable compression(s) (in addition to setting the a
			// protocol-specific accept-encoding header).
			return "", "", errorf(
				CodeUnimplemented,
				"unknown compression %q: supported encodings are %v",
				sent, availableCompressors.CommaSeparatedNames(),
			)
		}
	}
	// Support asymmetric compression. This logic follows
	// https://github.com/grpc/grpc/blob/master/doc/compression.md and common
	// sense.
	responseCompression = requestCompression
	// If we're not already planning to compress the response, check whether the
	// client requested a compression algorithm we support.
	if responseCompression == compressionIdentity && accept != "" {
		for _, name := range strings.FieldsFunc(accept, isCommaOrSpace) {
			if availableCompressors.Contains(name) {
				// We found a mutually supported compression algorithm. Unlike standard
				// HTTP, there's no preference weighting, so can bail out immediately.
				responseCompression = name
				break
			}
		}
	}
	return requestCompression, responseCompression, nil
}

// checkServerStreamsCanFlush ensures that bidi and server streaming handlers
// have received an http.ResponseWriter that implements http.Flusher, since
// they must flush data after sending each message.
func checkServerStreamsCanFlush(spec Spec, responseWriter http.ResponseWriter) *Error {
	requiresFlusher := (spec.StreamType & StreamTypeServer) == StreamTypeServer
	if _, flushable := responseWriter.(http.Flusher); requiresFlusher && !flushable {
		return NewError(CodeInternal, fmt.Errorf("%T does not implement http.Flusher", responseWriter))
	}
	return nil
}

func flushResponseWriter(w http.ResponseWriter) {
	if f, ok := w.(http.Flusher); ok {
		f.Flush()
	}
}

func canonicalizeContentType(ct string) string {
	base, params, err := mime.ParseMediaType(ct)
	if err != nil {
		return ct
	}
	return mime.FormatMediaType(base, params)
}
