package server

import (
	"context"
	"fmt"
	"log/slog"
	"strconv"

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

// MUST be lowercase because the code sets map entries directly, not via setters (to avoid overhead).
const (
	agentIDMDKey   = "kas2agentk-tunnel-agent-id"
	agentTypeMDKey = "kas2agentk-tunnel-agent-type"
)

var (
	_ modserver.AgentRPCAPI = (*agentRPCAPI)(nil)
)

type agentRPCAPI struct {
	modserver.AgentRPCAPI
	expectedAgentKey api.AgentKey
}

func (a *agentRPCAPI) AgentInfo(ctx context.Context, log *slog.Logger) (api.AgentInfo, error) {
	info, err := a.AgentRPCAPI.AgentInfo(ctx, log)
	if err != nil {
		return nil, err
	}
	if info.GetAgentKey() != a.expectedAgentKey {
		return nil, status.Error(codes.Unauthenticated, "provided agent token belongs to a different agent key")
	}
	return info, nil
}

func AgentRPCAPIFactory(delegate modserver.AgentRPCAPIFactory) modserver.AgentRPCAPIFactory {
	return func(ctx context.Context, fullMethodName string) (modserver.AgentRPCAPI, error) {
		rpcAPI, err := delegate(ctx, fullMethodName)
		if err != nil {
			return nil, err
		}
		expectedAgentIDStr := metadata.ValueFromIncomingContext(ctx, agentIDMDKey)
		if len(expectedAgentIDStr) != 1 {
			err = fmt.Errorf("expecting a single %s, got %d", agentIDMDKey, len(expectedAgentIDStr))
			rpcAPI.HandleProcessingError(rpcAPI.Log(), "AgentKey validation", err)
			// This is a coding mistake and shouldn't ever happen, so use the Internal error code.
			return nil, status.Error(codes.Internal, err.Error())
		}
		expectedAgentID, err := strconv.ParseInt(expectedAgentIDStr[0], 10, 64)
		if err != nil {
			return nil, status.Errorf(codes.InvalidArgument, "invalid %s", agentIDMDKey)
		}

		expectedAgentTypeStr := metadata.ValueFromIncomingContext(ctx, agentTypeMDKey)
		if len(expectedAgentTypeStr) > 1 {
			err = fmt.Errorf("expecting a single %s, got %d", agentTypeMDKey, len(expectedAgentTypeStr))
			rpcAPI.HandleProcessingError(rpcAPI.Log(), "AgentKey validation", err)
			// This is a coding mistake and shouldn't ever happen, so use the Internal error code.
			return nil, status.Error(codes.Internal, err.Error())
		}

		expectedAgentType := api.AgentTypeKubernetes
		if len(expectedAgentTypeStr) == 1 {
			expectedAgentType, err = api.ParseAgentType(expectedAgentTypeStr[0])
			if err != nil {
				return nil, status.Error(codes.InvalidArgument, err.Error())
			}
		}
		return &agentRPCAPI{
			AgentRPCAPI:      rpcAPI,
			expectedAgentKey: api.AgentKey{ID: expectedAgentID, Type: expectedAgentType},
		}, nil
	}
}
