package server

import (
	"context"
	"log/slog"
	"maps"
	"net/http"

	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/api"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/gitlab"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/module/workspaces/rpc"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/grpctool"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/httpz"
	"google.golang.org/grpc"
)

type tunnelHandlerInterface interface {
	tunnel(w http.ResponseWriter, r *http.Request)
}

type tunnelHandler struct {
	log                   *slog.Logger
	gitLabClient          gitlab.ClientInterface
	agentConnPool         func(agentKey api.AgentKey) grpc.ClientConnInterface
	handleProcessingError func(string, error)
	serverName            string
	serverVia             string
}

func newTunnelHandler(log *slog.Logger, gitLabClient gitlab.ClientInterface, agentConnPool func(agentKey api.AgentKey) grpc.ClientConnInterface, handleProcessingError func(string, error), serverName, serverVia string) *tunnelHandler {
	return &tunnelHandler{
		log:                   log,
		gitLabClient:          gitLabClient,
		agentConnPool:         agentConnPool,
		handleProcessingError: handleProcessingError,
		serverName:            serverName,
		serverVia:             serverVia,
	}
}

// tunnel handles all workspaces traffic and tunnels them into the appropriate agentw
func (h *tunnelHandler) tunnel(w http.ResponseWriter, r *http.Request) {
	workspace, err := getWorkspaceFromContext(r.Context())
	if err != nil {
		h.handleProcessingError("Failed to get workspace from context", err)
		http.Error(w, "Unauthorized", http.StatusUnauthorized)
		return
	}
	agentKey := api.AgentKey{
		Type: api.AgentTypeWorkspace,
		ID:   workspace.Id,
	}

	r.Header[httpz.ServerHeader] = []string{h.serverName} // It will be removed just before responding with actual headers from upstream
	r.Header[httpz.ViaHeader] = append(r.Header[httpz.ViaHeader], h.serverVia)

	cookies := r.Cookies()
	r.Header.Del("Cookie")
	for _, cookie := range cookies {
		if cookie.Name != userSessionCookieName { // Keep other cookies as is
			r.AddCookie(cookie)
		}
	}

	http2grpc := grpctool.InboundHTTPToOutboundGRPC{
		HTTP2GRPC: grpctool.HTTPToOutboundGRPC{
			Log:                       h.log,
			HandleProcessingErrorFunc: h.handleProcessingError,
			CheckHeader:               func(statusCode int32, header http.Header) error { return nil },
		},
		NewClient: func(ctx context.Context) (grpctool.HTTPRequestClient, error) {
			return rpc.NewWorkspacesClient(h.agentConnPool(agentKey)).TunnelHTTP(ctx, grpc.WaitForReady(true))
		},
		WriteErrorResponse: func(errResp *grpctool.ErrResp) {
			// errResp.Err.Error() contains the underlying message (e.g. connection refused)
			http.Error(w, errResp.Err.Error(), int(errResp.StatusCode))
		},
		MergeHeaders: func(outboundResponse, inboundResponse http.Header) {
			delete(inboundResponse, httpz.ServerHeader) // remove the header we've added above. We use Via instead.

			// Since the user is accessing their own server in the workspace, there is nothing to protect against.
			// Hence, we set the headers from the proxied response as is and
			// do not attempt to overrule anything (e.g. CORS).
			maps.Copy(inboundResponse, outboundResponse)

			inboundResponse[httpz.ViaHeader] = append(inboundResponse[httpz.ViaHeader], h.serverVia)
		},
		TrackResponse: func(statusCode int) {},
	}
	headerExtra := &rpc.HTTPHeaderExtra{
		Port: workspace.Port,
	}
	abort := http2grpc.Pipe(w, r, headerExtra)
	if abort {
		panic(http.ErrAbortHandler)
	}
}
