package store

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

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

type TaskBinaryResultBackfillItem struct {
	ID         string
	Result     map[string]any
	FinishedAt time.Time
}

func (s *Store) ListTaskBinaryResultBackfillBatch(ctx context.Context, afterID string, batchSize int) ([]TaskBinaryResultBackfillItem, error) {
	if batchSize < 1 || batchSize > 100 {
		batchSize = 100
	}
	var items []TaskBinaryResultBackfillItem
	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
		}
		rows, err := tx.Query(ctx, `
SELECT id::text, result, COALESCE(finished_at, updated_at)
FROM gateway_tasks
WHERE status = 'succeeded'
  AND result <> '{}'::jsonb
  AND (NULLIF($1::text, '') IS NULL OR id > NULLIF($1::text, '')::uuid)
ORDER BY id
LIMIT $2`, afterID, batchSize)
		if err != nil {
			return err
		}
		defer rows.Close()
		items = make([]TaskBinaryResultBackfillItem, 0, batchSize)
		for rows.Next() {
			var item TaskBinaryResultBackfillItem
			var resultJSON []byte
			if err := rows.Scan(&item.ID, &resultJSON, &item.FinishedAt); err != nil {
				return err
			}
			item.Result = decodeObject(resultJSON)
			items = append(items, item)
		}
		return rows.Err()
	})
	return items, err
}

func (s *Store) UpdateTaskBinaryResultBackfill(ctx context.Context, taskID string, result map[string]any) (bool, error) {
	report := sanitizeJSONForStorageWithReport(minimalTaskResult(result))
	if report.BinaryCount > 0 {
		return false, &taskPayloadBinaryError{
			target: ErrTaskResultBinaryNotMaterialized,
			code:   "result_binary_not_materialized",
			count:  report.BinaryCount,
		}
	}
	resultJSON, err := json.Marshal(report.Value)
	if err != nil {
		return false, err
	}
	updated := false
	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
		}
		tag, err := tx.Exec(ctx, `
UPDATE gateway_tasks
SET result = $2::jsonb,
    updated_at = now()
WHERE id = $1::uuid
  AND status = 'succeeded'`,
			taskID,
			string(resultJSON),
		)
		if err != nil {
			return err
		}
		updated = tag.RowsAffected() == 1
		return nil
	})
	return updated, err
}