package store

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

	"github.com/jackc/pgx/v5"
)

const workerHeartbeatStaleAfter = 15 * time.Second

type WorkerRegistrationInput struct {
	InstanceID              string
	WorkerID                string
	PoolID                  string
	Endpoint                string
	Labels                  map[string]string
	Capabilities            map[string]any
	ProtocolVersion         string
	OrchestratorInstanceRef string
	Revision                string
	DesiredCapacity         int
	CapacityLimit           int
	LoadMode                string
	SafeCapacity            int
	HeavyCapacity           int
	ActiveTasks             int
	PreparingTasks          int
	WaitingUpstreamTasks    int
	FinalizingTasks         int
	PressureState           string
	PressureReason          string
	LoadSampledAt           time.Time
	HeartbeatStaleAfter     time.Duration
}

type WorkerAllocation struct {
	InstanceID      string
	WorkerID        string
	PoolID          string
	DesiredCapacity int
	Allocated       int
	GlobalAllocated int
	ActiveInstances int
	HeartbeatAt     time.Time
}

type activeWorkerCapacity struct {
	InstanceID    string
	CapacityLimit int
}

// ActiveWorkerCapacity returns the cluster-wide execution capacity currently
// owned by live Worker instances.
func (s *Store) ActiveWorkerCapacity(ctx context.Context) (int, error) {
	var capacity int
	err := s.pool.QueryRow(ctx, `
SELECT COALESCE(SUM(allocated_capacity), 0)::int
FROM gateway_worker_instances
WHERE status = 'active'
  AND heartbeat_at > now() - $1::interval`, workerHeartbeatStaleAfter.String()).Scan(&capacity)
	return capacity, err
}

func (s *Store) RegisterWorkerInstance(ctx context.Context, input WorkerRegistrationInput) (WorkerAllocation, error) {
	input.InstanceID = strings.TrimSpace(input.InstanceID)
	if input.InstanceID == "" {
		return WorkerAllocation{}, errors.New("worker instance ID is required")
	}
	input.WorkerID = strings.TrimSpace(input.WorkerID)
	if input.WorkerID == "" {
		input.WorkerID = input.InstanceID
	}
	input.PoolID = strings.TrimSpace(input.PoolID)
	if input.PoolID == "" {
		input.PoolID = "legacy-default"
	}
	if strings.TrimSpace(input.ProtocolVersion) == "" {
		input.ProtocolVersion = "v1"
	}
	if input.Labels == nil {
		input.Labels = map[string]string{}
	}
	if input.Capabilities == nil {
		input.Capabilities = map[string]any{}
	}
	labels, err := json.Marshal(input.Labels)
	if err != nil {
		return WorkerAllocation{}, err
	}
	capabilities, err := json.Marshal(input.Capabilities)
	if err != nil {
		return WorkerAllocation{}, err
	}
	if input.DesiredCapacity < 0 {
		return WorkerAllocation{}, errors.New("worker desired capacity cannot be negative")
	}
	if input.CapacityLimit < 0 {
		return WorkerAllocation{}, errors.New("worker capacity limit cannot be negative")
	}
	if input.SafeCapacity < 0 || input.HeavyCapacity < 0 || input.ActiveTasks < 0 || input.PreparingTasks < 0 || input.WaitingUpstreamTasks < 0 || input.FinalizingTasks < 0 {
		return WorkerAllocation{}, errors.New("worker load values cannot be negative")
	}
	if input.ActiveTasks != input.PreparingTasks+input.WaitingUpstreamTasks+input.FinalizingTasks {
		return WorkerAllocation{}, errors.New("worker active task count must equal phase task counts")
	}
	if input.CapacityLimit == 0 {
		input.CapacityLimit = input.DesiredCapacity
	}
	hardCapacityLimit := input.CapacityLimit
	if strings.EqualFold(strings.TrimSpace(input.LoadMode), "adaptive") {
		input.CapacityLimit = min(input.CapacityLimit, input.SafeCapacity)
	}
	pressureState := strings.ToLower(strings.TrimSpace(input.PressureState))
	switch pressureState {
	case "normal", "busy", "critical":
	default:
		pressureState = "unknown"
	}
	if input.LoadSampledAt.IsZero() {
		input.LoadSampledAt = time.Now()
	}
	staleAfter := input.HeartbeatStaleAfter
	if staleAfter < workerHeartbeatStaleAfter {
		staleAfter = workerHeartbeatStaleAfter
	}
	tx, err := s.pool.Begin(ctx)
	if err != nil {
		return WorkerAllocation{}, err
	}
	defer rollbackTransaction(tx)
	// Worker heartbeats and allocations are ephemeral coordination state that is
	// refreshed every few seconds. Do not hold the global allocation lock while
	// waiting for a synchronous replica to acknowledge the commit.
	if _, err := tx.Exec(ctx, `SELECT set_config('synchronous_commit', 'off', true)`); err != nil {
		return WorkerAllocation{}, err
	}
	if _, err := tx.Exec(ctx, `SELECT pg_advisory_xact_lock(hashtextextended('gateway-worker-capacity', 0))`); err != nil {
		return WorkerAllocation{}, err
	}
	if _, err := tx.Exec(ctx, `
INSERT INTO gateway_execution_pools (pool_id, labels, capabilities, state, updated_at)
VALUES ($1, $2::jsonb, $3::jsonb, 'active', now())
ON CONFLICT (pool_id) DO UPDATE
SET labels = EXCLUDED.labels,
    capabilities = EXCLUDED.capabilities,
    updated_at = now()
WHERE gateway_execution_pools.state <> 'disabled'`, input.PoolID, labels, capabilities); err != nil {
		return WorkerAllocation{}, err
	}
	if _, err := tx.Exec(ctx, `
INSERT INTO gateway_worker_instances (
  instance_id, worker_id, pool_id, endpoint, labels, capabilities, protocol_version,
  orchestrator_instance_ref, revision, status,
  desired_capacity, capacity_limit, hard_capacity_limit, safe_capacity, heavy_capacity,
  active_tasks, preparing_tasks, waiting_upstream_tasks, finalizing_tasks,
  pressure_state, pressure_reason, load_sampled_at,
  allocated_capacity, started_at, heartbeat_at, updated_at
)
VALUES (
  $1, $2, $3, $4, $5::jsonb, $6::jsonb, $7,
  $8, $9, 'active', $10, $11, $12, $13, $14,
  $15, $16, $17, $18, $19, $20, $21,
  0, now(), now(), now()
)
ON CONFLICT (instance_id) DO UPDATE
SET worker_id = EXCLUDED.worker_id,
    pool_id = EXCLUDED.pool_id,
    endpoint = EXCLUDED.endpoint,
    labels = EXCLUDED.labels,
    capabilities = EXCLUDED.capabilities,
    protocol_version = EXCLUDED.protocol_version,
    orchestrator_instance_ref = EXCLUDED.orchestrator_instance_ref,
    revision = EXCLUDED.revision,
    status = CASE
      WHEN gateway_worker_instances.status = 'draining' THEN 'draining'
      ELSE 'active'
    END,
    desired_capacity = EXCLUDED.desired_capacity,
    capacity_limit = EXCLUDED.capacity_limit,
    hard_capacity_limit = EXCLUDED.hard_capacity_limit,
    safe_capacity = EXCLUDED.safe_capacity,
    heavy_capacity = EXCLUDED.heavy_capacity,
    active_tasks = EXCLUDED.active_tasks,
    preparing_tasks = EXCLUDED.preparing_tasks,
    waiting_upstream_tasks = EXCLUDED.waiting_upstream_tasks,
    finalizing_tasks = EXCLUDED.finalizing_tasks,
    pressure_state = EXCLUDED.pressure_state,
    pressure_reason = EXCLUDED.pressure_reason,
    load_sampled_at = EXCLUDED.load_sampled_at,
    heartbeat_at = now(),
    updated_at = now()`,
		input.InstanceID,
		input.WorkerID,
		input.PoolID,
		strings.TrimRight(strings.TrimSpace(input.Endpoint), "/"),
		labels,
		capabilities,
		strings.TrimSpace(input.ProtocolVersion),
		strings.TrimSpace(input.OrchestratorInstanceRef),
		strings.TrimSpace(input.Revision),
		input.DesiredCapacity,
		input.CapacityLimit,
		hardCapacityLimit,
		input.SafeCapacity,
		input.HeavyCapacity,
		input.ActiveTasks,
		input.PreparingTasks,
		input.WaitingUpstreamTasks,
		input.FinalizingTasks,
		pressureState,
		strings.TrimSpace(input.PressureReason),
		input.LoadSampledAt,
	); err != nil {
		return WorkerAllocation{}, err
	}
	if _, err := tx.Exec(ctx, `
UPDATE gateway_worker_instances
SET status = 'stale',
    allocated_capacity = 0,
    updated_at = now()
WHERE status = 'active'
  AND heartbeat_at <= now() - $1::interval`, staleAfter.String()); err != nil {
		return WorkerAllocation{}, err
	}

	rows, err := tx.Query(ctx, `
SELECT instance_id, capacity_limit
FROM gateway_worker_instances
WHERE status = 'active'
  AND heartbeat_at > now() - $1::interval
ORDER BY instance_id ASC
FOR UPDATE`, staleAfter.String())
	if err != nil {
		return WorkerAllocation{}, err
	}
	activeWorkers := make([]activeWorkerCapacity, 0)
	for rows.Next() {
		var worker activeWorkerCapacity
		if err := rows.Scan(&worker.InstanceID, &worker.CapacityLimit); err != nil {
			rows.Close()
			return WorkerAllocation{}, err
		}
		activeWorkers = append(activeWorkers, worker)
	}
	if err := rows.Err(); err != nil {
		rows.Close()
		return WorkerAllocation{}, err
	}
	rows.Close()
	if len(activeWorkers) == 0 {
		return WorkerAllocation{}, errors.New("worker registration was not active after heartbeat")
	}

	allocations, globalAllocated := allocateWorkerCapacities(activeWorkers, input.DesiredCapacity)
	allocated := 0
	for _, worker := range activeWorkers {
		capacity := allocations[worker.InstanceID]
		if _, err := tx.Exec(ctx, `
UPDATE gateway_worker_instances
SET desired_capacity = $2,
    allocated_capacity = $3,
    updated_at = now()
WHERE instance_id = $1
  AND (
    desired_capacity IS DISTINCT FROM $2
    OR allocated_capacity IS DISTINCT FROM $3
  )`, worker.InstanceID, input.DesiredCapacity, capacity); err != nil {
			return WorkerAllocation{}, err
		}
		if worker.InstanceID == input.InstanceID {
			allocated = capacity
		}
	}
	var heartbeatAt time.Time
	if err := tx.QueryRow(ctx, `
SELECT heartbeat_at
FROM gateway_worker_instances
WHERE instance_id = $1`, input.InstanceID).Scan(&heartbeatAt); err != nil {
		return WorkerAllocation{}, err
	}
	if err := tx.Commit(ctx); err != nil {
		return WorkerAllocation{}, err
	}
	return WorkerAllocation{
		InstanceID:      input.InstanceID,
		WorkerID:        input.WorkerID,
		PoolID:          input.PoolID,
		DesiredCapacity: input.DesiredCapacity,
		Allocated:       allocated,
		GlobalAllocated: globalAllocated,
		ActiveInstances: len(activeWorkers),
		HeartbeatAt:     heartbeatAt,
	}, nil
}

func allocateWorkerCapacities(workers []activeWorkerCapacity, desired int) (map[string]int, int) {
	allocations := make(map[string]int, len(workers))
	if desired <= 0 || len(workers) == 0 {
		return allocations, 0
	}
	remaining := desired
	for remaining > 0 {
		progressed := false
		for _, worker := range workers {
			limit := worker.CapacityLimit
			if limit <= 0 {
				continue
			}
			if allocations[worker.InstanceID] >= limit {
				continue
			}
			allocations[worker.InstanceID]++
			remaining--
			progressed = true
			if remaining == 0 {
				break
			}
		}
		if !progressed {
			break
		}
	}
	return allocations, desired - remaining
}

func (s *Store) MarkWorkerDraining(ctx context.Context, instanceID string) error {
	result, err := s.pool.Exec(ctx, `
UPDATE gateway_worker_instances
SET status = 'draining',
    allocated_capacity = 0,
    draining_at = COALESCE(draining_at, now()),
    heartbeat_at = now(),
    updated_at = now()
WHERE instance_id = $1`, strings.TrimSpace(instanceID))
	if err != nil {
		return err
	}
	if result.RowsAffected() == 0 {
		return pgx.ErrNoRows
	}
	return nil
}

func (s *Store) ReactivateWorkerInstance(ctx context.Context, instanceID string) error {
	result, err := s.pool.Exec(ctx, `
UPDATE gateway_worker_instances
SET status = 'active',
    draining_at = NULL,
    heartbeat_at = now(),
    updated_at = now()
WHERE instance_id = $1
  AND status = 'draining'`, strings.TrimSpace(instanceID))
	if err != nil {
		return err
	}
	if result.RowsAffected() == 0 {
		return pgx.ErrNoRows
	}
	return nil
}

type WorkerInstanceRuntime struct {
	InstanceID              string     `json:"instanceId"`
	PodUID                  string     `json:"podUid,omitempty"`
	PodName                 string     `json:"podName,omitempty"`
	OrchestratorInstanceRef string     `json:"-"`
	Site                    string     `json:"site,omitempty"`
	PoolID                  string     `json:"poolId,omitempty"`
	Revision                string     `json:"revision,omitempty"`
	Status                  string     `json:"status"`
	Allocated               int        `json:"allocatedCapacity"`
	CapacityLimit           int        `json:"capacityLimit"`
	HardCapacityLimit       int        `json:"hardCapacityLimit"`
	SafeCapacity            int        `json:"safeCapacity"`
	HeavyCapacity           int        `json:"heavyCapacity"`
	ReportedActiveTasks     int        `json:"reportedActiveTasks"`
	PreparingTasks          int        `json:"preparingTasks"`
	WaitingUpstreamTasks    int        `json:"waitingUpstreamTasks"`
	FinalizingTasks         int        `json:"finalizingTasks"`
	PressureState           string     `json:"pressureState"`
	PressureReason          string     `json:"pressureReason,omitempty"`
	LoadSampledAt           *time.Time `json:"loadSampledAt,omitempty"`
	RunningTasks            int        `json:"runningTasks"`
	ActiveLeases            int        `json:"activeLeases"`
	HeartbeatAt             time.Time  `json:"heartbeatAt"`
	DrainingAt              *time.Time `json:"drainingAt,omitempty"`
}

type WorkerQueueRuntime struct {
	Queued            int     `json:"queued"`
	Running           int     `json:"running"`
	OldestWaitSeconds float64 `json:"oldestWaitSeconds"`
}

type PoolQueueRuntime struct {
	PoolID  string `json:"poolId"`
	Queued  int    `json:"queued"`
	Running int    `json:"running"`
}

type WorkerClusterRuntime struct {
	Workers    []WorkerInstanceRuntime `json:"workers"`
	Queue      WorkerQueueRuntime      `json:"queue"`
	CapturedAt time.Time               `json:"capturedAt"`
}

func (s *Store) GetWorkerClusterRuntime(ctx context.Context) (WorkerClusterRuntime, error) {
	workers, err := s.ListWorkerInstanceRuntime(ctx)
	if err != nil {
		return WorkerClusterRuntime{}, err
	}
	queue, err := s.WorkerQueueRuntime(ctx)
	if err != nil {
		return WorkerClusterRuntime{}, err
	}
	return WorkerClusterRuntime{Workers: workers, Queue: queue, CapturedAt: time.Now()}, nil
}

type CapacityDatabaseHealth struct {
	Connections      int
	MaxConnections   int
	SynchronousPeers int
}

func (s *Store) WorkerQueueRuntime(ctx context.Context) (WorkerQueueRuntime, error) {
	var snapshot WorkerQueueRuntime
	err := s.pool.QueryRow(ctx, `
SELECT count(*) FILTER (WHERE status = 'queued')::int,
       count(*) FILTER (WHERE status = 'running')::int,
       COALESCE(EXTRACT(EPOCH FROM now() - MIN(created_at) FILTER (WHERE status = 'queued')), 0)::float8
FROM gateway_tasks
WHERE status IN ('queued', 'running')
  AND run_mode IN ('production', 'acceptance', 'acceptance_canary')`).Scan(
		&snapshot.Queued,
		&snapshot.Running,
		&snapshot.OldestWaitSeconds,
	)
	return snapshot, err
}

func (s *Store) ListPoolQueueRuntime(ctx context.Context) ([]PoolQueueRuntime, error) {
	rows, err := s.pool.Query(ctx, `
SELECT assigned_pool_id,
       count(*) FILTER (WHERE status = 'queued')::int,
       count(*) FILTER (WHERE status = 'running')::int
FROM gateway_tasks
WHERE assigned_pool_id IS NOT NULL
  AND status IN ('queued', 'running')
  AND run_mode IN ('production', 'acceptance', 'acceptance_canary')
GROUP BY assigned_pool_id
ORDER BY assigned_pool_id`)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	items := make([]PoolQueueRuntime, 0)
	for rows.Next() {
		var item PoolQueueRuntime
		if err := rows.Scan(&item.PoolID, &item.Queued, &item.Running); err != nil {
			return nil, err
		}
		items = append(items, item)
	}
	return items, rows.Err()
}

func (s *Store) ListWorkerInstanceRuntime(ctx context.Context) ([]WorkerInstanceRuntime, error) {
	rows, err := s.pool.Query(ctx, `
SELECT worker.instance_id,
       worker.pod_uid,
       worker.pod_name,
	   worker.orchestrator_instance_ref,
       worker.site,
	       worker.pool_id,
       worker.revision,
       worker.status,
       worker.allocated_capacity,
       worker.capacity_limit,
       worker.hard_capacity_limit,
       worker.safe_capacity,
       worker.heavy_capacity,
       worker.active_tasks,
       worker.preparing_tasks,
       worker.waiting_upstream_tasks,
       worker.finalizing_tasks,
       worker.pressure_state,
       worker.pressure_reason,
       worker.load_sampled_at,
       count(DISTINCT task.id) FILTER (WHERE task.status = 'running')::int,
       count(DISTINCT lease.id) FILTER (WHERE lease.released_at IS NULL)::int,
       worker.heartbeat_at,
       worker.draining_at
FROM gateway_worker_instances worker
LEFT JOIN gateway_tasks task
  ON task.locked_by = worker.instance_id
 AND task.status IN ('queued', 'running')
LEFT JOIN gateway_concurrency_leases lease
  ON lease.task_id = task.id
 AND lease.released_at IS NULL
WHERE worker.status IN ('active', 'draining')
  AND worker.heartbeat_at > now() - $1::interval
GROUP BY worker.instance_id
ORDER BY worker.pool_id ASC, worker.status DESC, worker.instance_id ASC`,
		runtimeWorkerStaleAfter.String(),
	)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	instances := make([]WorkerInstanceRuntime, 0)
	for rows.Next() {
		var instance WorkerInstanceRuntime
		if err := rows.Scan(
			&instance.InstanceID,
			&instance.PodUID,
			&instance.PodName,
			&instance.OrchestratorInstanceRef,
			&instance.Site,
			&instance.PoolID,
			&instance.Revision,
			&instance.Status,
			&instance.Allocated,
			&instance.CapacityLimit,
			&instance.HardCapacityLimit,
			&instance.SafeCapacity,
			&instance.HeavyCapacity,
			&instance.ReportedActiveTasks,
			&instance.PreparingTasks,
			&instance.WaitingUpstreamTasks,
			&instance.FinalizingTasks,
			&instance.PressureState,
			&instance.PressureReason,
			&instance.LoadSampledAt,
			&instance.RunningTasks,
			&instance.ActiveLeases,
			&instance.HeartbeatAt,
			&instance.DrainingAt,
		); err != nil {
			return nil, err
		}
		instances = append(instances, instance)
	}
	return instances, rows.Err()
}

func (s *Store) CapacityDatabaseHealth(ctx context.Context) (CapacityDatabaseHealth, error) {
	var health CapacityDatabaseHealth
	err := s.pool.QueryRow(ctx, `
SELECT
  (SELECT count(*)::int FROM pg_stat_activity),
  current_setting('max_connections')::int,
  (SELECT count(*)::int
   FROM pg_stat_replication
   WHERE state = 'streaming'
     AND sync_state IN ('sync', 'quorum'))`).Scan(
		&health.Connections,
		&health.MaxConnections,
		&health.SynchronousPeers,
	)
	return health, err
}

// YieldStaleAsyncTaskAdmissions removes waiting FIFO entries whose River job
// is still marked running even though the task no longer owns a live execution
// lease. The job and task remain untouched: a surviving worker can recreate
// the admission, while unrelated queue followers are no longer head-of-line
// blocked behind stale runtime ownership.
func (s *Store) YieldStaleAsyncTaskAdmissions(
	ctx context.Context,
	staleAfter time.Duration,
	limit int,
) (int64, error) {
	if staleAfter < workerHeartbeatStaleAfter {
		staleAfter = workerHeartbeatStaleAfter
	}
	if limit <= 0 || limit > 1000 {
		limit = 100
	}
	tx, err := s.pool.Begin(ctx)
	if err != nil {
		return 0, err
	}
	defer rollbackTransaction(tx)
	var yielded int64
	if err := tx.QueryRow(ctx, `
WITH stale AS MATERIALIZED (
  SELECT admission.task_id
  FROM gateway_task_admissions admission
  JOIN gateway_tasks task ON task.id = admission.task_id
  JOIN river_job job ON job.id = task.river_job_id
  WHERE admission.status = 'waiting'
    AND admission.mode = 'async'
    AND task.async_mode = true
    AND task.status = 'queued'
    AND (task.execution_lease_expires_at IS NULL OR task.execution_lease_expires_at <= now())
    AND job.queue = 'gateway_tasks'
    AND job.kind = 'gateway_task_run'
    AND job.state = 'running'
    AND job.attempted_at <= now() - $1::interval
  ORDER BY admission.priority ASC, admission.enqueued_at ASC, admission.task_id ASC
  LIMIT $2
  FOR UPDATE OF admission SKIP LOCKED
),
locked AS MATERIALIZED (
  SELECT task_id,
         pg_advisory_xact_lock(hashtextextended('task-admission:' || task_id::text, 0))
  FROM stale
),
deleted AS (
  DELETE FROM gateway_task_admissions admission
  USING locked
  WHERE admission.task_id = locked.task_id
    AND admission.status = 'waiting'
    AND admission.mode = 'async'
  RETURNING admission.task_id
)
SELECT count(*)::bigint
FROM deleted`, staleAfter.String(), limit).Scan(&yielded); err != nil {
		return 0, err
	}
	if err := tx.Commit(ctx); err != nil {
		return 0, err
	}
	if yielded > 0 {
		s.notifyTaskAdmissionBestEffort(ctx, "*")
	}
	return yielded, nil
}

// RecoverOrphanedAsyncRiverJobs retries stale River jobs after the gateway
// execution lease has already returned the task to queued. A process heartbeat
// does not prove that an individual River job goroutine is still alive; the
// task execution lease is the ownership fence.
func (s *Store) RecoverOrphanedAsyncRiverJobs(
	ctx context.Context,
	workerStaleAfter time.Duration,
	limit int,
) (int64, error) {
	if workerStaleAfter < workerHeartbeatStaleAfter {
		workerStaleAfter = workerHeartbeatStaleAfter
	}
	if limit <= 0 || limit > 1000 {
		limit = 100
	}
	tx, err := s.pool.Begin(ctx)
	if err != nil {
		return 0, err
	}
	defer rollbackTransaction(tx)
	var recovered int64
	if err := tx.QueryRow(ctx, `
WITH orphaned AS MATERIALIZED (
  SELECT job.id AS job_id,
         task.id AS task_id,
         job.attempt
  FROM river_job job
  JOIN gateway_tasks task ON task.river_job_id = job.id
  WHERE job.queue = 'gateway_tasks'
    AND job.kind = 'gateway_task_run'
    AND job.state = 'running'
    AND job.attempted_at <= now() - $1::interval
    AND task.async_mode = true
    AND task.status = 'queued'
    AND (task.execution_lease_expires_at IS NULL OR task.execution_lease_expires_at <= now())
  ORDER BY job.attempted_at ASC, job.id ASC
  LIMIT $2
  FOR UPDATE OF job SKIP LOCKED
),
released_leases AS (
  UPDATE gateway_concurrency_leases lease
  SET released_at = statement_timestamp()
  FROM orphaned
  WHERE lease.task_id = orphaned.task_id
    AND lease.released_at IS NULL
),
cleared_admissions AS (
  DELETE FROM gateway_task_admissions admission
  USING orphaned
  WHERE admission.task_id = orphaned.task_id
),
reset_tasks AS (
  UPDATE gateway_tasks task
  SET locked_by = NULL,
      locked_at = NULL,
      heartbeat_at = NULL,
      execution_token = NULL,
      execution_lease_expires_at = NULL,
      next_run_at = now(),
      updated_at = now()
  FROM orphaned
  WHERE task.id = orphaned.task_id
),
recovered_jobs AS (
  UPDATE river_job job
  SET errors = array_append(
        job.errors,
        jsonb_build_object(
          'at', now(),
          'attempt', orphaned.attempt,
          'error', 'Orphaned gateway job recovered after worker heartbeat expired',
          'trace', ''
        )
      ),
      finalized_at = NULL,
      scheduled_at = now(),
      state = 'retryable'
  FROM orphaned
  WHERE job.id = orphaned.job_id
  RETURNING job.id
)
SELECT count(*)::bigint
FROM recovered_jobs`, workerStaleAfter.String(), limit).Scan(&recovered); err != nil {
		return 0, err
	}
	if err := tx.Commit(ctx); err != nil {
		return 0, err
	}
	return recovered, nil
}

// FinalizeOrphanedTerminalAsyncRiverJobs closes River jobs whose gateway task
// is already terminal after the owning Worker disappeared. River deliberately
// uses a one-hour rescue horizon for long provider polls, so relying on its
// generic rescue loop would leave rollout-interrupted terminal jobs reported
// as running long after their execution lease and business resources were
// released.
func (s *Store) FinalizeOrphanedTerminalAsyncRiverJobs(
	ctx context.Context,
	workerStaleAfter time.Duration,
	limit int,
) (int64, error) {
	if workerStaleAfter < workerHeartbeatStaleAfter {
		workerStaleAfter = workerHeartbeatStaleAfter
	}
	if limit <= 0 || limit > 1000 {
		limit = 100
	}
	tx, err := s.pool.Begin(ctx)
	if err != nil {
		return 0, err
	}
	defer rollbackTransaction(tx)
	var finalized int64
	if err := tx.QueryRow(ctx, `
WITH orphaned AS MATERIALIZED (
  SELECT job.id AS job_id,
         task.id AS task_id,
         job.attempt
  FROM river_job job
  JOIN gateway_tasks task ON task.river_job_id = job.id
  WHERE job.queue = 'gateway_tasks'
    AND job.kind = 'gateway_task_run'
    AND job.state = 'running'
    AND job.attempted_at <= now() - $1::interval
    AND task.async_mode = true
    AND task.status IN ('succeeded', 'failed', 'cancelled')
    AND task.execution_token IS NULL
    AND (task.execution_lease_expires_at IS NULL OR task.execution_lease_expires_at <= now())
    AND NOT EXISTS (
      SELECT 1
      FROM gateway_worker_instances worker
      WHERE worker.status = 'active'
        AND worker.heartbeat_at > now() - $2::interval
        AND EXISTS (
          SELECT 1
          FROM unnest(job.attempted_by) attempted_owner
          WHERE attempted_owner LIKE worker.instance_id || '-exec-%'
        )
    )
  ORDER BY job.attempted_at ASC, job.id ASC
  LIMIT $3
  FOR UPDATE OF job SKIP LOCKED
), finalized_jobs AS (
  UPDATE river_job job
  SET errors = array_append(
        COALESCE(job.errors, ARRAY[]::jsonb[]),
        jsonb_build_object(
          'at', now(),
          'attempt', orphaned.attempt,
          'error', 'Terminal gateway job finalized after Worker owner disappeared',
          'trace', ''
        )
      ),
      finalized_at = now(),
      state = 'completed'
  FROM orphaned
  WHERE job.id = orphaned.job_id
  RETURNING job.id
)
SELECT count(*)::bigint
FROM finalized_jobs`, workerStaleAfter.String(), workerHeartbeatStaleAfter.String(), limit).Scan(&finalized); err != nil {
		return 0, err
	}
	if err := tx.Commit(ctx); err != nil {
		return 0, err
	}
	return finalized, nil
}
