package tunserver

import (
	"context"
	"io"
	"log/slog"
	"net/url"

	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/module/modshared"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/logz"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/metadata"
	"google.golang.org/grpc/status"
)

var (
	_ grpc.ClientConnInterface = (*RoutingClientConn[*url.URL])(nil)
)

type RoutingClientConn[T Target] struct {
	Log    *slog.Logger
	API    modshared.API
	Plugin RouterPlugin[T]
}

func (c *RoutingClientConn[T]) Invoke(ctx context.Context, method string, args, reply any, opts ...grpc.CallOption) error {
	ctx, cancel := context.WithCancel(ctx)
	defer cancel()
	stream, err := c.NewStream(ctx, nil, method, opts...)
	if err != nil {
		return err
	}
	err = stream.SendMsg(args)
	if err != nil {
		if err == io.EOF { //nolint:errorlint
			// It's useless to return the EOF error as the calling code has no means to get the actual server error.
			// So, while this behavior is different from gRPC, we do the right thing here - get and return the actual error.
			// See https://github.com/grpc/grpc-go/issues/7230 for more info.
			return stream.RecvMsg(reply) // return the actual underlying error
		}
		return err
	}
	err = stream.CloseSend() // We send close send even for unary RPCs.
	if err != nil {
		return err
	}
	err = stream.RecvMsg(reply)
	if err != nil {
		return err
	}
	var x any
	err = stream.RecvMsg(x) // consume EOF
	if err != io.EOF {      //nolint:errorlint
		return err
	}
	return nil
}

func (c *RoutingClientConn[T]) NewStream(ctx context.Context, desc *grpc.StreamDesc, method string, opts ...grpc.CallOption) (grpc.ClientStream, error) {
	// 0. process call options
	ro, err := newRoutingOpts(opts)
	if err != nil {
		return nil, err
	}

	// 1. find a ready, suitable gateway tunserver
	rg, log, agentKey, err := c.Plugin.FindReadyGateway(ctx, c.Log, method)
	if err != nil {
		return nil, err
	}

	// 2. start streaming via the found gateway tunserver
	return &routingStream{
		log:               log.With(logz.GatewayURL(rg.URL.String())),
		api:               c.API,
		ro:                ro,
		agentKey:          agentKey,
		destination:       rg.Stream,
		destinationCancel: rg.StreamCancel,
		codec:             rg.Codec,
	}, nil
}

type routingOpts struct {
	header  []*metadata.MD
	trailer []*metadata.MD
}

func newRoutingOpts(opts []grpc.CallOption) (routingOpts, error) {
	r := routingOpts{}
	for _, opt := range opts {
		switch o := opt.(type) {
		case grpc.EmptyCallOption, grpc.StaticMethodCallOption:
		// ignore
		case grpc.HeaderCallOption:
			r.header = append(r.header, o.HeaderAddr)
		case grpc.TrailerCallOption:
			r.trailer = append(r.trailer, o.TrailerAddr)
		case grpc.FailFastCallOption:
			// Do nothing.
		default:
			// We could support PeerCallOption and OnFinishCallOption, but we don't need them, so we don't.
			// PeerCallOption, OnFinishCallOption, MaxRecvMsgSizeCallOption, MaxSendMsgSizeCallOption
			// PerRPCCredsCallOption, CompressorCallOption, ContentSubtypeCallOption, ForceCodecCallOption,
			// CustomCodecCallOption, MaxRetryRPCBufferSizeCallOption
			return routingOpts{}, status.Errorf(codes.Internal, "unsupported call option: %T", opt)
		}
	}
	return r, nil
}
