package modserver

import (
	"context"
	"log/slog"

	grpc_middleware "github.com/grpc-ecosystem/go-grpc-middleware/v2"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/api"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/module/modshared"
	"google.golang.org/grpc"
)

type agentRPCAPIKeyType int

const (
	agentRPCAPIKey agentRPCAPIKeyType = iota
)

// AgentRPCAPI provides the API for the module's gRPC handlers to use.
// It should be used only by modules, that handle requests from agents.
type AgentRPCAPI interface {
	modshared.RPCAPI
	// AgentTokenWithType returns the token and type of the agent making the RPC.
	AgentTokenWithType() api.AgentTokenWithType
	// AgentInfo returns information about the agent making the RPC.
	// Returns a gRPC-compatible error.
	// Returns an error with the Unavailable code if there was a retriable error.
	// If there was an error, it takes care of tracking it via HandleProcessingError().
	AgentInfo(ctx context.Context, log *slog.Logger) (api.AgentInfo, error)
}

type AgentRPCAPIFactory func(ctx context.Context, fullMethodName string) (AgentRPCAPI, error)

func InjectAgentRPCAPI(ctx context.Context, rpcAPI AgentRPCAPI) context.Context {
	return context.WithValue(ctx, agentRPCAPIKey, rpcAPI)
}

func AgentRPCAPIFromContext(ctx context.Context) AgentRPCAPI {
	rpcAPI, ok := ctx.Value(agentRPCAPIKey).(AgentRPCAPI)
	if !ok {
		// This is a programmer error, so panic.
		panic("modserver.AgentRPCAPI not attached to context. Make sure you are using interceptors")
	}
	return rpcAPI
}

// UnaryAgentRPCAPIInterceptor returns a new unary server interceptor that augments connection context with a AgentRPCAPI.
func UnaryAgentRPCAPIInterceptor(factory AgentRPCAPIFactory) grpc.UnaryServerInterceptor {
	return func(ctx context.Context, req any, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (any, error) {
		rpcAPI, err := factory(ctx, info.FullMethod)
		if err != nil {
			return nil, err
		}
		return handler(InjectAgentRPCAPI(ctx, rpcAPI), req)
	}
}

// StreamAgentRPCAPIInterceptor returns a new stream server interceptor that augments connection context with a AgentRPCAPI.
func StreamAgentRPCAPIInterceptor(factory AgentRPCAPIFactory) grpc.StreamServerInterceptor {
	return func(srv any, ss grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
		wrapper := grpc_middleware.WrapServerStream(ss)
		rpcAPI, err := factory(wrapper.WrappedContext, info.FullMethod)
		if err != nil {
			return err
		}
		wrapper.WrappedContext = InjectAgentRPCAPI(wrapper.WrappedContext, rpcAPI)
		return handler(srv, wrapper)
	}
}
