package grpctool

import (
	"errors"
	"fmt"
	"io"
	"log/slog"
	"net/http"

	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/logz"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/memz"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/prototool"
	"gitlab.com/gitlab-org/cluster-integration/tunnel/tool/streamvisitor"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/protobuf/proto"
	"google.golang.org/protobuf/types/known/anypb"
)

var (
	// errEarlyExit is a sentinel error value to make stream visitor exit early.
	errEarlyExit = errors.New("early exit sentinel error")
)

type ResponseWriter interface {
	io.Writer
	WriteHeader(statusCode int32, status string, header http.Header) (successfulUpgrade bool, hasContentTypeHeader bool, err error)
}
type HTTPRequestClient = grpc.BidiStreamingClient[HttpRequest, HttpResponse]

// CheckHeader checks the status code and header of the upstream server's HTTP response.
// It may return an error or alter the header.
type CheckHeader func(statusCode int32, header http.Header) error

type ErrResp struct {
	StatusCode int32
	Msg        string
	// Err can be nil.
	Err error
}

type HTTPToOutboundGRPC struct {
	Log                       *slog.Logger
	HandleProcessingErrorFunc HandleProcessingErrorFunc
	CheckHeader               CheckHeader
}

func (x *HTTPToOutboundGRPC) PipeOutboundToInbound(outboundClient HTTPRequestClient, w ResponseWriter) *ErrResp {
	writeFailed := false
	checkHeaderFailed := false
	successfulUpgrade := false
	hasContentTypeHeader := false
	err := HTTPResponseStreamVisitor().Visit(outboundClient,
		streamvisitor.WithCallback(HTTPResponseHeaderFieldNumber, func(header *HttpResponse_Header) error {
			responseStatusCode := header.Response.StatusCode
			outboundResponse := header.Response.HTTPHeader()
			err := x.CheckHeader(responseStatusCode, outboundResponse)
			if err != nil {
				checkHeaderFailed = true
				return err
			}
			successfulUpgrade, hasContentTypeHeader, err = w.WriteHeader(responseStatusCode, header.Response.Status, outboundResponse)
			if err != nil {
				writeFailed = true
				return err
			}
			return nil
		}),
		streamvisitor.WithCallback(HTTPResponseDataFieldNumber, func(data *HttpResponse_Data) error {
			if !hasContentTypeHeader {
				checkHeaderFailed = true
				return errors.New("response is missing Content-Type header, but attempted to send content. Dropping connection due to security concerns")
			}

			_, err := w.Write(data.Data)
			if err != nil {
				writeFailed = true
				return err
			}
			return nil
		}),
		streamvisitor.WithCallback(HTTPResponseTrailerFieldNumber, func(trailer *HttpResponse_Trailer) error {
			if successfulUpgrade {
				return errEarlyExit
			}
			return nil
		}),
		// if it's a successful upgrade, then this field is unreachable because of the early exit above.
		// otherwise, (unsuccessful upgrade or not an upgrade) the remote must not send this field.
		streamvisitor.WithNotExpectingToGet(codes.Internal, HTTPResponseUpgradeDataFieldNumber),
	)
	if err != nil {
		switch {
		case err == errEarlyExit: //nolint: errorlint
			// Return no error.
		case writeFailed:
			// there is likely a connection problem so the client will likely not receive this
			return x.HandleIOError("failed to write HTTP response", err)
		case checkHeaderFailed:
			return x.HandleProcessingError(http.StatusBadGateway, "bad server response", err)
		default:
			return x.HandleIOError("failed to read gRPC response", err)
		}
	}
	return nil
}

func (x *HTTPToOutboundGRPC) PipeInboundToOutbound(outboundClient HTTPRequestClient, r *http.Request, headerExtra proto.Message) *ErrResp {
	var extra *anypb.Any
	if headerExtra != nil {
		var err error
		extra, err = anypb.New(headerExtra)
		if err != nil {
			return x.HandleInternalError("failed to marshal header extra proto", err)
		}
	}
	eResp := x.send(outboundClient, "failed to send request header", &HttpRequest{
		Message: &HttpRequest_Header_{
			Header: &HttpRequest_Header{
				Request: &prototool.HttpRequest{
					Method:  r.Method,
					Header:  prototool.HTTPHeaderToHeaderKV(r.Header),
					UrlPath: r.URL.Path,
					Query:   prototool.URLValuesToQueryKV(r.URL.Query()),
				},
				Extra:         extra,
				ContentLength: &r.ContentLength,
			},
		},
	})
	if eResp != nil {
		return eResp
	}

	eResp = x.sendRequestBody(outboundClient, r.Body)
	if eResp != nil {
		return eResp
	}
	return x.send(outboundClient, "failed to send trailer", &HttpRequest{
		Message: &HttpRequest_Trailer_{
			Trailer: &HttpRequest_Trailer{},
		},
	})
}

func (x *HTTPToOutboundGRPC) sendRequestBody(outboundClient HTTPRequestClient, body io.Reader) *ErrResp {
	if body == nil {
		// http.Request#Body can be nil for outbound GET/HEAD requests.
		return nil
	}
	bp := memz.Get32k()
	defer memz.Put32k(bp)
	buffer := *bp
	for {
		n, readErr := body.Read(buffer)
		if n > 0 { // handle n>0 before readErr != nil to ensure any consumed data gets forwarded
			eResp := x.send(outboundClient, "failed to send request data", &HttpRequest{
				Message: &HttpRequest_Data_{
					Data: &HttpRequest_Data{
						Data: buffer[:n],
					},
				},
			})
			if eResp != nil {
				return eResp
			}
		}
		if readErr != nil {
			if readErr == io.EOF {
				break
			}
			// There is likely a connection problem so the client will likely not receive this
			return x.HandleIOError("failed to read request body", readErr)
		}
	}
	return nil
}

func (x *HTTPToOutboundGRPC) PipeInboundToOutboundUpgraded(outboundClient HTTPRequestClient, inboundStream io.Reader) error {
	bp := memz.Get32k()
	defer memz.Put32k(bp)
	buffer := *bp
	for {
		n, readErr := inboundStream.Read(buffer)
		if n > 0 { // handle n>0 before readErr != nil to ensure any consumed data gets forwarded
			sendErr := outboundClient.Send(&HttpRequest{
				Message: &HttpRequest_UpgradeData_{
					UpgradeData: &HttpRequest_UpgradeData{
						Data: buffer[:n],
					},
				},
			})
			if sendErr != nil {
				if readErr == io.EOF {
					return nil // the other goroutine will receive the error in RecvMsg()
				}
				return fmt.Errorf("Send(HttpRequest_UpgradeData): %w", sendErr)
			}
		}
		if readErr != nil {
			if readErr == io.EOF {
				break
			}
			// There is likely a connection problem so the client will likely not receive this
			return fmt.Errorf("read failed: %w", readErr)
		}
	}
	_ = outboundClient.CloseSend() // never returns an error
	return nil
}

func (x *HTTPToOutboundGRPC) PipeOutboundToInboundUpgraded(outboundClient HTTPRequestClient, inboundStream io.Writer) error {
	var writeFailed bool
	err := HTTPResponseStreamVisitor().Visit(outboundClient,
		streamvisitor.WithStartState(HTTPResponseTrailerFieldNumber),
		streamvisitor.WithCallback(HTTPResponseUpgradeDataFieldNumber, func(data *HttpResponse_UpgradeData) error {
			_, err := inboundStream.Write(data.Data)
			if err != nil {
				writeFailed = true
			}
			return err
		}),
	)
	if err != nil {
		if writeFailed {
			// there is likely a connection problem so the client will likely not receive this
			return fmt.Errorf("failed to write upgraded HTTP response: %w", err)
		}
		return fmt.Errorf("failed to read upgraded gRPC response: %w", err)
	}
	return nil
}

func (x *HTTPToOutboundGRPC) SendCloseSend(outboundClient HTTPRequestClient) {
	_ = outboundClient.CloseSend() // never returns an error
}

func (x *HTTPToOutboundGRPC) send(client HTTPRequestClient, errMsg string, msg *HttpRequest) *ErrResp {
	err := client.Send(msg)
	if err != nil {
		if err == io.EOF { //nolint:errorlint
			_, err = client.Recv()
		}
		return x.HandleIOError(errMsg, err)
	}
	return nil
}

func (x *HTTPToOutboundGRPC) HandleIOError(msg string, err error) *ErrResp {
	msg = "HTTP->gRPC: " + msg
	x.Log.Debug(msg, logz.Error(err))
	return &ErrResp{
		// See https://datatracker.ietf.org/doc/html/rfc7231#section-6.6.3
		StatusCode: http.StatusBadGateway,
		Msg:        msg,
		Err:        err,
	}
}

func (x *HTTPToOutboundGRPC) HandleInternalError(msg string, err error) *ErrResp {
	// See https://datatracker.ietf.org/doc/html/rfc7231#section-6.6.1
	return x.HandleProcessingError(http.StatusInternalServerError, msg, err)
}

func (x *HTTPToOutboundGRPC) HandleProcessingError(statusCode int32, msg string, err error) *ErrResp {
	msg = "HTTP->gRPC: " + msg
	x.HandleProcessingErrorFunc(msg, err)
	return &ErrResp{
		StatusCode: statusCode,
		Msg:        msg,
		Err:        err,
	}
}
