package tunserver

import (
	"io"
	"log/slog"

	grpc_middleware "github.com/grpc-ecosystem/go-grpc-middleware/v2"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/module/modshared"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/grpctool"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tunnel/rpc"
	statuspb "google.golang.org/genproto/googleapis/rpc/status"
	"google.golang.org/grpc"
	"google.golang.org/grpc/metadata"
)

func (r *Router[T]) routeToTunclient(srv any, stream grpc.ServerStream) error { //nolint:unparam
	ctx := stream.Context()
	rpcAPI := modshared.RPCAPIFromContext[modshared.RPCAPI](ctx)
	tunnelFound, log, findHandle, err := r.Plugin.FindTunnel(stream, rpcAPI)
	if err != nil {
		return err
	}
	defer findHandle.Done(ctx)

	if !tunnelFound {
		err = stream.SendMsg(&rpc.GatewayResponse{
			Msg: &rpc.GatewayResponse_NoTunnel_{
				NoTunnel: &rpc.GatewayResponse_NoTunnel{},
			},
		})
		if err != nil {
			return rpcAPI.HandleIOError(log, "SendMsg(GatewayResponse_NoTunnel) failed", err)
		}
	}
	ageCtx := grpctool.MaxConnectionAgeContextFromStreamContext(ctx)
	tun, err := findHandle.Get(ageCtx)
	if err != nil {
		if grpctool.RequestCanceledOrTimedOut(err) {
			// Gateway finder handles clean early return.
			return nil
		}
		return err
	}
	defer tun.Done(ctx)
	err = stream.SendMsg(&rpc.GatewayResponse{
		Msg: &rpc.GatewayResponse_TunnelReady_{
			TunnelReady: &rpc.GatewayResponse_TunnelReady{},
		},
	})
	if err != nil {
		return rpcAPI.HandleIOError(log, "SendMsg(GatewayResponse_TunnelReady) failed", err)
	}
	var start rpc.StartStreaming
	err = stream.RecvMsg(&start)
	if err != nil {
		if err == io.EOF { //nolint:errorlint
			// Routing kas decided not to proceed
			return nil
		}
		return err
	}
	stream = removeReservedHeaders(stream)
	stream, err = r.Plugin.PrepareStreamForForwarding(stream)
	if err != nil {
		return err
	}
	return tun.ForwardStream(log, rpcAPI, stream, newWrappingCallback(log, rpcAPI, stream))
}

var (
	// MUST be lowercase.
	reservedHeaders = []string{
		// Contains the authentication token from the routing tunserver.
		"authorization",

		// gRPC-specific headers.
		":authority",
		"content-type",
		"grpc-accept-encoding",
	}
)

func removeReservedHeaders(stream grpc.ServerStream) grpc.ServerStream {
	md, _ := metadata.FromIncomingContext(stream.Context())
	for _, header := range reservedHeaders {
		delete(md, header)
	}

	// Overwrite incoming MD with sanitized MD
	wrappedStream := grpc_middleware.WrapServerStream(stream)
	wrappedStream.WrappedContext = metadata.NewIncomingContext(
		wrappedStream.WrappedContext,
		md,
	)
	return wrappedStream
}

type wrappingCallback struct {
	log    *slog.Logger
	rpcAPI modshared.RPCAPI
	stream grpc.ServerStream
}

func newWrappingCallback(log *slog.Logger, rpcAPI modshared.RPCAPI, stream grpc.ServerStream) *wrappingCallback {
	return &wrappingCallback{
		log:    log,
		rpcAPI: rpcAPI,
		stream: stream,
	}
}

func (c *wrappingCallback) Header(md []*rpc.MetadataKV) error {
	return c.sendMsg("SendMsg(GatewayResponse_Header) failed", &rpc.GatewayResponse{
		Msg: &rpc.GatewayResponse_Header_{
			Header: &rpc.GatewayResponse_Header{
				Meta: md,
			},
		},
	})
}

func (c *wrappingCallback) Message(data []byte) error {
	return c.sendMsg("SendMsg(GatewayResponse_Message) failed", &rpc.GatewayResponse{
		Msg: &rpc.GatewayResponse_Message_{
			Message: &rpc.GatewayResponse_Message{
				Data: data,
			},
		},
	})
}

func (c *wrappingCallback) Trailer(md []*rpc.MetadataKV) error {
	return c.sendMsg("SendMsg(GatewayResponse_Trailer) failed", &rpc.GatewayResponse{
		Msg: &rpc.GatewayResponse_Trailer_{
			Trailer: &rpc.GatewayResponse_Trailer{
				Meta: md,
			},
		},
	})
}

func (c *wrappingCallback) Error(stat *statuspb.Status) error {
	return c.sendMsg("SendMsg(GatewayResponse_Error) failed", &rpc.GatewayResponse{
		Msg: &rpc.GatewayResponse_Error_{
			Error: &rpc.GatewayResponse_Error{
				Status: stat,
			},
		},
	})
}

func (c *wrappingCallback) sendMsg(errMsg string, msg *rpc.GatewayResponse) error {
	err := c.stream.SendMsg(msg)
	if err != nil {
		return c.rpcAPI.HandleIOError(c.log, errMsg, err)
	}
	return nil
}
