package store

import (
	"context"
	"time"

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

const (
	TaskCallbackBatchSize = 100
	TaskCallbackLockTTL   = 60 * time.Second
)

type TaskCallbackDelivery struct {
	ID          string
	TaskID      string
	EventID     string
	Seq         int64
	CallbackURL string
	Status      string
	Attempts    int
	LockToken   string
	EventType   string
	TaskStatus  string
	CreatedAt   time.Time
}

type TaskHistoryCleanupResult struct {
	CompactedTasks        int64
	CompactedCheckpoints  int64
	CompactedAttempts     int64
	CompactedEvents       int64
	CompactedCallbacks    int64
	RetiredCallbacks      int64
	CompactedParamLogs    int64
	CompactedClonedVoices int64
	DeletedCallbacks      int64
	DeletedParamLogs      int64
	DeletedEvents         int64
	DeletedAttempts       int64
	DeletedTasks          int64
}

func (result TaskHistoryCleanupResult) Total() int64 {
	return result.CompactedTasks +
		result.CompactedCheckpoints +
		result.CompactedAttempts +
		result.CompactedEvents +
		result.CompactedCallbacks +
		result.RetiredCallbacks +
		result.CompactedParamLogs +
		result.CompactedClonedVoices +
		result.DeletedCallbacks +
		result.DeletedParamLogs +
		result.DeletedEvents +
		result.DeletedAttempts +
		result.DeletedTasks
}

func (s *Store) ClaimTaskCallbacks(ctx context.Context, workerID string, limit int, staleAfter time.Duration) ([]TaskCallbackDelivery, error) {
	if limit <= 0 || limit > TaskCallbackBatchSize {
		limit = TaskCallbackBatchSize
	}
	if staleAfter <= 0 {
		staleAfter = TaskCallbackLockTTL
	}
	rows, err := s.pool.Query(ctx, `
WITH picked AS (
  SELECT outbox.id, event.event_type, COALESCE(event.status, '') AS task_status, event.created_at
  FROM gateway_task_callback_outbox outbox
  JOIN gateway_task_events event ON event.id = outbox.event_id
  WHERE outbox.created_at >= COALESCE((
      SELECT applied_at
      FROM schema_migrations
      WHERE version = '0083_task_history_minimal_storage'
    ), now())
    AND (
      (
        outbox.status = 'pending'
        AND outbox.next_attempt_at <= now()
      )
      OR (
        outbox.status = 'processing'
        AND outbox.locked_at < now() - ($3::int * interval '1 second')
      )
    )
  ORDER BY outbox.next_attempt_at, outbox.created_at
  LIMIT $2
  FOR UPDATE OF outbox SKIP LOCKED
)
UPDATE gateway_task_callback_outbox outbox
SET status = 'processing',
    attempts = outbox.attempts + 1,
    locked_by = $1,
    lock_token = gen_random_uuid(),
    locked_at = now(),
    updated_at = now()
FROM picked
WHERE outbox.id = picked.id
RETURNING outbox.id::text, outbox.task_id::text, COALESCE(outbox.event_id::text, ''),
          outbox.seq, outbox.callback_url, outbox.status, outbox.attempts,
          COALESCE(outbox.lock_token::text, ''), picked.event_type, picked.task_status,
          picked.created_at`,
		workerID, limit, int(staleAfter/time.Second))
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	items := make([]TaskCallbackDelivery, 0)
	for rows.Next() {
		var item TaskCallbackDelivery
		if err := rows.Scan(
			&item.ID,
			&item.TaskID,
			&item.EventID,
			&item.Seq,
			&item.CallbackURL,
			&item.Status,
			&item.Attempts,
			&item.LockToken,
			&item.EventType,
			&item.TaskStatus,
			&item.CreatedAt,
		); err != nil {
			return nil, err
		}
		items = append(items, item)
	}
	return items, rows.Err()
}

func (s *Store) MarkTaskCallbackDelivered(ctx context.Context, item TaskCallbackDelivery) error {
	_, err := s.pool.Exec(ctx, `
UPDATE gateway_task_callback_outbox
SET status = 'delivered',
    delivered_at = now(),
    failed_at = NULL,
    last_error = NULL,
    locked_by = NULL,
    lock_token = NULL,
    locked_at = NULL,
    updated_at = now()
WHERE id = $1::uuid
  AND status = 'processing'
  AND lock_token = $2::uuid`, item.ID, item.LockToken)
	return err
}

func (s *Store) MarkTaskCallbackFailed(ctx context.Context, item TaskCallbackDelivery, retry bool, nextAttemptAt time.Time, message string) error {
	status := "failed"
	if retry {
		status = "pending"
	}
	_, err := s.pool.Exec(ctx, `
UPDATE gateway_task_callback_outbox
SET status = $3,
    next_attempt_at = CASE WHEN $3 = 'pending' THEN $4::timestamptz ELSE next_attempt_at END,
    failed_at = CASE WHEN $3 = 'failed' THEN now() ELSE NULL END,
    last_error = NULLIF(left($5, 2048), ''),
    locked_by = NULL,
    lock_token = NULL,
    locked_at = NULL,
    updated_at = now()
WHERE id = $1::uuid
  AND status = 'processing'
  AND lock_token = $2::uuid`,
		item.ID, item.LockToken, status, nextAttemptAt, message)
	return err
}

func (s *Store) CleanupTaskHistory(ctx context.Context, analysisCutoff time.Time, taskCutoff time.Time, batchSize int) (TaskHistoryCleanupResult, error) {
	if batchSize < 1 {
		batchSize = 1000
	}
	var result TaskHistoryCleanupResult
	err := pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error {
		if _, err := tx.Exec(ctx, `SET LOCAL lock_timeout = '5s'; SET LOCAL statement_timeout = '5s'`); err != nil {
			return err
		}
		steps := []struct {
			target *int64
			sql    string
			args   []any
		}{
			{&result.CompactedTasks, `
WITH picked AS (
  SELECT id
  FROM gateway_tasks
  WHERE COALESCE(normalized_request, '{}'::jsonb) <> '{}'::jsonb
     OR compatibility_public_id IS NOT NULL
     OR compatibility_submit_http_status IS NOT NULL
     OR COALESCE(compatibility_submit_headers, '{}'::jsonb) <> '{}'::jsonb
     OR COALESCE(compatibility_submit_body, '{}'::jsonb) <> '{}'::jsonb
     OR result ?| ARRAY[
       'raw', 'raw_data', 'rawData', 'provider_response', 'providerResponse',
       'submit', 'file_retrieve', 'fileRetrieve', 'upstream_task_id', 'remote_task_id'
     ]
     OR error IS NOT NULL
     OR octet_length(COALESCE(error_message, '')) > 2048
  ORDER BY updated_at
  LIMIT $1
  FOR UPDATE SKIP LOCKED
)
UPDATE gateway_tasks task
SET normalized_request = '{}'::jsonb,
    compatibility_public_id = NULL,
    compatibility_submit_http_status = NULL,
    compatibility_submit_headers = '{}'::jsonb,
    compatibility_submit_body = '{}'::jsonb,
    result = COALESCE(task.result, '{}'::jsonb)
      - 'raw' - 'raw_data' - 'rawData'
      - 'provider_response' - 'providerResponse'
      - 'submit' - 'file_retrieve' - 'fileRetrieve'
      - 'upstream_task_id' - 'remote_task_id',
    error = NULL,
    error_message = CASE
      WHEN octet_length(COALESCE(task.error_message, '')) > 2048 THEN left(task.error_message, 512)
      ELSE task.error_message
    END,
    remote_task_payload = CASE
      WHEN task.status IN ('succeeded', 'failed', 'cancelled') THEN '{}'::jsonb
      ELSE (
        SELECT COALESCE(jsonb_object_agg(entry.key, entry.value), '{}'::jsonb)
        FROM jsonb_each(COALESCE(task.remote_task_payload, '{}'::jsonb)) entry
        WHERE entry.key IN (
          'endpoint', 'pollEndpoint', 'phase', 'mode', 'taskType',
          'targetResolution', 'imageToken', 'receipt', 'cleanupElementIds'
        )
      )
    END
FROM picked
WHERE task.id = picked.id`, []any{batchSize}},
			{&result.CompactedCheckpoints, `
WITH picked AS (
  SELECT id
  FROM gateway_tasks
  WHERE COALESCE(remote_task_payload, '{}'::jsonb) <> '{}'::jsonb
    AND (
      status IN ('succeeded', 'failed', 'cancelled')
      OR EXISTS (
        SELECT 1
        FROM jsonb_object_keys(COALESCE(remote_task_payload, '{}'::jsonb)) AS key
        WHERE key NOT IN (
          'endpoint', 'pollEndpoint', 'phase', 'mode', 'taskType',
          'targetResolution', 'imageToken', 'receipt', 'cleanupElementIds'
        )
      )
    )
  ORDER BY updated_at
  LIMIT $1
  FOR UPDATE SKIP LOCKED
)
UPDATE gateway_tasks task
SET remote_task_payload = CASE
      WHEN task.status IN ('succeeded', 'failed', 'cancelled') THEN '{}'::jsonb
      ELSE (
        SELECT COALESCE(jsonb_object_agg(entry.key, entry.value), '{}'::jsonb)
        FROM jsonb_each(COALESCE(task.remote_task_payload, '{}'::jsonb)) entry
        WHERE entry.key IN (
          'endpoint', 'pollEndpoint', 'phase', 'mode', 'taskType',
          'targetResolution', 'imageToken', 'receipt', 'cleanupElementIds'
        )
      )
    END
FROM picked
WHERE task.id = picked.id`, []any{batchSize}},
			{&result.CompactedAttempts, `
WITH picked AS (
  SELECT id
  FROM gateway_task_attempts
  WHERE request_snapshot <> '{}'::jsonb
     OR COALESCE(response_snapshot, '{}'::jsonb) <> '{}'::jsonb
     OR COALESCE(usage, '{}'::jsonb) <> '{}'::jsonb
     OR COALESCE(metrics, '{}'::jsonb) <> '{}'::jsonb
     OR COALESCE(pricing_snapshot, '{}'::jsonb) <> '{}'::jsonb
     OR request_fingerprint IS NOT NULL
  ORDER BY started_at
  LIMIT $1
  FOR UPDATE SKIP LOCKED
)
UPDATE gateway_task_attempts attempt
SET request_snapshot = '{}'::jsonb,
    response_snapshot = '{}'::jsonb,
    usage = '{}'::jsonb,
    metrics = '{}'::jsonb,
    pricing_snapshot = '{}'::jsonb,
    request_fingerprint = NULL,
    error_message = NULLIF(left(error_message, 2048), '')
FROM picked
WHERE attempt.id = picked.id`, []any{batchSize}},
			{&result.CompactedEvents, `
WITH picked AS (
  SELECT id
  FROM gateway_task_events
  WHERE payload <> '{}'::jsonb
     OR phase IS NOT NULL
     OR COALESCE(progress, 0) <> 0
     OR message IS NOT NULL
  ORDER BY created_at
  LIMIT $1
  FOR UPDATE SKIP LOCKED
)
UPDATE gateway_task_events event
SET platform_id = CASE
      WHEN event.platform_id IS NULL
       AND event.payload->>'platformId' ~* '^[0-9a-f]{8}-[0-9a-f]{4}-[1-5][0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}$'
      THEN (event.payload->>'platformId')::uuid
      ELSE event.platform_id
    END,
    phase = NULL,
    progress = 0,
    message = NULL,
    payload = '{}'::jsonb
FROM picked
WHERE event.id = picked.id`, []any{batchSize}},
			{&result.CompactedCallbacks, `
WITH picked AS (
  SELECT id
  FROM gateway_task_callback_outbox
  WHERE payload <> '{}'::jsonb
  ORDER BY created_at
  LIMIT $1
  FOR UPDATE SKIP LOCKED
)
UPDATE gateway_task_callback_outbox outbox
SET payload = '{}'::jsonb
FROM picked
WHERE outbox.id = picked.id`, []any{batchSize}},
			{&result.RetiredCallbacks, `
WITH migration_boundary AS (
  SELECT applied_at
  FROM schema_migrations
  WHERE version = '0083_task_history_minimal_storage'
), picked AS (
  SELECT outbox.id
  FROM gateway_task_callback_outbox outbox
  CROSS JOIN migration_boundary boundary
  WHERE outbox.created_at < boundary.applied_at
    AND outbox.status IN ('pending', 'processing')
  ORDER BY outbox.created_at
  LIMIT $1
  FOR UPDATE OF outbox SKIP LOCKED
)
UPDATE gateway_task_callback_outbox outbox
SET status = 'failed',
    failed_at = outbox.created_at,
    last_error = 'legacy_callback_not_replayed',
    locked_by = NULL,
    lock_token = NULL,
    locked_at = NULL,
    updated_at = now()
FROM picked
WHERE outbox.id = picked.id`, []any{batchSize}},
			{&result.CompactedParamLogs, `
WITH picked AS (
  SELECT id
  FROM gateway_task_param_preprocessing_logs
  WHERE actual_input <> '{}'::jsonb
     OR converted_output <> '{}'::jsonb
     OR model_snapshot <> '{}'::jsonb
     OR octet_length(changes::text) > 1024
  ORDER BY created_at
  LIMIT $1
  FOR UPDATE SKIP LOCKED
)
UPDATE gateway_task_param_preprocessing_logs log
SET actual_input = '{}'::jsonb,
    converted_output = '{}'::jsonb,
    model_snapshot = '{}'::jsonb,
    changes = CASE WHEN octet_length(log.changes::text) > 1024 THEN '[]'::jsonb ELSE log.changes END
FROM picked
WHERE log.id = picked.id`, []any{batchSize}},
			{&result.CompactedClonedVoices, `
WITH picked AS (
  SELECT id
  FROM gateway_cloned_voices
  WHERE metadata ?| ARRAY['raw', 'raw_data', 'rawData', 'provider_response', 'providerResponse', 'request']
  ORDER BY updated_at
  LIMIT $1
  FOR UPDATE SKIP LOCKED
)
UPDATE gateway_cloned_voices voice
SET metadata = COALESCE(voice.metadata, '{}'::jsonb)
      - 'raw' - 'raw_data' - 'rawData'
      - 'provider_response' - 'providerResponse' - 'request',
    updated_at = now()
FROM picked
WHERE voice.id = picked.id`, []any{batchSize}},
			{&result.DeletedCallbacks, `
WITH picked AS (
  SELECT id
  FROM gateway_task_callback_outbox
  WHERE (status = 'delivered' AND delivered_at < now() - interval '24 hours')
     OR (status = 'failed' AND COALESCE(failed_at, updated_at) < $1::timestamptz)
  ORDER BY updated_at
  LIMIT $2
  FOR UPDATE SKIP LOCKED
)
DELETE FROM gateway_task_callback_outbox outbox
USING picked
WHERE outbox.id = picked.id`, []any{analysisCutoff, batchSize}},
			{&result.DeletedParamLogs, `
WITH picked AS (
  SELECT id
  FROM gateway_task_param_preprocessing_logs
  WHERE created_at < $1::timestamptz
  ORDER BY created_at
  LIMIT $2
  FOR UPDATE SKIP LOCKED
)
DELETE FROM gateway_task_param_preprocessing_logs log
USING picked
WHERE log.id = picked.id`, []any{analysisCutoff, batchSize}},
			{&result.DeletedEvents, `
WITH picked AS (
  SELECT event.id
  FROM gateway_task_events event
  WHERE event.created_at < $1::timestamptz
    AND NOT EXISTS (
      SELECT 1
      FROM gateway_task_callback_outbox outbox
      WHERE outbox.event_id = event.id
        AND outbox.status IN ('pending', 'processing')
    )
  ORDER BY event.created_at
  LIMIT $2
  FOR UPDATE OF event SKIP LOCKED
)
DELETE FROM gateway_task_events event
USING picked
WHERE event.id = picked.id`, []any{analysisCutoff, batchSize}},
			{&result.DeletedAttempts, `
WITH picked AS (
  SELECT id
  FROM gateway_task_attempts
  WHERE COALESCE(finished_at, started_at) < $1::timestamptz
    AND status <> 'running'
  ORDER BY COALESCE(finished_at, started_at)
  LIMIT $2
  FOR UPDATE SKIP LOCKED
)
DELETE FROM gateway_task_attempts attempt
USING picked
WHERE attempt.id = picked.id`, []any{analysisCutoff, batchSize}},
			{&result.DeletedTasks, `
WITH picked AS (
  SELECT task.id
  FROM gateway_tasks task
  WHERE task.finished_at < $1::timestamptz
    AND task.status IN ('succeeded', 'failed', 'cancelled')
    AND task.billing_status IN ('settled', 'released', 'not_required')
    AND (task.execution_lease_expires_at IS NULL OR task.execution_lease_expires_at <= now())
    AND NOT EXISTS (
      SELECT 1
      FROM settlement_outbox settlement
      WHERE settlement.task_id = task.id
        AND settlement.status IN ('pending', 'processing', 'retryable_failed', 'manual_review')
    )
    AND NOT EXISTS (
      SELECT 1
      FROM gateway_task_callback_outbox callback
      WHERE callback.task_id = task.id
    )
    AND NOT EXISTS (
      SELECT 1
      FROM gateway_concurrency_leases lease
      WHERE lease.task_id = task.id
        AND lease.released_at IS NULL
        AND lease.expires_at > now()
    )
    AND NOT EXISTS (
      SELECT 1
      FROM gateway_rate_limit_reservations reservation
      WHERE reservation.task_id = task.id
        AND reservation.status = 'reserved'
    )
  ORDER BY task.finished_at
  LIMIT $2
  FOR UPDATE OF task SKIP LOCKED
)
DELETE FROM gateway_tasks task
USING picked
WHERE task.id = picked.id`, []any{taskCutoff, batchSize}},
		}
		for _, step := range steps {
			tag, err := tx.Exec(ctx, step.sql, step.args...)
			if err != nil {
				return err
			}
			*step.target = tag.RowsAffected()
		}
		return nil
	})
	return result, err
}