package server //nolint:dupl

import (
	"context"
	"strconv"

	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/module/agent_registrar/runnerc_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/runnerc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
	"google.golang.org/protobuf/types/known/timestamppb"
)

type RegisterRunnerControllerEvent 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 *RegisterRunnerControllerEvent) DeduplicateKey() string {
	return strconv.FormatInt(e.ConnectionID, 10)
}

type runnerControllerServer struct {
	runnerc_rpc.UnsafeAgentRegistrarServer
	agentRegisterer           agent_tracker.ExpiringRegisterer
	registerAgentEventTracker event_tracker.EventsInterface
}

func (s *runnerControllerServer) Register(ctx context.Context, req *runnerc_rpc.RegisterRequest) (*runnerc_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.registerRunnerController(ctx, req, a, req.Meta)
	if err != nil {
		return nil, err
	}

	return &runnerc_rpc.RegisterResponse{}, nil
}

func (s *runnerControllerServer) registerRunnerController(ctx context.Context, req *runnerc_rpc.RegisterRequest, a server_api.AgentInfo, meta *runnerc.Meta) error {
	agentInfo, ok := a.(*server_api.RunnerControllerInfo)
	if !ok {
		return status.Error(codes.InvalidArgument, "expected runnerc details to register agent")
	}

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

	// Register agent
	err := s.agentRegisterer.RegisterRunnerControllerExpiring(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 := &RegisterRunnerControllerEvent{
		AgentVersion:       meta.Version,
		Architecture:       meta.Architecture,
		ConnectionID:       req.InstanceId,
		AgentID:            agentInfo.Key.ID,
		ExtraTelemetryData: meta.ExtraTelemetryData,
	}
	s.registerAgentEventTracker.EmitEvent(event)
	return nil
}

func (s *runnerControllerServer) Unregister(ctx context.Context, req *runnerc_rpc.UnregisterRequest) (*runnerc_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.unregisterRunnerController(ctx, req, a, req.Meta)
	if err != nil {
		return nil, err
	}

	return &runnerc_rpc.UnregisterResponse{}, nil
}

func (s *runnerControllerServer) unregisterRunnerController(ctx context.Context, req *runnerc_rpc.UnregisterRequest, a server_api.AgentInfo, meta *runnerc.Meta) error {
	agentInfo, ok := a.(*server_api.RunnerControllerInfo)
	if !ok {
		return status.Error(codes.InvalidArgument, "expected runnerc details to unregister agent")
	}

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

	// Unregister agent
	err := s.agentRegisterer.UnregisterRunnerController(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
}
