package store

import (
	"context"
	"errors"
	"strings"
	"time"

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

const workerHeartbeatStaleAfter = 30 * time.Second

type WorkerRegistrationInput struct {
	InstanceID           string
	PodUID               string
	PodName              string
	Site                 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
	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")
	}
	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_worker_instances (
  instance_id, pod_uid, pod_name, site, 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, 'active', $6, $7, $8, $9, $10,
  $11, $12, $13, $14, $15, $16, $17,
  0, now(), now(), now()
)
ON CONFLICT (instance_id) DO UPDATE
SET pod_uid = EXCLUDED.pod_uid,
    pod_name = EXCLUDED.pod_name,
    site = EXCLUDED.site,
    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,
		strings.TrimSpace(input.PodUID),
		strings.TrimSpace(input.PodName),
		strings.TrimSpace(input.Site),
		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,
		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"`
	Site                 string     `json:"site,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 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) ListWorkerInstanceRuntime(ctx context.Context) ([]WorkerInstanceRuntime, error) {
	rows, err := s.pool.Query(ctx, `
SELECT worker.instance_id,
       worker.pod_uid,
       worker.pod_name,
       worker.site,
       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.site 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.Site,
			&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
}