geo/server/internal/ops/repository/scheduler.go

package repository

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

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

	"github.com/geo-platform/tenant-api/internal/ops/domain"
)

var ErrSchedulerJobNotFound = errors.New("scheduler job not found")

type SchedulerRepository struct {
	pool *pgxpool.Pool
}

func NewSchedulerRepository(pool *pgxpool.Pool) *SchedulerRepository {
	return &SchedulerRepository{pool: pool}
}

type SchedulerJobFilter struct {
	Category string
	Keyword  string
	Enabled  *bool
	Limit    int
	Offset   int
}

type SchedulerJobUpdate struct {
	Enabled         *bool
	ScheduleType    *string
	IntervalSeconds *int
	CronExpr        *string
	Timezone        *string
	TimeoutSeconds  *int
	BatchSize       *int
	MaxConcurrency  *int
	Config          map[string]any
	UpdatedBy       *int64
}

type SchedulerRunFilter struct {
	JobKey string
	Status string
	Limit  int
	Offset int
}

type SchedulerRunStart struct {
	JobKey              string
	TriggerID           *int64
	SchedulerInstanceID string
	TriggerType         string
	ConfigVersion       int
	ConfigSnapshot      map[string]any
}

type SchedulerRunFinish struct {
	RunID        int64
	Status       string
	Stats        map[string]any
	ErrorMessage *string
}

type SchedulerTriggerCreate struct {
	JobKey         string
	TriggerType    string
	RequestedBy    *int64
	Reason         string
	ConfigOverride map[string]any
}

type SchedulerTriggerClaim struct {
	ID             int64
	JobKey         string
	TriggerType    string
	ConfigOverride map[string]any
}

const schedulerJobSelect = `
	SELECT
		j.job_key,
		j.display_name,
		j.category,
		j.description,
		j.enabled,
		j.schedule_type,
		j.interval_seconds,
		j.cron_expr,
		j.timezone,
		j.timeout_seconds,
		j.batch_size,
		j.max_concurrency,
		j.config,
		j.version,
		j.updated_by,
		j.updated_at,
		j.created_at,
		COALESCE(pending.pending_count, 0)
	FROM ops.scheduler_jobs j
	LEFT JOIN LATERAL (
		SELECT COUNT(*)::int AS pending_count
		FROM ops.scheduler_job_triggers t
		WHERE t.job_key = j.job_key AND t.status = 'pending'
	) pending ON true
`

func (r *SchedulerRepository) ListJobs(ctx context.Context, f SchedulerJobFilter) ([]domain.SchedulerJob, int64, error) {
	args := []any{}
	where := "1=1"
	if f.Category != "" {
		args = append(args, f.Category)
		where += fmt.Sprintf(" AND j.category = $%d", len(args))
	}
	if f.Keyword != "" {
		args = append(args, "%"+f.Keyword+"%")
		where += fmt.Sprintf(" AND (j.job_key ILIKE $%d OR j.display_name ILIKE $%d OR COALESCE(j.description, '') ILIKE $%d)", len(args), len(args), len(args))
	}
	if f.Enabled != nil {
		args = append(args, *f.Enabled)
		where += fmt.Sprintf(" AND j.enabled = $%d", len(args))
	}

	var total int64
	if err := r.pool.QueryRow(ctx, "SELECT COUNT(*) FROM ops.scheduler_jobs j WHERE "+where, args...).Scan(&total); err != nil {
		return nil, 0, err
	}

	limit := f.Limit
	if limit <= 0 || limit > 200 {
		limit = 50
	}
	offset := f.Offset
	if offset < 0 {
		offset = 0
	}
	args = append(args, limit, offset)
	limitArg := len(args) - 1
	offsetArg := len(args)

	rows, err := r.pool.Query(ctx, schedulerJobSelect+fmt.Sprintf(`
		WHERE %s
		ORDER BY j.category ASC, j.job_key ASC
		LIMIT $%d OFFSET $%d`, where, limitArg, offsetArg), args...)
	if err != nil {
		return nil, 0, err
	}
	defer rows.Close()

	items := make([]domain.SchedulerJob, 0, limit)
	for rows.Next() {
		item, scanErr := scanSchedulerJob(rows)
		if scanErr != nil {
			return nil, 0, scanErr
		}
		items = append(items, *item)
	}
	if err := rows.Err(); err != nil {
		return nil, 0, err
	}

	if err := r.attachRunSnapshots(ctx, items); err != nil {
		return nil, 0, err
	}
	return items, total, nil
}

func (r *SchedulerRepository) GetJob(ctx context.Context, key string) (*domain.SchedulerJob, error) {
	item, err := scanSchedulerJob(r.pool.QueryRow(ctx, schedulerJobSelect+`
		WHERE j.job_key = $1`, key))
	if err != nil {
		return nil, err
	}
	items := []domain.SchedulerJob{*item}
	if err := r.attachRunSnapshots(ctx, items); err != nil {
		return nil, err
	}
	*item = items[0]
	return item, nil
}

func (r *SchedulerRepository) GetJobForRuntime(ctx context.Context, key string) (*domain.SchedulerJob, error) {
	return scanSchedulerJob(r.pool.QueryRow(ctx, schedulerJobSelect+`
		WHERE j.job_key = $1`, key))
}

func (r *SchedulerRepository) LastAutoRunStartedAt(ctx context.Context, key string) (time.Time, bool, error) {
	var startedAt time.Time
	err := r.pool.QueryRow(ctx, `
		SELECT started_at
		FROM ops.scheduler_job_runs
		WHERE job_key = $1
		  AND trigger_type = 'auto'
		ORDER BY started_at DESC, id DESC
		LIMIT 1
	`, key).Scan(&startedAt)
	if err != nil {
		if errors.Is(err, pgx.ErrNoRows) {
			return time.Time{}, false, nil
		}
		return time.Time{}, false, err
	}
	return startedAt, true, nil
}

func (r *SchedulerRepository) UpdateJob(ctx context.Context, key string, in SchedulerJobUpdate) (*domain.SchedulerJob, error) {
	configJSON, err := marshalSchedulerJSONMap(in.Config)
	if err != nil {
		return nil, err
	}
	return scanSchedulerJob(r.pool.QueryRow(ctx, `
		UPDATE ops.scheduler_jobs
		SET enabled = COALESCE($2, enabled),
		    schedule_type = COALESCE(NULLIF($3, ''), schedule_type),
		    interval_seconds = COALESCE($4, interval_seconds),
		    cron_expr = COALESCE($5, cron_expr),
		    timezone = COALESCE(NULLIF($6, ''), timezone),
		    timeout_seconds = COALESCE($7, timeout_seconds),
		    batch_size = COALESCE($8, batch_size),
		    max_concurrency = COALESCE($9, max_concurrency),
		    config = COALESCE($10::jsonb, config),
		    updated_by = COALESCE($11, updated_by),
		    version = version + 1,
		    updated_at = NOW()
		WHERE job_key = $1
		RETURNING job_key, display_name, category, description, enabled, schedule_type,
			interval_seconds, cron_expr, timezone, timeout_seconds, batch_size,
			max_concurrency, config, version, updated_by, updated_at, created_at, 0::int
	`, key, in.Enabled, nullableStringPtr(in.ScheduleType), in.IntervalSeconds, in.CronExpr,
		nullableStringPtr(in.Timezone), in.TimeoutSeconds, in.BatchSize, in.MaxConcurrency, configJSON, in.UpdatedBy))
}

func (r *SchedulerRepository) SetJobEnabled(ctx context.Context, key string, enabled bool, updatedBy *int64) (*domain.SchedulerJob, error) {
	return r.UpdateJob(ctx, key, SchedulerJobUpdate{
		Enabled:   &enabled,
		UpdatedBy: updatedBy,
	})
}

func (r *SchedulerRepository) CreateTrigger(ctx context.Context, in SchedulerTriggerCreate) (*domain.SchedulerTrigger, error) {
	configJSON, err := marshalSchedulerJSONMap(in.ConfigOverride)
	if err != nil {
		return nil, err
	}
	reason := nullableString(in.Reason)
	return scanSchedulerTrigger(r.pool.QueryRow(ctx, `
		INSERT INTO ops.scheduler_job_triggers (job_key, trigger_type, requested_by, reason, config_override)
		VALUES ($1, $2, $3, $4, COALESCE($5::jsonb, '{}'::jsonb))
		RETURNING id, job_key, trigger_type, status, requested_by, reason, config_override,
			scheduler_instance_id, run_id, requested_at, claimed_at, finished_at, error_message
	`, in.JobKey, in.TriggerType, in.RequestedBy, reason, configJSON))
}

func (r *SchedulerRepository) ClaimPendingTrigger(ctx context.Context, jobKey, instanceID string) (*SchedulerTriggerClaim, bool, error) {
	var claim SchedulerTriggerClaim
	var configRaw []byte
	err := r.pool.QueryRow(ctx, `
		UPDATE ops.scheduler_job_triggers t
		SET status = 'claimed',
		    claimed_at = NOW(),
		    scheduler_instance_id = $2
		WHERE t.id = (
			SELECT id
			FROM ops.scheduler_job_triggers
			WHERE job_key = $1 AND status = 'pending'
			ORDER BY requested_at ASC, id ASC
			FOR UPDATE SKIP LOCKED
			LIMIT 1
		)
		RETURNING id, job_key, trigger_type, config_override
	`, jobKey, instanceID).Scan(&claim.ID, &claim.JobKey, &claim.TriggerType, &configRaw)
	if err != nil {
		if errors.Is(err, pgx.ErrNoRows) {
			return nil, false, nil
		}
		return nil, false, err
	}
	claim.ConfigOverride = decodeSchedulerJSONMap(configRaw)
	return &claim, true, nil
}

func (r *SchedulerRepository) FinishTrigger(ctx context.Context, triggerID int64, runID *int64, status string, errorMessage *string) error {
	_, err := r.pool.Exec(ctx, `
		UPDATE ops.scheduler_job_triggers
		SET status = $3,
		    run_id = $2,
		    finished_at = NOW(),
		    error_message = $4
		WHERE id = $1
	`, triggerID, runID, status, errorMessage)
	return err
}

func (r *SchedulerRepository) StartRun(ctx context.Context, in SchedulerRunStart) (*domain.SchedulerRun, error) {
	configJSON, err := marshalSchedulerJSONMap(in.ConfigSnapshot)
	if err != nil {
		return nil, err
	}
	return scanSchedulerRun(r.pool.QueryRow(ctx, `
		INSERT INTO ops.scheduler_job_runs (
			job_key, trigger_id, scheduler_instance_id, trigger_type,
			status, config_version, config_snapshot, started_at
		)
		VALUES ($1, $2, $3, $4, 'running', $5, COALESCE($6::jsonb, '{}'::jsonb), NOW())
		RETURNING id, job_key, trigger_id, scheduler_instance_id, trigger_type, status,
			config_version, config_snapshot, stats, error_message, started_at,
			finished_at, duration_ms
	`, in.JobKey, in.TriggerID, nullableString(in.SchedulerInstanceID), in.TriggerType, in.ConfigVersion, configJSON))
}

func (r *SchedulerRepository) FinishRun(ctx context.Context, in SchedulerRunFinish) (*domain.SchedulerRun, error) {
	statsJSON, err := marshalSchedulerJSONMap(in.Stats)
	if err != nil {
		return nil, err
	}
	return scanSchedulerRun(r.pool.QueryRow(ctx, `
		UPDATE ops.scheduler_job_runs
		SET status = $2,
		    stats = COALESCE($3::jsonb, '{}'::jsonb),
		    error_message = $4,
		    finished_at = NOW(),
		    duration_ms = GREATEST(0, FLOOR(EXTRACT(EPOCH FROM (NOW() - started_at)) * 1000)::bigint)
		WHERE id = $1
		RETURNING id, job_key, trigger_id, scheduler_instance_id, trigger_type, status,
			config_version, config_snapshot, stats, error_message, started_at,
			finished_at, duration_ms
	`, in.RunID, in.Status, statsJSON, in.ErrorMessage))
}

func (r *SchedulerRepository) ListRuns(ctx context.Context, f SchedulerRunFilter) ([]domain.SchedulerRun, int64, error) {
	args := []any{}
	where := "1=1"
	if f.JobKey != "" {
		args = append(args, f.JobKey)
		where += fmt.Sprintf(" AND job_key = $%d", len(args))
	}
	if f.Status != "" {
		args = append(args, f.Status)
		where += fmt.Sprintf(" AND status = $%d", len(args))
	}
	var total int64
	if err := r.pool.QueryRow(ctx, "SELECT COUNT(*) FROM ops.scheduler_job_runs WHERE "+where, args...).Scan(&total); err != nil {
		return nil, 0, err
	}
	limit := f.Limit
	if limit <= 0 || limit > 200 {
		limit = 50
	}
	offset := f.Offset
	if offset < 0 {
		offset = 0
	}
	args = append(args, limit, offset)
	limitArg := len(args) - 1
	offsetArg := len(args)
	rows, err := r.pool.Query(ctx, fmt.Sprintf(`
		SELECT id, job_key, trigger_id, scheduler_instance_id, trigger_type, status,
			config_version, config_snapshot, stats, error_message, started_at,
			finished_at, duration_ms
		FROM ops.scheduler_job_runs
		WHERE %s
		ORDER BY started_at DESC, id DESC
		LIMIT $%d OFFSET $%d
	`, where, limitArg, offsetArg), args...)
	if err != nil {
		return nil, 0, err
	}
	defer rows.Close()
	items := make([]domain.SchedulerRun, 0, limit)
	for rows.Next() {
		item, scanErr := scanSchedulerRun(rows)
		if scanErr != nil {
			return nil, 0, scanErr
		}
		items = append(items, *item)
	}
	return items, total, rows.Err()
}

func (r *SchedulerRepository) UpsertInstance(ctx context.Context, instance domain.SchedulerInstance) error {
	metadataJSON, err := marshalSchedulerJSONMap(instance.Metadata)
	if err != nil {
		return err
	}
	_, err = r.pool.Exec(ctx, `
		INSERT INTO ops.scheduler_instances (instance_id, hostname, process, build_version, metadata, started_at, last_seen_at)
		VALUES ($1, $2, $3, $4, COALESCE($5::jsonb, '{}'::jsonb), $6, NOW())
		ON CONFLICT (instance_id) DO UPDATE
		SET hostname = EXCLUDED.hostname,
		    process = EXCLUDED.process,
		    build_version = EXCLUDED.build_version,
		    metadata = EXCLUDED.metadata,
		    last_seen_at = NOW()
	`, instance.InstanceID, instance.Hostname, instance.Process, instance.BuildVersion, metadataJSON, instance.StartedAt)
	return err
}

func (r *SchedulerRepository) ListInstances(ctx context.Context) ([]domain.SchedulerInstance, error) {
	rows, err := r.pool.Query(ctx, `
		SELECT instance_id, hostname, process, build_version, started_at, last_seen_at, metadata
		FROM ops.scheduler_instances
		ORDER BY last_seen_at DESC
		LIMIT 200
	`)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	items := make([]domain.SchedulerInstance, 0, 16)
	for rows.Next() {
		item, scanErr := scanSchedulerInstance(rows)
		if scanErr != nil {
			return nil, scanErr
		}
		items = append(items, *item)
	}
	return items, rows.Err()
}

func (r *SchedulerRepository) attachRunSnapshots(ctx context.Context, items []domain.SchedulerJob) error {
	if len(items) == 0 {
		return nil
	}
	keys := make([]string, 0, len(items))
	for _, item := range items {
		keys = append(keys, item.JobKey)
	}
	rows, err := r.pool.Query(ctx, `
		WITH ranked AS (
			SELECT
				id, job_key, trigger_id, scheduler_instance_id, trigger_type, status,
				config_version, config_snapshot, stats, error_message, started_at,
				finished_at, duration_ms,
				ROW_NUMBER() OVER (PARTITION BY job_key ORDER BY started_at DESC, id DESC) AS rn
			FROM ops.scheduler_job_runs
			WHERE job_key = ANY($1)
		)
		SELECT id, job_key, trigger_id, scheduler_instance_id, trigger_type, status,
			config_version, config_snapshot, stats, error_message, started_at,
			finished_at, duration_ms
		FROM ranked
		WHERE rn = 1
	`, keys)
	if err != nil {
		return err
	}
	defer rows.Close()
	lastRuns := map[string]*domain.SchedulerRun{}
	for rows.Next() {
		run, scanErr := scanSchedulerRun(rows)
		if scanErr != nil {
			return scanErr
		}
		lastRuns[run.JobKey] = run
	}
	if err := rows.Err(); err != nil {
		return err
	}

	runningRows, err := r.pool.Query(ctx, `
		SELECT DISTINCT ON (job_key)
			id, job_key, trigger_id, scheduler_instance_id, trigger_type, status,
			config_version, config_snapshot, stats, error_message, started_at,
			finished_at, duration_ms
		FROM ops.scheduler_job_runs
		WHERE job_key = ANY($1) AND status = 'running'
		ORDER BY job_key, started_at DESC, id DESC
	`, keys)
	if err != nil {
		return err
	}
	defer runningRows.Close()
	runningRuns := map[string]*domain.SchedulerRun{}
	for runningRows.Next() {
		run, scanErr := scanSchedulerRun(runningRows)
		if scanErr != nil {
			return scanErr
		}
		runningRuns[run.JobKey] = run
	}
	if err := runningRows.Err(); err != nil {
		return err
	}

	for i := range items {
		items[i].LastRun = lastRuns[items[i].JobKey]
		items[i].RunningRun = runningRuns[items[i].JobKey]
	}
	return nil
}

func scanSchedulerJob(row pgx.Row) (*domain.SchedulerJob, error) {
	var item domain.SchedulerJob
	var configRaw []byte
	err := row.Scan(
		&item.JobKey,
		&item.DisplayName,
		&item.Category,
		&item.Description,
		&item.Enabled,
		&item.ScheduleType,
		&item.IntervalSeconds,
		&item.CronExpr,
		&item.Timezone,
		&item.TimeoutSeconds,
		&item.BatchSize,
		&item.MaxConcurrency,
		&configRaw,
		&item.Version,
		&item.UpdatedBy,
		&item.UpdatedAt,
		&item.CreatedAt,
		&item.PendingTriggers,
	)
	if err != nil {
		if errors.Is(err, pgx.ErrNoRows) {
			return nil, ErrSchedulerJobNotFound
		}
		return nil, err
	}
	item.Config = decodeSchedulerJSONMap(configRaw)
	return &item, nil
}

func scanSchedulerRun(row pgx.Row) (*domain.SchedulerRun, error) {
	var item domain.SchedulerRun
	var configRaw, statsRaw []byte
	err := row.Scan(
		&item.ID,
		&item.JobKey,
		&item.TriggerID,
		&item.SchedulerInstanceID,
		&item.TriggerType,
		&item.Status,
		&item.ConfigVersion,
		&configRaw,
		&statsRaw,
		&item.ErrorMessage,
		&item.StartedAt,
		&item.FinishedAt,
		&item.DurationMilliseconds,
	)
	if err != nil {
		return nil, err
	}
	item.ConfigSnapshot = decodeSchedulerJSONMap(configRaw)
	item.Stats = decodeSchedulerJSONMap(statsRaw)
	return &item, nil
}

func scanSchedulerTrigger(row pgx.Row) (*domain.SchedulerTrigger, error) {
	var item domain.SchedulerTrigger
	var configRaw []byte
	err := row.Scan(
		&item.ID,
		&item.JobKey,
		&item.TriggerType,
		&item.Status,
		&item.RequestedBy,
		&item.Reason,
		&configRaw,
		&item.SchedulerInstanceID,
		&item.RunID,
		&item.RequestedAt,
		&item.ClaimedAt,
		&item.FinishedAt,
		&item.ErrorMessage,
	)
	if err != nil {
		return nil, err
	}
	item.ConfigOverride = decodeSchedulerJSONMap(configRaw)
	return &item, nil
}

func scanSchedulerInstance(row pgx.Row) (*domain.SchedulerInstance, error) {
	var item domain.SchedulerInstance
	var metadataRaw []byte
	if err := row.Scan(
		&item.InstanceID,
		&item.Hostname,
		&item.Process,
		&item.BuildVersion,
		&item.StartedAt,
		&item.LastSeenAt,
		&metadataRaw,
	); err != nil {
		return nil, err
	}
	item.Metadata = decodeSchedulerJSONMap(metadataRaw)
	return &item, nil
}

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

func decodeSchedulerJSONMap(raw []byte) map[string]any {
	if len(raw) == 0 {
		return map[string]any{}
	}
	var out map[string]any
	if err := json.Unmarshal(raw, &out); err != nil || out == nil {
		return map[string]any{}
	}
	return out
}

func nullableStringPtr(value *string) *string {
	if value == nil {
		return nil
	}
	trimmed := *value
	return &trimmed
}