package grpctool

import (
	"bufio"
	"bytes"
	"context"
	"fmt"
	"io"
	"net"
	"net/http"
	"time"

	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/httpz"
	"google.golang.org/protobuf/proto"
)

var (
	// See https://datatracker.ietf.org/doc/html/rfc9110#name-connection
	// See https://datatracker.ietf.org/doc/html/rfc2616#section-13.5.1
	// See https://github.com/golang/go/blob/81ea89adf38b90c3c3a8c4eed9e6c093a8634d59/src/net/http/httputil/reverseproxy.go#L169-L184
	// Must be in canonical form.
	hopHeaders = []string{
		httpz.ConnectionHeader,
		httpz.ProxyConnectionHeader,
		httpz.KeepAliveHeader,
		httpz.ProxyAuthenticateHeader,
		httpz.ProxyAuthorizationHeader,
		httpz.TeHeader,
		httpz.TrailerHeader,
		httpz.TransferEncodingHeader,
		httpz.UpgradeHeader,
	}
)

type MergeHeadersFunc func(outboundResponse, inboundResponse http.Header)
type WriteErrorResponse func(*ErrResp)
type TrackResponseFunc func(statusCode int)

type InboundHTTPToOutboundGRPC struct {
	HTTP2GRPC          HTTPToOutboundGRPC
	NewClient          func(context.Context) (HTTPRequestClient, error)
	WriteErrorResponse WriteErrorResponse
	MergeHeaders       MergeHeadersFunc
	TrackResponse      TrackResponseFunc
}

// Pipe pipes.
// headerExtra can be nil.
func (x *InboundHTTPToOutboundGRPC) Pipe(w http.ResponseWriter, r *http.Request, headerExtra proto.Message) bool /* abort */ {
	headerWritten, eResp := x.pipe(w, r, headerExtra)
	if eResp != nil {
		if headerWritten {
			// HTTP status has been written already as part of the normal response flow.
			// But then something went wrong and an error happened. To let the client know that something isn't right
			// we have only one thing we can do - abruptly close the connection.
			// We return a bool flag to indicate this.
			// The user of this method may panic with a special error value that the "http" package provides.
			// http.ErrAbortHandler -> See its description.
			// If we try to write the status again here, http package would log a warning, which is not nice.
			return true
		}

		x.WriteErrorResponse(eResp)
	}
	return false
}

func (x *InboundHTTPToOutboundGRPC) pipe(w http.ResponseWriter, r *http.Request, headerExtra proto.Message) (bool /* headerWritten */, *ErrResp) {
	// 0. Construct the outbound client
	ctx, cancel := context.WithCancelCause(r.Context())
	defer cancel(nil)
	outboundClient, err := x.NewClient(ctx)
	if err != nil {
		return false, &ErrResp{
			StatusCode: http.StatusInternalServerError,
			Msg:        "Proxy failed to make outbound request",
			Err:        err,
		}
	}

	// http.ResponseWriter does not support concurrent request body reads and response writes so
	// consume the request body first and then write the response from remote.
	// See https://github.com/golang/go/issues/15527
	// See https://github.com/golang/go/blob/go1.17.2/src/net/http/server.go#L118-L139

	// 1. Pipe client -> remote
	eResp := x.pipeInboundToOutbound(outboundClient, r, headerExtra)
	if eResp != nil {
		return false, eResp
	}
	isUpgrade := len(r.Header[httpz.UpgradeHeader]) > 0
	if !isUpgrade { // Close outbound connection for writes if it's not an upgraded connection
		eResp = x.HTTP2GRPC.SendCloseSend(outboundClient)
		if eResp != nil {
			return false, eResp
		}
	}
	// 2. Pipe remote -> client
	headerWritten, responseStatusCode, eResp := x.pipeOutboundToInbound(outboundClient, w, isUpgrade)
	if eResp != nil {
		return headerWritten, eResp
	}
	// 3. Pipe client <-> remote if connection upgrade is requested
	if !isUpgrade { // nothing to do
		return true, nil
	}
	if responseStatusCode != http.StatusSwitchingProtocols {
		// Remote doesn't want to upgrade the connection
		return true, nil
	}
	return true, x.pipeUpgradedConnection(outboundClient, w, cancel)
}

func (x *InboundHTTPToOutboundGRPC) pipeInboundToOutbound(outboundClient HTTPRequestClient, r *http.Request, headerExtra proto.Message) *ErrResp {
	rx := *r                         // shallow clone to avoid mutating the original's header
	header := rx.Header.Clone()      // can be nil
	delete(header, httpz.HostHeader) // Use the destination host name
	cleanHeader(header)
	rx.Header = header

	return x.HTTP2GRPC.PipeInboundToOutbound(outboundClient, &rx, headerExtra)
}

func (x *InboundHTTPToOutboundGRPC) pipeOutboundToInbound(outboundClient HTTPRequestClient, w http.ResponseWriter, isUpgrade bool) (bool /* headerWritten */, int32, *ErrResp) {
	wr := &rw2rw{
		w:              w,
		outboundClient: outboundClient,
		mergeHeaders:   x.MergeHeaders,
		isUpgrade:      isUpgrade,
		trackResponse:  x.TrackResponse,
	}
	eResp := x.HTTP2GRPC.PipeOutboundToInbound(outboundClient, wr) // keep it on a separate line as we are reading the fields below.
	return wr.headerWritten, wr.responseStatusCode, eResp
}

func (x *InboundHTTPToOutboundGRPC) pipeUpgradedConnection(outboundClient HTTPRequestClient, w http.ResponseWriter,
	cancelOutbound context.CancelCauseFunc) (errRet *ErrResp) {
	// Connection upgrade requested. For that ResponseWriter must support hijacking.
	conn, bufrw, err := http.NewResponseController(w).Hijack()
	if err != nil {
		return x.HTTP2GRPC.HandleInternalError("unable to upgrade connection: error hijacking response", err)
	}
	defer func() {
		err = conn.Close()
		if err != nil && errRet == nil {
			errRet = x.HTTP2GRPC.HandleIOError("failed to close upgraded connection", err)
		}
	}()
	// Hijack() docs say we are responsible for managing connection deadlines and a deadline may be set already.
	// We clear the read deadline here because we don't know if the client will be sending any data to us soon.
	err = conn.SetReadDeadline(time.Time{})
	if err != nil {
		return x.HTTP2GRPC.HandleIOError("failed to clear connection read deadline", err)
	}
	// We don't care if a write deadline is set already, we just wrap the connection in a wrapper that
	// will each time set a new deadline before performing an actual write.
	conn = &httpz.WriteTimeoutConn{
		Conn:    conn,
		Timeout: 20 * time.Second,
	}
	r, err := decoupleReader(bufrw.Reader, conn)
	if err != nil {
		return x.HTTP2GRPC.HandleIOError("failed to read buffered data", err)
	}
	p := InboundStreamToOutboundStream{
		PipeInboundToOutbound: func() error {
			in2outErr := x.HTTP2GRPC.PipeInboundToOutboundUpgraded(outboundClient, r)
			if in2outErr != nil {
				// Hijacked connections do not cancel the request context when they break.
				// See https://github.com/golang/go/blob/go1.23.1/src/net/http/server.go#L203-L205.
				// We have to manually unblock the outbound client since it relies on context cancellation to be unblocked.
				cancelOutbound(in2outErr)
			}
			return in2outErr
		},
		PipeOutboundToInbound: func() error {
			return x.HTTP2GRPC.PipeOutboundToInboundUpgraded(outboundClient, conn)
		},
	}
	err = p.Pipe()
	if err != nil {
		return x.HTTP2GRPC.HandleIOError("failed to pipe upgraded connection streams", err)
	}
	return nil
}

// decoupleReader returns an io.Reader that is decoupled from the buffered reader, returned by Hijack.
// This is a workaround for https://github.com/golang/go/issues/32314.
func decoupleReader(r *bufio.Reader, conn net.Conn) (io.Reader, error) {
	buffered := r.Buffered()
	if buffered == 0 {
		return conn, nil
	}
	b, err := r.Peek(buffered)
	if err != nil {
		return nil, err
	}
	return io.MultiReader(bytes.NewReader(b), conn), nil
}

// header may be nil.
func cleanHeader(header http.Header) {
	upgrade := header[httpz.UpgradeHeader]

	// 1. Remove hop-by-hop headers listed in the Connection header. See https://datatracker.ietf.org/doc/html/rfc7230#section-6.1
	httpz.RemoveConnectionHeaders(header)
	// 2. Remove well-known hop-by-hop headers
	for _, name := range hopHeaders {
		delete(header, name)
	}
	// 3. Fix up Connection and Upgrade headers if upgrade is requested/confirmed
	if len(upgrade) > 0 {
		header[httpz.UpgradeHeader] = upgrade                // put it back
		header[httpz.ConnectionHeader] = []string{"upgrade"} // this discards any other connection options if they were there
	}
}

var (
	_ ResponseWriter = (*rw2rw)(nil)
)

// rw2rw is a ResponseWriter -> http.ResponseWriter adapter.
type rw2rw struct {
	w                  http.ResponseWriter
	outboundClient     HTTPRequestClient
	mergeHeaders       MergeHeadersFunc
	trackResponse      TrackResponseFunc
	responseStatusCode int32
	isUpgrade          bool
	headerWritten      bool
}

func (r *rw2rw) Write(data []byte) (int, error) {
	n, err := r.w.Write(data)
	if err == nil { // flush when there was no error
		// http.ResponseWriter buffers headers and response body writes and that may break use cases like long polling or streaming.
		// Flush() is used so that when HTTP headers and response body chunks are received from the outbound connection,
		// they are flushed to the inbound stream ASAP.
		err = http.NewResponseController(r.w).Flush()
	}
	return n, err
}

func (r *rw2rw) WriteHeader(statusCode int32, status string, header http.Header) (bool, error) {
	successfulUpgrade := false

	if r.isUpgrade {
		if statusCode == http.StatusSwitchingProtocols {
			successfulUpgrade = true
		} else {
			// We were trying to upgrade but the server doesn't want to switch protocols.
			// Hence, we need to close the outbound client stream as we are not going to be sending
			// any traffic via the upgraded connection.
			// This would allow server to finish reading from the incoming stream and unblock.
			// Then the server would return, closing the stream the caller of this method is reading - that visitor will unblock.
			err := r.outboundClient.CloseSend()
			if err != nil {
				return false, fmt.Errorf("failed to send close frame: %w", err)
			}
		}
	}

	cleanHeader(header)
	r.mergeHeaders(header, r.w.Header())
	r.w.WriteHeader(int(statusCode))
	r.responseStatusCode = statusCode
	r.headerWritten = true

	// NOTE: the HTTP standard library doesn't no-op for a flush when WriteHeader() was already called with a 1xx status code
	// and already flushed. This leads to the response being sent twice once with the correct status code and once with `200 OK`.
	// Thus, we avoid flushing manually for all 1xx responses.
	// This seems to be a regression in Go 1.19, introduced with https://go-review.googlesource.com/c/go/+/269997
	var err error
	if statusCode >= 200 {
		err = http.NewResponseController(r.w).Flush()
	}
	r.trackResponse(int(statusCode))
	return successfulUpgrade, err
}
