package server //nolint:dupl

import (
	"context"
	"strconv"

	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/module/agent_registrar/agentw_rpc"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/module/agent_tracker"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/module/event_tracker"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/module/modserver"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/server_api"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/logz"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/pkg/entity/agentw"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
	"google.golang.org/protobuf/types/known/timestamppb"
)

type RegisterAgentwEvent struct {
	AgentVersion       string            `json:"agent_version"`
	Architecture       string            `json:"architecture"`
	ConnectionID       int64             `json:"-"`
	AgentID            int64             `json:"agent_id"`
	ExtraTelemetryData map[string]string `json:"extra_telemetry_data"`
}

func (e *RegisterAgentwEvent) DeduplicateKey() string {
	return strconv.FormatInt(e.ConnectionID, 10)
}

type agentwServer struct {
	agentw_rpc.UnsafeAgentRegistrarServer
	agentRegisterer           agent_tracker.ExpiringRegisterer
	registerAgentEventTracker event_tracker.EventsInterface
}

func (s *agentwServer) Register(ctx context.Context, req *agentw_rpc.RegisterRequest) (*agentw_rpc.RegisterResponse, error) { //nolint:dupl
	rpcAPI := modserver.AgentRPCAPIFromContext(ctx)
	log := rpcAPI.Log()

	ctx, cancel := context.WithTimeout(context.WithoutCancel(ctx), registerTimeout) // we got the request, let's register even if the client cancels.
	defer cancel()

	// Get agent info
	a, err := rpcAPI.AgentInfo(ctx, log)
	if err != nil {
		return nil, err
	}

	err = s.registerAgentw(ctx, req, a, req.Meta)
	if err != nil {
		return nil, err
	}
	return &agentw_rpc.RegisterResponse{}, nil
}

func (s *agentwServer) registerAgentw(ctx context.Context, req *agentw_rpc.RegisterRequest, a server_api.AgentInfo, meta *agentw.Meta) error {
	agentInfo, ok := a.(*server_api.AgentwInfo)
	if !ok {
		return status.Error(codes.InvalidArgument, "expected agentw details to register agent")
	}

	connectedAgentInfo := &agent_tracker.ConnectedAgentwInfo{
		AgentMeta:    meta,
		ConnectedAt:  timestamppb.Now(),
		ConnectionId: req.InstanceId,
		WorkspaceId:  agentInfo.Key.ID,
	}

	// Register agent
	err := s.agentRegisterer.RegisterAgentwExpiring(ctx, connectedAgentInfo)
	if err != nil {
		rpcAPI := modserver.AgentRPCAPIFromContext(ctx)
		rpcAPI.HandleProcessingError(rpcAPI.Log(), "Failed to register agent", err, logz.AgentKey(agentInfo.Key))
		return status.Error(codes.Unavailable, "failed to register agent")
	}

	event := &RegisterAgentwEvent{
		AgentVersion:       meta.Version,
		Architecture:       meta.Architecture,
		ConnectionID:       req.InstanceId,
		AgentID:            agentInfo.Key.ID,
		ExtraTelemetryData: meta.ExtraTelemetryData,
	}
	s.registerAgentEventTracker.EmitEvent(event)
	return nil
}

func (s *agentwServer) Unregister(ctx context.Context, req *agentw_rpc.UnregisterRequest) (*agentw_rpc.UnregisterResponse, error) { //nolint:dupl
	rpcAPI := modserver.AgentRPCAPIFromContext(ctx)
	log := rpcAPI.Log()

	ctx, cancel := context.WithTimeout(context.WithoutCancel(ctx), registerTimeout) // we got the request, let's unregister even if the client cancels.
	defer cancel()

	// Get agent info
	a, err := rpcAPI.AgentInfo(ctx, log)
	if err != nil {
		return nil, err
	}

	err = s.unregisterAgentw(ctx, req, a, req.Meta)
	if err != nil {
		return nil, err
	}

	return &agentw_rpc.UnregisterResponse{}, nil
}

func (s *agentwServer) unregisterAgentw(ctx context.Context, req *agentw_rpc.UnregisterRequest, a server_api.AgentInfo, meta *agentw.Meta) error {
	agentInfo, ok := a.(*server_api.AgentwInfo)
	if !ok {
		return status.Error(codes.InvalidArgument, "expected agentw details to unregister agent")
	}

	disconnectAgentInfo := &agent_tracker.DisconnectAgentwInfo{
		AgentMeta:    meta,
		ConnectionId: req.InstanceId,
		WorkspaceId:  agentInfo.Key.ID,
	}

	// Unregister agent
	err := s.agentRegisterer.UnregisterAgentw(ctx, disconnectAgentInfo)
	if err != nil {
		rpcAPI := modserver.AgentRPCAPIFromContext(ctx)
		rpcAPI.HandleProcessingError(rpcAPI.Log(), "Failed to unregister agent", err, logz.AgentKey(agentInfo.Key))
		return status.Error(codes.Unavailable, "failed to unregister agent")
	}
	return nil
}
