geo/server/internal/tenant/app/desktop_task_service.go

package app

import (
	"context"
	"encoding/json"
	"errors"
	"strings"
	"time"

	"github.com/google/uuid"
	"github.com/jackc/pgx/v5"
	"go.uber.org/zap"

	"github.com/geo-platform/tenant-api/internal/shared/auth"
	"github.com/geo-platform/tenant-api/internal/shared/messaging/rabbitmq"
	"github.com/geo-platform/tenant-api/internal/shared/response"
	"github.com/geo-platform/tenant-api/internal/tenant/repository"
)

type DesktopTaskService struct {
	repo      repository.DesktopTaskRepository
	messaging *rabbitmq.Client
	logger    *zap.Logger
}

func NewDesktopTaskService(repo repository.DesktopTaskRepository, messaging *rabbitmq.Client, logger *zap.Logger) *DesktopTaskService {
	return &DesktopTaskService{
		repo:      repo,
		messaging: messaging,
		logger:    logger,
	}
}

type DesktopTaskView struct {
	ID              string          `json:"id"`
	JobID           string          `json:"job_id"`
	TenantID        int64           `json:"tenant_id"`
	WorkspaceID     int64           `json:"workspace_id"`
	TargetAccountID string          `json:"target_account_id"`
	TargetClientID  string          `json:"target_client_id"`
	Platform        string          `json:"platform"`
	Kind            string          `json:"kind"`
	Payload         json.RawMessage `json:"payload"`
	Status          string          `json:"status"`
	DedupKey        *string         `json:"dedup_key"`
	ActiveAttemptID *string         `json:"active_attempt_id"`
	LeaseExpiresAt  *time.Time      `json:"lease_expires_at"`
	Attempts        int             `json:"attempts"`
	ParkedReason    *string         `json:"parked_reason"`
	Result          json.RawMessage `json:"result"`
	Error           json.RawMessage `json:"error"`
	CreatedAt       time.Time       `json:"created_at"`
	UpdatedAt       time.Time       `json:"updated_at"`
}

type LeaseDesktopTaskRequest struct {
	TaskID *string `json:"task_id"`
	Kind   *string `json:"kind" binding:"omitempty,oneof=publish monitor"`
}

type LeaseDesktopTaskResponse struct {
	Task           *DesktopTaskView `json:"task"`
	AttemptID      *string          `json:"attempt_id,omitempty"`
	LeaseToken     *string          `json:"lease_token,omitempty"`
	LeaseExpiresAt *time.Time       `json:"lease_expires_at,omitempty"`
}

type ExtendDesktopTaskRequest struct {
	LeaseToken string `json:"lease_token" binding:"required"`
}

type ParkDesktopTaskRequest struct {
	LeaseToken string `json:"lease_token" binding:"required"`
	Reason     string `json:"reason" binding:"required,oneof=waiting_user captcha"`
}

type CompleteDesktopTaskRequest struct {
	LeaseToken string         `json:"lease_token" binding:"required"`
	Status     string         `json:"status" binding:"required,oneof=succeeded failed unknown"`
	Payload    map[string]any `json:"payload"`
	Error      map[string]any `json:"error"`
}

type CancelDesktopTaskRequest struct {
	LeaseToken *string `json:"lease_token"`
	Reason     *string `json:"reason"`
}

type ReconcileDesktopTaskRequest struct {
	Status string         `json:"status" binding:"required,oneof=succeeded failed aborted retry"`
	Result map[string]any `json:"result"`
	Error  map[string]any `json:"error"`
}

func (s *DesktopTaskService) Lease(ctx context.Context, client *repository.DesktopClient, req LeaseDesktopTaskRequest, routeTaskID *uuid.UUID, fromParked bool) (*LeaseDesktopTaskResponse, error) {
	if client == nil {
		return nil, response.ErrUnauthorized(40108, "desktop_client_missing", "desktop client context is required")
	}

	taskID := routeTaskID
	if taskID == nil && req.TaskID != nil && strings.TrimSpace(*req.TaskID) != "" {
		parsed, err := uuid.Parse(strings.TrimSpace(*req.TaskID))
		if err != nil {
			return nil, response.ErrBadRequest(40084, "invalid_desktop_task_id", "task id must be a uuid")
		}
		taskID = &parsed
	}

	rawToken, tokenHash, err := newDesktopClientToken()
	if err != nil {
		return nil, response.ErrInternal(50087, "desktop_task_lease_token_failed", "failed to generate desktop task lease token")
	}

	attemptID := uuid.New()
	leaseParams := repository.DesktopTaskLeaseParams{
		AttemptID:      attemptID,
		LeaseTokenHash: tokenHash,
	}

	var task *repository.DesktopTask
	switch {
	case fromParked:
		if taskID == nil {
			return nil, response.ErrBadRequest(40085, "missing_desktop_task_id", "task id is required when from_parked=true")
		}
		task, err = s.repo.LeaseParked(ctx, *taskID, client.ID, leaseParams)
	case taskID != nil:
		task, err = s.repo.LeaseQueuedByDesktopID(ctx, *taskID, client.ID, leaseParams)
	default:
		task, err = s.repo.LeaseNextQueued(ctx, client.ID, req.Kind, leaseParams)
	}
	if err != nil {
		if errors.Is(err, pgx.ErrNoRows) {
			if taskID == nil && !fromParked {
				return &LeaseDesktopTaskResponse{}, nil
			}
			return nil, s.classifyLeaseError(ctx, client, taskID, fromParked)
		}
		return nil, response.ErrInternal(50088, "desktop_task_lease_failed", "failed to lease desktop task")
	}

	if _, attemptErr := s.repo.CreateAttempt(ctx, repository.CreateDesktopTaskAttemptParams{
		DesktopID:      attemptID,
		TaskID:         task.DesktopID,
		ClientID:       client.ID,
		LeaseTokenHash: tokenHash,
	}); attemptErr != nil {
		s.logWarn("desktop task attempt insert failed", attemptErr, zap.String("task_id", task.DesktopID.String()))
	}

	s.publishTaskEvent(ctx, task, "task_leased")

	view := buildDesktopTaskView(task)
	attemptIDText := attemptID.String()
	return &LeaseDesktopTaskResponse{
		Task:           &view,
		AttemptID:      &attemptIDText,
		LeaseToken:     &rawToken,
		LeaseExpiresAt: task.LeaseExpiresAt,
	}, nil
}

func (s *DesktopTaskService) Extend(ctx context.Context, client *repository.DesktopClient, desktopID uuid.UUID, req ExtendDesktopTaskRequest) (*DesktopTaskView, error) {
	if client == nil {
		return nil, response.ErrUnauthorized(40108, "desktop_client_missing", "desktop client context is required")
	}

	task, err := s.repo.ExtendLease(ctx, desktopID, HashDesktopClientToken(strings.TrimSpace(req.LeaseToken)))
	if err != nil {
		if errors.Is(err, pgx.ErrNoRows) {
			return nil, response.ErrConflict(40983, "desktop_task_lease_invalid", "desktop task lease token is invalid or expired")
		}
		return nil, response.ErrInternal(50089, "desktop_task_extend_failed", "failed to extend desktop task lease")
	}

	s.publishTaskEvent(ctx, task, "task_extended")

	view := buildDesktopTaskView(task)
	return &view, nil
}

func (s *DesktopTaskService) Park(ctx context.Context, client *repository.DesktopClient, desktopID uuid.UUID, req ParkDesktopTaskRequest) (*DesktopTaskView, error) {
	if client == nil {
		return nil, response.ErrUnauthorized(40108, "desktop_client_missing", "desktop client context is required")
	}

	previousTask, _ := s.repo.GetByDesktopID(ctx, desktopID, client.WorkspaceID)
	previousAttemptID := activeAttemptID(previousTask)

	task, err := s.repo.Park(ctx, desktopID, HashDesktopClientToken(strings.TrimSpace(req.LeaseToken)), req.Reason)
	if err != nil {
		if errors.Is(err, pgx.ErrNoRows) {
			return nil, response.ErrConflict(40983, "desktop_task_lease_invalid", "desktop task lease token is invalid or expired")
		}
		return nil, response.ErrInternal(50090, "desktop_task_park_failed", "failed to park desktop task")
	}

	if previousAttemptID != nil {
		if finishErr := s.repo.FinishAttempt(ctx, repository.FinishDesktopTaskAttemptParams{
			TaskID:      task.DesktopID,
			AttemptID:   *previousAttemptID,
			FinalStatus: literalStringPtr("waiting_user"),
		}); finishErr != nil {
			s.logWarn("desktop task attempt finish failed after park", finishErr, zap.String("task_id", task.DesktopID.String()))
		}
	}

	s.publishTaskEvent(ctx, task, "task_parked")

	view := buildDesktopTaskView(task)
	return &view, nil
}

func (s *DesktopTaskService) Complete(ctx context.Context, client *repository.DesktopClient, desktopID uuid.UUID, req CompleteDesktopTaskRequest) (*DesktopTaskView, error) {
	if client == nil {
		return nil, response.ErrUnauthorized(40108, "desktop_client_missing", "desktop client context is required")
	}

	previousTask, _ := s.repo.GetByDesktopID(ctx, desktopID, client.WorkspaceID)
	previousAttemptID := activeAttemptID(previousTask)

	resultJSON, err := marshalOptionalJSON(req.Payload)
	if err != nil {
		return nil, response.ErrBadRequest(40086, "invalid_desktop_task_payload", "payload must be serializable")
	}
	errorJSON, err := marshalOptionalJSON(req.Error)
	if err != nil {
		return nil, response.ErrBadRequest(40087, "invalid_desktop_task_error", "error must be serializable")
	}

	task, err := s.repo.Complete(ctx, desktopID, HashDesktopClientToken(strings.TrimSpace(req.LeaseToken)), req.Status, resultJSON, errorJSON)
	if err != nil {
		if errors.Is(err, pgx.ErrNoRows) {
			return nil, response.ErrConflict(40983, "desktop_task_lease_invalid", "desktop task lease token is invalid or expired")
		}
		return nil, response.ErrInternal(50091, "desktop_task_complete_failed", "failed to complete desktop task")
	}

	if previousAttemptID != nil {
		if finishErr := s.repo.FinishAttempt(ctx, repository.FinishDesktopTaskAttemptParams{
			TaskID:      task.DesktopID,
			AttemptID:   *previousAttemptID,
			FinalStatus: &req.Status,
			Error:       errorJSON,
		}); finishErr != nil {
			s.logWarn("desktop task attempt finish failed after complete", finishErr, zap.String("task_id", task.DesktopID.String()))
		}
	}

	s.publishTaskEvent(ctx, task, "task_completed")

	view := buildDesktopTaskView(task)
	return &view, nil
}

func (s *DesktopTaskService) Cancel(ctx context.Context, client *repository.DesktopClient, desktopID uuid.UUID, req CancelDesktopTaskRequest) (*DesktopTaskView, error) {
	if client == nil {
		return nil, response.ErrUnauthorized(40108, "desktop_client_missing", "desktop client context is required")
	}

	previousTask, _ := s.repo.GetByDesktopID(ctx, desktopID, client.WorkspaceID)
	previousAttemptID := activeAttemptID(previousTask)

	errorJSON, err := marshalOptionalJSON(map[string]any{
		"reason": optionalStringValue(req.Reason),
		"source": "desktop_client",
	})
	if err != nil {
		return nil, response.ErrBadRequest(40088, "invalid_desktop_task_cancel", "cancel payload must be serializable")
	}

	var task *repository.DesktopTask
	if req.LeaseToken != nil && strings.TrimSpace(*req.LeaseToken) != "" {
		task, err = s.repo.CancelByLease(ctx, desktopID, HashDesktopClientToken(strings.TrimSpace(*req.LeaseToken)), errorJSON)
	} else {
		task, err = s.repo.CancelByClient(ctx, desktopID, client.ID, errorJSON)
	}
	if err != nil {
		if errors.Is(err, pgx.ErrNoRows) {
			return nil, response.ErrConflict(40984, "desktop_task_cancel_conflict", "desktop task cannot be canceled in its current state")
		}
		return nil, response.ErrInternal(50092, "desktop_task_cancel_failed", "failed to cancel desktop task")
	}

	if previousAttemptID != nil && req.LeaseToken != nil && strings.TrimSpace(*req.LeaseToken) != "" {
		if finishErr := s.repo.FinishAttempt(ctx, repository.FinishDesktopTaskAttemptParams{
			TaskID:      task.DesktopID,
			AttemptID:   *previousAttemptID,
			FinalStatus: literalStringPtr("aborted"),
			Error:       errorJSON,
		}); finishErr != nil {
			s.logWarn("desktop task attempt finish failed after cancel", finishErr, zap.String("task_id", task.DesktopID.String()))
		}
	}

	s.publishTaskEvent(ctx, task, "task_canceled")

	view := buildDesktopTaskView(task)
	return &view, nil
}

func (s *DesktopTaskService) TenantCancel(ctx context.Context, actor auth.Actor, desktopID uuid.UUID, req CancelDesktopTaskRequest) (*DesktopTaskView, error) {
	if actor.PrimaryWorkspaceID == 0 {
		return nil, response.ErrUnauthorized(40132, "workspace_scope_missing", "workspace context is required")
	}

	errorJSON, err := marshalOptionalJSON(map[string]any{
		"reason": optionalStringValue(req.Reason),
		"source": "tenant_web",
	})
	if err != nil {
		return nil, response.ErrBadRequest(40088, "invalid_desktop_task_cancel", "cancel payload must be serializable")
	}

	task, err := s.repo.TenantCancel(ctx, desktopID, actor.PrimaryWorkspaceID, errorJSON)
	if err != nil {
		if errors.Is(err, pgx.ErrNoRows) {
			return nil, response.ErrConflict(40984, "desktop_task_cancel_conflict", "desktop task cannot be canceled in its current state")
		}
		return nil, response.ErrInternal(50093, "desktop_task_tenant_cancel_failed", "failed to cancel desktop task")
	}

	s.publishTaskEvent(ctx, task, "task_canceled")

	view := buildDesktopTaskView(task)
	return &view, nil
}

func (s *DesktopTaskService) Reconcile(ctx context.Context, actor auth.Actor, desktopID uuid.UUID, req ReconcileDesktopTaskRequest) (*DesktopTaskView, error) {
	if actor.PrimaryWorkspaceID == 0 {
		return nil, response.ErrUnauthorized(40132, "workspace_scope_missing", "workspace context is required")
	}

	resultJSON, err := marshalOptionalJSON(req.Result)
	if err != nil {
		return nil, response.ErrBadRequest(40089, "invalid_desktop_task_reconcile_result", "result must be serializable")
	}
	errorJSON, err := marshalOptionalJSON(req.Error)
	if err != nil {
		return nil, response.ErrBadRequest(40090, "invalid_desktop_task_reconcile_error", "error must be serializable")
	}

	task, err := s.repo.Reconcile(ctx, desktopID, actor.PrimaryWorkspaceID, req.Status, resultJSON, errorJSON)
	if err != nil {
		if errors.Is(err, pgx.ErrNoRows) {
			return nil, response.ErrConflict(40985, "desktop_task_reconcile_conflict", "desktop task is not waiting for reconcile")
		}
		return nil, response.ErrInternal(50094, "desktop_task_reconcile_failed", "failed to reconcile desktop task")
	}

	s.publishTaskEvent(ctx, task, "task_reconciled")

	view := buildDesktopTaskView(task)
	return &view, nil
}

func (s *DesktopTaskService) classifyLeaseError(ctx context.Context, client *repository.DesktopClient, taskID *uuid.UUID, fromParked bool) error {
	if taskID == nil {
		return response.ErrConflict(40986, "desktop_task_unavailable", "no desktop task is currently available")
	}

	task, err := s.repo.GetByDesktopID(ctx, *taskID, client.WorkspaceID)
	if err != nil {
		if errors.Is(err, pgx.ErrNoRows) {
			return response.ErrNotFound(40482, "desktop_task_not_found", "desktop task not found")
		}
		return response.ErrInternal(50095, "desktop_task_lookup_failed", "failed to inspect desktop task state")
	}

	if task.TargetClientID != client.ID {
		return response.ErrForbidden(40381, "desktop_task_not_target_client", "desktop task is assigned to another desktop client")
	}

	if fromParked {
		switch task.Status {
		case "aborted":
			return response.ErrConflict(40987, "desktop_task_aborted", "desktop task has already been aborted")
		case "succeeded", "failed", "unknown":
			return response.ErrConflict(40988, "desktop_task_terminal", "desktop task is already in a terminal state")
		case "in_progress":
			return response.ErrConflict(40989, "desktop_task_already_in_progress", "desktop task is already in progress")
		default:
			return response.ErrConflict(40990, "desktop_task_not_parked", "desktop task is not waiting_user")
		}
	}

	if task.Status == "in_progress" {
		return response.ErrConflict(40989, "desktop_task_already_in_progress", "desktop task is already in progress")
	}

	return response.ErrConflict(40991, "desktop_task_not_queued", "desktop task is not queued")
}

func buildDesktopTaskView(task *repository.DesktopTask) DesktopTaskView {
	var activeAttemptID *string
	if task.ActiveAttemptID != nil {
		value := task.ActiveAttemptID.String()
		activeAttemptID = &value
	}

	return DesktopTaskView{
		ID:              task.DesktopID.String(),
		JobID:           task.JobID.String(),
		TenantID:        task.TenantID,
		WorkspaceID:     task.WorkspaceID,
		TargetAccountID: task.TargetAccountID.String(),
		TargetClientID:  task.TargetClientID.String(),
		Platform:        task.Platform,
		Kind:            task.Kind,
		Payload:         json.RawMessage(task.Payload),
		Status:          task.Status,
		DedupKey:        task.DedupKey,
		ActiveAttemptID: activeAttemptID,
		LeaseExpiresAt:  task.LeaseExpiresAt,
		Attempts:        task.Attempts,
		ParkedReason:    task.ParkedReason,
		Result:          json.RawMessage(task.Result),
		Error:           json.RawMessage(task.Error),
		CreatedAt:       task.CreatedAt,
		UpdatedAt:       task.UpdatedAt,
	}
}

func marshalOptionalJSON(value any) ([]byte, error) {
	if value == nil {
		return nil, nil
	}
	return json.Marshal(value)
}

func optionalStringValue(value *string) string {
	if value == nil {
		return ""
	}
	return strings.TrimSpace(*value)
}

func literalStringPtr(value string) *string {
	return &value
}

func activeAttemptID(task *repository.DesktopTask) *uuid.UUID {
	if task == nil || task.ActiveAttemptID == nil {
		return nil
	}
	value := *task.ActiveAttemptID
	return &value
}

func (s *DesktopTaskService) publishTaskEvent(ctx context.Context, task *repository.DesktopTask, eventType string) {
	if task == nil {
		return
	}

	err := publishDesktopTaskEvent(ctx, s.messaging, DesktopTaskEvent{
		Type:           eventType,
		TaskID:         task.DesktopID.String(),
		JobID:          task.JobID.String(),
		WorkspaceID:    task.WorkspaceID,
		TargetClientID: task.TargetClientID.String(),
		Status:         task.Status,
		Kind:           task.Kind,
		UpdatedAt:      task.UpdatedAt,
	})
	if err != nil {
		s.logWarn("desktop task event publish failed", err, zap.String("task_id", task.DesktopID.String()), zap.String("event_type", eventType))
	}
}

func (s *DesktopTaskService) logWarn(message string, err error, fields ...zap.Field) {
	if s.logger == nil || err == nil {
		return
	}
	fields = append(fields, zap.Error(err))
	s.logger.Warn(message, fields...)
}