huangguo_server/common/shortrecommend/shortrecommend.go

package shortrecommend

import (
	"bytes"
	"context"
	cryptorand "crypto/rand"
	"crypto/sha256"
	"encoding/hex"
	"errors"
	"fmt"
	"hash/fnv"
	"math/rand"
	"sort"
	"strconv"
	"strings"
	"sync/atomic"
	"time"

	"91porn-server/common/redis"
	"91porn-server/models/v/vidmod"
)

const (
	CurrentKey              = "recommend:short:current"
	BuildLockKey            = "recommend:short:build-lock"
	defaultKeyTTL           = 72 * time.Hour
	minKeyTTLHours          = 24
	maxKeyTTLHours          = 30 * 24
	buildingKeyTTL          = time.Hour
	takeIdempotencyTTL      = 30 * time.Second
	reservationTTL          = 15 * time.Second
	maxTakeRequestIDLength  = 128
	healthTTLDriftTolerance = 2 * time.Second
	queuePushBatchSize      = 1000
	queuePipelineBatchCount = 16
	publishCleanupTimeout   = 5 * time.Second
)

var keyTTLSeconds atomic.Int64

func init() {
	keyTTLSeconds.Store(int64(defaultKeyTTL / time.Second))
}

// SetKeyTTLHours 配置正式队列、元数据和用户偏移的统一绝对过期时间。
func SetKeyTTLHours(hours int) error {
	if hours < minKeyTTLHours || hours > maxKeyTTLHours {
		return fmt.Errorf(
			"short recommend key TTL hours must be between %d and %d",
			minKeyTTLHours, maxKeyTTLHours,
		)
	}
	keyTTLSeconds.Store(int64(time.Duration(hours) * time.Hour / time.Second))
	return nil
}

// KeyTTL 返回当前并发安全的短推荐键过期时间配置。
func KeyTTL() time.Duration {
	seconds := keyTTLSeconds.Load()
	if seconds <= 0 {
		return defaultKeyTTL
	}
	return time.Duration(seconds) * time.Second
}

type QueueMeta struct {
	Version     string
	Length      int
	GeneratedAt time.Time
	StartOffset int
}

func QueueKey(version string) string  { return "recommend:short:queue:" + version }
func OffsetKey(version string) string { return "recommend:short:offset:" + version }
func MetaKey(version string) string   { return "recommend:short:meta:" + version }
func LockKey(_ string) string         { return BuildLockKey }

func VersionAt(t time.Time) string {
	return t.In(time.FixedZone("CST", 8*60*60)).Format("20060102")
}

// VersionWithRevision 为同一天的每次实际重建生成独立队列版本。
// revision 必须由调用方提供非空且唯一的构建标识。
func VersionWithRevision(t time.Time, revision string) string {
	return VersionAt(t) + "-" + revision
}

func versionMatchesDate(version, date string) bool {
	return version == date ||
		(len(version) > len(date)+1 && strings.HasPrefix(version, date+"-"))
}

type AssembleStats struct {
	TotalCount int
	HighCount  int
	NewCount   int
	BlockCount int
}

// Assemble 按每块17条高分+3条新视频组装，视频全局去重，尾块保留实际数量。
func Assemble(input []vidmod.RecommendCandidate, generatedAt time.Time) []string {
	ids, _ := AssembleWithStats(input, generatedAt)
	return ids
}

// AssembleWithStats 在组装队列的同时返回高分池、新视频池和块数指标。
func AssembleWithStats(
	input []vidmod.RecommendCandidate,
	generatedAt time.Time,
) ([]string, AssembleStats) {
	cutoff := generatedAt.Add(-24 * time.Hour)
	news := make([]*vidmod.RecommendCandidate, 0)
	high := make([]*vidmod.RecommendCandidate, 0)
	for i := range input {
		c := &input[i]
		if c.ReviewAt.After(generatedAt) {
			continue
		}
		if !c.ReviewAt.Before(cutoff) && !c.ReviewAt.After(generatedAt) {
			news = append(news, c)
		} else {
			high = append(high, c)
		}
	}
	sort.Slice(news, func(i, j int) bool {
		if news[i].ReviewAt.Equal(news[j].ReviewAt) {
			return bytes.Compare(news[i].ID[:], news[j].ID[:]) > 0
		}
		return news[i].ReviewAt.After(news[j].ReviewAt)
	})
	sort.Slice(high, func(i, j int) bool {
		if high[i].RecommendScore != high[j].RecommendScore {
			return high[i].RecommendScore > high[j].RecommendScore
		}
		if !high[i].ReviewAt.Equal(high[j].ReviewAt) {
			return high[i].ReviewAt.After(high[j].ReviewAt)
		}
		return bytes.Compare(high[i].ID[:], high[j].ID[:]) > 0
	})

	out := make([]string, 0, len(input))
	hi, ni, block := 0, 0, 0
	for hi < len(high) || ni < len(news) {
		blockStart := len(out)
		highCount := 17
		newCount := min(3, len(news)-ni)
		highCount += 3 - newCount
		highCount = min(highCount, len(high)-hi)
		for n := 0; n < highCount; n++ {
			out = append(out, high[hi].ID.Hex())
			hi++
		}
		for n := 0; n < newCount; n++ {
			out = append(out, news[ni].ID.Hex())
			ni++
		}
		if len(out) == blockStart { // 高分耗尽后继续消费剩余新视频，仍不复制凑数
			for ni < len(news) && len(out)-blockStart < 3 {
				out = append(out, news[ni].ID.Hex())
				ni++
			}
		}
		seed := deterministicSeed(VersionAt(generatedAt), block)
		part := out[blockStart:]
		rand.New(rand.NewSource(seed)).Shuffle(len(part), func(i, j int) {
			part[i], part[j] = part[j], part[i]
		})
		block++
	}
	return out, AssembleStats{
		TotalCount: len(out),
		HighCount:  len(high),
		NewCount:   len(news),
		BlockCount: block,
	}
}

func deterministicSeed(version string, block int) int64 {
	h := fnv.New64a()
	_, _ = h.Write([]byte(version + ":" + strconv.Itoa(block)))
	return int64(h.Sum64())
}

func min(a, b int) int {
	if a < b {
		return a
	}
	return b
}

type publishClient interface {
	DelContext(ctx context.Context, keys ...string) (int64, error)
	EvalContext(
		ctx context.Context,
		script string,
		keys []string,
		args ...interface{},
	) (interface{}, error)
	RPushPipelineContext(
		ctx context.Context,
		key string,
		values []string,
		batchSize int,
		expirations ...time.Duration,
	) (int64, error)
}

// Publish 完整写入新版本后再切换current，构建失败不会破坏旧版本。
func Publish(ctx context.Context, client *redis.Client, meta QueueMeta, ids []string, lockKey, lockToken string) error {
	return publish(ctx, client, meta, ids, lockKey, lockToken)
}

func publish(
	ctx context.Context,
	client publishClient,
	meta QueueMeta,
	ids []string,
	lockKey, lockToken string,
) error {
	if len(ids) == 0 {
		return fmt.Errorf("short recommend queue is empty")
	}
	if lockKey == "" || lockToken == "" {
		return fmt.Errorf("short recommend build lock is required")
	}
	if !versionMatchesDate(meta.Version, VersionAt(meta.GeneratedAt)) {
		return fmt.Errorf("short recommend version does not match generatedAt")
	}
	if meta.Length != len(ids) {
		return fmt.Errorf("short recommend metadata length does not match queue length")
	}
	queueTTL := KeyTTL()
	queueKey, metaKey := QueueKey(meta.Version), MetaKey(meta.Version)
	suffix := ":" + strconv.FormatInt(time.Now().UnixNano(), 10)
	tmpQueueKey, tmpMetaKey := queueKey+":building"+suffix, metaKey+":building"+suffix
	const firstPushScript = `
for i = 1, #ARGV - 1 do
  redis.call('RPUSH', KEYS[1], ARGV[i])
end
redis.call('EXPIRE', KEYS[1], ARGV[#ARGV])
return 1
`
	if err := ctx.Err(); err != nil {
		return err
	}
	firstEnd := min(queuePushBatchSize, len(ids))
	firstValues := make([]interface{}, firstEnd)
	for i := 0; i < firstEnd; i++ {
		firstValues[i] = ids[i]
	}
	firstArgs := append(firstValues, int64(buildingKeyTTL/time.Second))
	if _, err := client.EvalContext(ctx, firstPushScript, []string{tmpQueueKey}, firstArgs...); err != nil {
		cleanupPublishKeys(client, tmpQueueKey, tmpMetaKey)
		return err
	}

	pipelineWindowSize := queuePushBatchSize * queuePipelineBatchCount
	for start := firstEnd; start < len(ids); start += pipelineWindowSize {
		if err := ctx.Err(); err != nil {
			cleanupPublishKeys(client, tmpQueueKey, tmpMetaKey)
			return err
		}
		end := min(start+pipelineWindowSize, len(ids))
		length, err := client.RPushPipelineContext(
			ctx, tmpQueueKey, ids[start:end], queuePushBatchSize, buildingKeyTTL,
		)
		if err != nil {
			cleanupPublishKeys(client, tmpQueueKey, tmpMetaKey)
			return err
		}
		if length != int64(end) {
			cleanupPublishKeys(client, tmpQueueKey, tmpMetaKey)
			return fmt.Errorf(
				"short recommend queue length mismatch after pipeline: got %d, want %d",
				length, end,
			)
		}
	}
	if err := ctx.Err(); err != nil {
		cleanupPublishKeys(client, tmpQueueKey, tmpMetaKey)
		return err
	}
	const writeMetaScript = `
redis.call('HSET', KEYS[1],
  'length', ARGV[1],
  'generatedAt', ARGV[2],
  'startOffset', ARGV[3])
redis.call('EXPIRE', KEYS[1], ARGV[4])
return 1
`
	if _, err := client.EvalContext(ctx, writeMetaScript, []string{tmpMetaKey},
		len(ids),
		meta.GeneratedAt.UTC().Format(time.RFC3339),
		meta.StartOffset,
		int64(buildingKeyTTL/time.Second)); err != nil {
		cleanupPublishKeys(client, tmpQueueKey, tmpMetaKey)
		return err
	}
	if err := ctx.Err(); err != nil {
		cleanupPublishKeys(client, tmpQueueKey, tmpMetaKey)
		return err
	}
	const swapScript = `
if redis.call('GET', KEYS[7]) ~= ARGV[3] then
  return 0
end
local redisTime = redis.call('TIME')
if tonumber(redisTime[1]) >= tonumber(ARGV[4]) then
  return -1
end
if redis.call('LLEN', KEYS[1]) ~= tonumber(ARGV[5]) then
  return -2
end
if tonumber(redis.call('HGET', KEYS[2], 'length') or '-1') ~= tonumber(ARGV[5]) then
  return -3
end
local oldVersion = redis.call('GET', KEYS[6])
redis.call('UNLINK', KEYS[5])
if oldVersion and oldVersion ~= ARGV[1] then
  redis.call('UNLINK', ARGV[6] .. oldVersion)
end
redis.call('RENAME', KEYS[1], KEYS[3])
redis.call('RENAME', KEYS[2], KEYS[4])
local expiresAtMs =
  tonumber(redisTime[1]) * 1000 +
  math.floor(tonumber(redisTime[2]) / 1000) +
  tonumber(ARGV[2]) * 1000
redis.call('HSET', KEYS[4], 'expiresAtMs', expiresAtMs)
redis.call('PEXPIREAT', KEYS[3], expiresAtMs)
redis.call('PEXPIREAT', KEYS[4], expiresAtMs)
redis.call('SET', KEYS[6], ARGV[1])
return 1
`
	swapped, err := client.EvalContext(ctx, swapScript,
		[]string{
			tmpQueueKey, tmpMetaKey, queueKey, metaKey,
			OffsetKey(meta.Version), CurrentKey, lockKey,
		},
		meta.Version, int64(queueTTL/time.Second), lockToken,
		versionValidUntil(meta).Unix(), len(ids), "recommend:short:offset:")
	if err != nil || fmt.Sprint(swapped) != "1" {
		cleanupPublishKeys(client, tmpQueueKey, tmpMetaKey)
		if err != nil {
			return err
		}
		switch fmt.Sprint(swapped) {
		case "0":
			return fmt.Errorf("short recommend build lock lost before queue swap")
		case "-1":
			return fmt.Errorf("short recommend version expired before queue swap")
		case "-2":
			return fmt.Errorf("short recommend queue length mismatch before queue swap")
		case "-3":
			return fmt.Errorf("short recommend metadata length mismatch before queue swap")
		default:
			return fmt.Errorf("short recommend queue swap rejected: %v", swapped)
		}
	}
	return nil
}

func cleanupPublishKeys(client publishClient, keys ...string) {
	cleanupCtx, cancel := context.WithTimeout(context.Background(), publishCleanupTimeout)
	defer cancel()
	_, _ = client.DelContext(cleanupCtx, keys...)
}

func versionValidUntil(meta QueueMeta) time.Time {
	cst := time.FixedZone("CST", 8*60*60)
	generatedAt := meta.GeneratedAt.In(cst)
	return time.Date(
		generatedAt.Year(), generatedAt.Month(), generatedAt.Day()+1,
		0, 0, 0, 0, cst,
	)
}

var (
	// ErrVersionChanged 表示调用方固定的队列版本已被每日切版替换。
	ErrVersionChanged = errors.New("short recommend queue version changed")
	// ErrQueueUnhealthy 表示current指向的队列、元数据或TTL不完整。
	ErrQueueUnhealthy = errors.New("short recommend queue is unhealthy")
	// ErrReservationBusy 表示同一用户已有另一个请求正在读取当前偏移。
	ErrReservationBusy = errors.New("short recommend offset reservation is busy")
	// ErrReservationExpired 表示预留在提交前已过期，偏移未推进。
	ErrReservationExpired = errors.New("short recommend offset reservation expired")
	// ErrReservationConflict 表示预留内容或当前偏移与提交参数不一致。
	ErrReservationConflict = errors.New("short recommend offset reservation conflict")
)

const takeScriptSource = `
local function readBatch(queueKey, version, length, offset, wanted)
  local result = {version, tostring(length), tostring(offset)}
  local firstEnd = math.min(offset + wanted - 1, length - 1)
  local first = redis.call('LRANGE', queueKey, offset, firstEnd)
  for _, item in ipairs(first) do
    table.insert(result, item)
  end
  local remaining = wanted - #first
  if remaining > 0 then
    local wrapped = redis.call('LRANGE', queueKey, 0, remaining - 1)
    for _, item in ipairs(wrapped) do
      table.insert(result, item)
    end
  end
  return result
end
local version = redis.call('GET', KEYS[1])
if not version then return {} end
if ARGV[6] ~= '' and version ~= ARGV[6] then
  return {version, '-1', '0'}
end
local queueKey = ARGV[1] .. version
local offsetKey = ARGV[2] .. version
local metaKey = ARGV[3] .. version
local offsetKeyTTL = redis.call('PTTL', offsetKey)
local length = redis.call('LLEN', queueKey)
local metaLength = tonumber(redis.call('HGET', metaKey, 'length') or '-1')
local queueTTL = redis.call('PTTL', queueKey)
local metaTTL = redis.call('PTTL', metaKey)
if length == 0 or metaLength ~= length or queueTTL <= 0 or metaTTL <= 0 then
  return {version, '-2', '0'}
end
local idempotencyKey = ''
if ARGV[7] ~= '' then
  idempotencyKey = ARGV[7] .. ':' .. version
  local cached = redis.call('LRANGE', idempotencyKey, 0, -1)
  if #cached == 4 then
    local cachedVersion = cached[1]
    local cachedLength = tonumber(cached[2])
    local cachedOffset = tonumber(cached[3])
    local cachedWanted = tonumber(cached[4])
    if cachedVersion == version and cachedLength == length and
       cachedOffset and cachedWanted then
      return readBatch(
        queueKey, cachedVersion, cachedLength, cachedOffset, cachedWanted
      )
    end
  end
  if #cached > 0 then
    redis.call('DEL', idempotencyKey)
  end
end
local wanted = tonumber(ARGV[5])
if wanted > length then wanted = length end
local storedOffset = redis.call('HGET', offsetKey, ARGV[4])
local offset
if storedOffset then
  offset = tonumber(storedOffset) or 0
else
  offset = tonumber(redis.call('HGET', metaKey, 'startOffset') or '0')
end
offset = offset % length
local result = readBatch(queueKey, version, length, offset, wanted)
redis.call('HSET', offsetKey, ARGV[4], (offset + wanted) % length)
local expiresAtMs = tonumber(redis.call('HGET', metaKey, 'expiresAtMs') or '0')
if expiresAtMs <= 0 then
  local redisTime = redis.call('TIME')
  expiresAtMs =
    tonumber(redisTime[1]) * 1000 +
    math.floor(tonumber(redisTime[2]) / 1000) +
    math.min(queueTTL, metaTTL)
end
if offsetKeyTTL <= 0 then
  redis.call('PEXPIREAT', offsetKey, expiresAtMs)
end
if idempotencyKey ~= '' then
  redis.call(
    'RPUSH', idempotencyKey,
    version, tostring(length), tostring(offset), tostring(wanted)
  )
  redis.call('EXPIRE', idempotencyKey, tonumber(ARGV[8]))
end
return result
`

var takeRedisScript = redis.NewScript(takeScriptSource)

type TakeResult struct {
	Version string
	Length  int
	Offset  int
	IDs     []string
}

type scriptClient interface {
	RunScriptContext(
		ctx context.Context,
		script *redis.Script,
		keys []string,
		args ...interface{},
	) (interface{}, error)
}

// Take 原子读取并推进单个用户的下一偏移，保留原调用方式。
func Take(client *redis.Client, uid uint64, size int) (TakeResult, error) {
	return TakeVersionContext(context.Background(), client, uid, size, "")
}

// TakeVersion 原子读取指定版本；切版后返回 ErrVersionChanged 且不推进新队列。
func TakeVersion(
	client *redis.Client,
	uid uint64,
	size int,
	expectedVersion string,
) (TakeResult, error) {
	return TakeVersionContext(context.Background(), client, uid, size, expectedVersion)
}

// TakeVersionContext 与 TakeVersion 相同，并将请求取消和超时传递到Redis。
func TakeVersionContext(
	ctx context.Context,
	client *redis.Client,
	uid uint64,
	size int,
	expectedVersion string,
) (TakeResult, error) {
	return takeContext(ctx, client, uid, size, expectedVersion, "")
}

// TakeIdempotentContext 为一次HTTP请求的单个补位批次提供短时幂等保护。
// 同一uid/requestID/batchIndex重试会返回原批次，不会再次推进偏移；
// expectedVersion可在首批留空，后续批次应固定为首批返回的Version。
func TakeIdempotentContext(
	ctx context.Context,
	client *redis.Client,
	uid uint64,
	size int,
	expectedVersion, requestID string,
	batchIndex int,
) (TakeResult, error) {
	return takeIdempotentContext(
		ctx, client, uid, size, expectedVersion, requestID, batchIndex,
	)
}

// TakeCurrentIdempotentContext 用于首批读取：Lua原子读取current，并按该版本隔离幂等缓存，
// 因而不需要额外执行Health/GET；切版后不会复用上一版本缓存。
func TakeCurrentIdempotentContext(
	ctx context.Context,
	client *redis.Client,
	uid uint64,
	size int,
	requestID string,
	batchIndex int,
) (TakeResult, error) {
	return takeIdempotentContext(ctx, client, uid, size, "", requestID, batchIndex)
}

func takeIdempotentContext(
	ctx context.Context,
	client scriptClient,
	uid uint64,
	size int,
	expectedVersion, requestID string,
	batchIndex int,
) (TakeResult, error) {
	requestID = strings.TrimSpace(requestID)
	if requestID == "" {
		return TakeResult{}, fmt.Errorf("short recommend request ID is required")
	}
	if len(requestID) > maxTakeRequestIDLength {
		return TakeResult{}, fmt.Errorf(
			"short recommend request ID exceeds %d bytes",
			maxTakeRequestIDLength,
		)
	}
	if batchIndex < 0 {
		return TakeResult{}, fmt.Errorf("short recommend batch index must not be negative")
	}
	requestHash := sha256.Sum256([]byte(requestID))
	idempotencyKey := fmt.Sprintf(
		"recommend:short:take-cache:%d:%x:%d",
		uid, requestHash[:16], batchIndex,
	)
	return takeContext(ctx, client, uid, size, expectedVersion, idempotencyKey)
}

func takeContext(
	ctx context.Context,
	client scriptClient,
	uid uint64,
	size int,
	expectedVersion, idempotencyKey string,
) (TakeResult, error) {
	if size <= 0 {
		return TakeResult{}, nil
	}
	if ctx == nil {
		return TakeResult{}, fmt.Errorf("short recommend take context must not be nil")
	}
	if client == nil {
		return TakeResult{}, fmt.Errorf("short recommend take Redis client must not be nil")
	}
	if err := ctx.Err(); err != nil {
		return TakeResult{}, err
	}
	raw, err := client.RunScriptContext(ctx, takeRedisScript, []string{CurrentKey},
		"recommend:short:queue:",
		"recommend:short:offset:",
		"recommend:short:meta:",
		strconv.FormatUint(uid, 10),
		size,
		expectedVersion,
		idempotencyKey,
		int64(takeIdempotencyTTL/time.Second))
	if err != nil {
		return TakeResult{}, err
	}
	items, ok := raw.([]interface{})
	if !ok || len(items) < 3 {
		return TakeResult{}, nil
	}
	result := TakeResult{Version: fmt.Sprint(items[0])}
	if result.Version == "" {
		return TakeResult{}, fmt.Errorf("short recommend take returned empty version")
	}
	result.Length, err = strconv.Atoi(fmt.Sprint(items[1]))
	if err != nil {
		return TakeResult{}, fmt.Errorf("short recommend take returned invalid length: %w", err)
	}
	result.Offset, err = strconv.Atoi(fmt.Sprint(items[2]))
	if err != nil {
		return TakeResult{}, fmt.Errorf("short recommend take returned invalid offset: %w", err)
	}
	switch result.Length {
	case -1:
		result.Length = 0
		return result, ErrVersionChanged
	case -2:
		result.Length = 0
		return result, ErrQueueUnhealthy
	}
	if result.Length <= 0 || result.Offset < 0 || result.Offset >= result.Length {
		return TakeResult{}, fmt.Errorf(
			"short recommend take returned invalid bounds: length=%d offset=%d",
			result.Length, result.Offset,
		)
	}
	for _, item := range items[3:] {
		result.IDs = append(result.IDs, fmt.Sprint(item))
	}
	return result, nil
}

// Reservation 是一次请求对用户偏移的临时预留。Reserve 只读取队列，
// Commit 成功后才推进偏移；预留过期或 Abort 均不会改变用户进度。
type Reservation struct {
	Version          string
	Length           int
	Offset           int
	Reserved         int
	IDs              []string
	LeaseToken       string
	ReceiptID        string
	AlreadyCommitted bool
}

const reserveScriptSource = `
local function readBatch(queueKey, version, status, length, offset, wanted, token)
  local result = {
    version, status, tostring(length), tostring(offset), tostring(wanted), token
  }
  local firstEnd = math.min(offset + wanted - 1, length - 1)
  local first = redis.call('LRANGE', queueKey, offset, firstEnd)
  for _, item in ipairs(first) do
    table.insert(result, item)
  end
  local remaining = wanted - #first
  if remaining > 0 then
    local wrapped = redis.call('LRANGE', queueKey, 0, remaining - 1)
    for _, item in ipairs(wrapped) do
      table.insert(result, item)
    end
  end
  return result
end

local receiptKey = ARGV[6] .. ARGV[4] .. ':' .. ARGV[10]
local receipt = redis.call('LRANGE', receiptKey, 0, -1)
if #receipt == 4 then
  local receiptVersion = receipt[1]
  local receiptLength = tonumber(receipt[2])
  local receiptOffset = tonumber(receipt[3])
  local receiptConsumed = tonumber(receipt[4])
  local receiptQueueKey = ARGV[1] .. receiptVersion
  if receiptLength and receiptOffset and receiptConsumed and
     receiptLength > 0 and receiptConsumed > 0 and
     receiptConsumed <= receiptLength and
     redis.call('LLEN', receiptQueueKey) == receiptLength and
     redis.call('PTTL', receiptQueueKey) > 0 then
    return readBatch(
      receiptQueueKey, receiptVersion, 'COMMITTED',
      receiptLength, receiptOffset, receiptConsumed, ARGV[9]
    )
  end
  redis.call('DEL', receiptKey)
elseif #receipt > 0 then
  redis.call('DEL', receiptKey)
end

local version = redis.call('GET', KEYS[1])
if not version then return {} end
if ARGV[8] ~= '' and version ~= ARGV[8] then
  return {version, 'VERSION_CHANGED', '0', '0', '0', ''}
end
local queueKey = ARGV[1] .. version
local offsetKey = ARGV[2] .. version
local metaKey = ARGV[3] .. version
local length = redis.call('LLEN', queueKey)
local metaLength = tonumber(redis.call('HGET', metaKey, 'length') or '-1')
local queueTTL = redis.call('PTTL', queueKey)
local metaTTL = redis.call('PTTL', metaKey)
if length == 0 or metaLength ~= length or queueTTL <= 0 or metaTTL <= 0 then
  return {version, 'UNHEALTHY', '0', '0', '0', ''}
end

local leaseKey = ARGV[5] .. version .. ':' .. ARGV[4]
local lease = redis.call(
  'HMGET', leaseKey, 'token', 'receiptID', 'length', 'offset', 'reserved'
)
if lease[1] then
  if lease[1] ~= ARGV[9] or lease[2] ~= ARGV[10] then
    return {version, 'BUSY', tostring(length), '0', '0', ''}
  end
  local leaseLength = tonumber(lease[3])
  local leaseOffset = tonumber(lease[4])
  local leaseReserved = tonumber(lease[5])
  if leaseLength == length and leaseOffset and leaseReserved and
     leaseOffset >= 0 and leaseOffset < length and
     leaseReserved > 0 and leaseReserved <= length then
    redis.call('PEXPIRE', leaseKey, tonumber(ARGV[11]))
    return readBatch(
      queueKey, version, 'RESERVED',
      length, leaseOffset, leaseReserved, ARGV[9]
    )
  end
  redis.call('DEL', leaseKey)
end

local wanted = tonumber(ARGV[7])
if wanted > length then wanted = length end
local storedOffset = redis.call('HGET', offsetKey, ARGV[4])
local offset
if storedOffset then
  offset = tonumber(storedOffset) or 0
else
  offset = tonumber(redis.call('HGET', metaKey, 'startOffset') or '0')
end
offset = offset % length
redis.call('HSET', leaseKey,
  'token', ARGV[9],
  'receiptID', ARGV[10],
  'length', tostring(length),
  'offset', tostring(offset),
  'reserved', tostring(wanted))
redis.call('PEXPIRE', leaseKey, tonumber(ARGV[11]))
return readBatch(
  queueKey, version, 'RESERVED', length, offset, wanted, ARGV[9]
)
`

const commitReservationScriptSource = `
local receiptKey = ARGV[14] .. ARGV[5] .. ':' .. ARGV[8]
local receipt = redis.call('LRANGE', receiptKey, 0, -1)
if #receipt == 4 then
  if receipt[1] == ARGV[6] and
     tonumber(receipt[2]) == tonumber(ARGV[9]) and
     tonumber(receipt[3]) == tonumber(ARGV[10]) and
     tonumber(receipt[4]) == tonumber(ARGV[12]) then
    return 2
  end
  return -5
elseif #receipt > 0 then
  return -5
end

local version = redis.call('GET', KEYS[1])
if version ~= ARGV[6] then return -1 end
local queueKey = ARGV[1] .. version
local offsetKey = ARGV[2] .. version
local metaKey = ARGV[3] .. version
local length = redis.call('LLEN', queueKey)
local metaLength = tonumber(redis.call('HGET', metaKey, 'length') or '-1')
if length <= 0 or metaLength ~= length or length ~= tonumber(ARGV[9]) or
   redis.call('PTTL', queueKey) <= 0 or redis.call('PTTL', metaKey) <= 0 then
  return -2
end

local leaseKey = ARGV[4] .. version .. ':' .. ARGV[5]
local lease = redis.call(
  'HMGET', leaseKey, 'token', 'receiptID', 'length', 'offset', 'reserved'
)
if not lease[1] then return -3 end
if lease[1] ~= ARGV[7] or lease[2] ~= ARGV[8] or
   tonumber(lease[3]) ~= tonumber(ARGV[9]) or
   tonumber(lease[4]) ~= tonumber(ARGV[10]) or
   tonumber(lease[5]) ~= tonumber(ARGV[11]) then
  return -4
end
local consumed = tonumber(ARGV[12])
if not consumed or consumed <= 0 or consumed > tonumber(lease[5]) then
  return -4
end

local storedOffset = redis.call('HGET', offsetKey, ARGV[5])
local currentOffset
if storedOffset then
  currentOffset = tonumber(storedOffset) or 0
else
  currentOffset = tonumber(redis.call('HGET', metaKey, 'startOffset') or '0')
end
currentOffset = currentOffset % length
if currentOffset ~= tonumber(lease[4]) then return -4 end

redis.call('HSET', offsetKey, ARGV[5], (currentOffset + consumed) % length)
local offsetKeyTTL = redis.call('PTTL', offsetKey)
local expiresAtMs = tonumber(redis.call('HGET', metaKey, 'expiresAtMs') or '0')
if expiresAtMs <= 0 then
  local redisTime = redis.call('TIME')
  expiresAtMs =
    tonumber(redisTime[1]) * 1000 +
    math.floor(tonumber(redisTime[2]) / 1000) +
    math.min(redis.call('PTTL', queueKey), redis.call('PTTL', metaKey))
end
if offsetKeyTTL <= 0 then
  redis.call('PEXPIREAT', offsetKey, expiresAtMs)
end
redis.call('DEL', leaseKey)
redis.call(
  'RPUSH', receiptKey,
  version, tostring(length), tostring(currentOffset), tostring(consumed)
)
redis.call('EXPIRE', receiptKey, tonumber(ARGV[13]))
return 1
`

const abortReservationScriptSource = `
local leaseKey = ARGV[1] .. ARGV[2] .. ':' .. ARGV[3]
local lease = redis.call('HMGET', leaseKey, 'token', 'receiptID')
if lease[1] ~= ARGV[4] or lease[2] ~= ARGV[5] then
  return 0
end
return redis.call('DEL', leaseKey)
`

var (
	reserveRedisScript           = redis.NewScript(reserveScriptSource)
	commitReservationRedisScript = redis.NewScript(commitReservationScriptSource)
	abortReservationRedisScript  = redis.NewScript(abortReservationScriptSource)
)

// ReserveCurrentContext 原子预留当前版本的一段队列，但不推进用户偏移。
func ReserveCurrentContext(
	ctx context.Context,
	client *redis.Client,
	uid uint64,
	size int,
	requestID string,
) (Reservation, error) {
	if client == nil {
		return Reservation{}, fmt.Errorf("short recommend reserve Redis client must not be nil")
	}
	return reserveContext(ctx, client, uid, size, "", requestID)
}

// ReserveVersionContext 与 ReserveCurrentContext 相同，但要求current仍为指定版本。
func ReserveVersionContext(
	ctx context.Context,
	client *redis.Client,
	uid uint64,
	size int,
	expectedVersion, requestID string,
) (Reservation, error) {
	if client == nil {
		return Reservation{}, fmt.Errorf("short recommend reserve Redis client must not be nil")
	}
	return reserveContext(ctx, client, uid, size, expectedVersion, requestID)
}

func reserveContext(
	ctx context.Context,
	client scriptClient,
	uid uint64,
	size int,
	expectedVersion, requestID string,
) (Reservation, error) {
	if size <= 0 {
		return Reservation{}, nil
	}
	if ctx == nil {
		return Reservation{}, fmt.Errorf("short recommend reserve context must not be nil")
	}
	if client == nil {
		return Reservation{}, fmt.Errorf("short recommend reserve Redis client must not be nil")
	}
	if uid == 0 {
		return Reservation{}, fmt.Errorf("short recommend reserve uid must be positive")
	}
	if err := ctx.Err(); err != nil {
		return Reservation{}, err
	}
	requestID = strings.TrimSpace(requestID)
	if len(requestID) > maxTakeRequestIDLength {
		return Reservation{}, fmt.Errorf(
			"short recommend request ID exceeds %d bytes",
			maxTakeRequestIDLength,
		)
	}
	leaseToken, err := newReservationToken()
	if err != nil {
		return Reservation{}, err
	}
	receiptID := leaseToken
	if requestID != "" {
		sum := sha256.Sum256([]byte(strconv.FormatUint(uid, 10) + ":" + requestID))
		receiptID = hex.EncodeToString(sum[:16])
	}
	keys := []string{CurrentKey}
	args := []interface{}{
		"recommend:short:queue:",
		"recommend:short:offset:",
		"recommend:short:meta:",
		strconv.FormatUint(uid, 10),
		"recommend:short:reservation:",
		"recommend:short:commit-receipt:",
		size,
		expectedVersion,
		leaseToken,
		receiptID,
		int64(reservationTTL / time.Millisecond),
	}
	raw, err := client.RunScriptContext(ctx, reserveRedisScript, keys, args...)
	if err != nil && ctx.Err() == nil {
		// 预留脚本可能已经执行成功但响应在网络中丢失。使用完全相同的
		// token/receipt 重试可命中原租约，不会创建第二个预留或推进偏移。
		raw, err = client.RunScriptContext(ctx, reserveRedisScript, keys, args...)
	}
	if err != nil {
		return Reservation{}, err
	}
	return parseReservation(raw, receiptID)
}

func newReservationToken() (string, error) {
	value := make([]byte, 16)
	if _, err := cryptorand.Read(value); err != nil {
		return "", fmt.Errorf("generate short recommend reservation token: %w", err)
	}
	return hex.EncodeToString(value), nil
}

func parseReservation(raw interface{}, receiptID string) (Reservation, error) {
	items, ok := raw.([]interface{})
	if !ok || len(items) == 0 {
		return Reservation{}, nil
	}
	if len(items) < 6 {
		return Reservation{}, fmt.Errorf("short recommend reserve returned malformed result")
	}
	result := Reservation{
		Version:    fmt.Sprint(items[0]),
		LeaseToken: fmt.Sprint(items[5]),
		ReceiptID:  receiptID,
	}
	status := fmt.Sprint(items[1])
	var err error
	result.Length, err = strconv.Atoi(fmt.Sprint(items[2]))
	if err != nil {
		return Reservation{}, fmt.Errorf("short recommend reserve returned invalid length: %w", err)
	}
	result.Offset, err = strconv.Atoi(fmt.Sprint(items[3]))
	if err != nil {
		return Reservation{}, fmt.Errorf("short recommend reserve returned invalid offset: %w", err)
	}
	result.Reserved, err = strconv.Atoi(fmt.Sprint(items[4]))
	if err != nil {
		return Reservation{}, fmt.Errorf("short recommend reserve returned invalid size: %w", err)
	}
	switch status {
	case "VERSION_CHANGED":
		return result, ErrVersionChanged
	case "UNHEALTHY":
		return result, ErrQueueUnhealthy
	case "BUSY":
		return result, ErrReservationBusy
	case "COMMITTED":
		result.AlreadyCommitted = true
	case "RESERVED":
	default:
		return Reservation{}, fmt.Errorf("short recommend reserve returned unknown status %q", status)
	}
	if result.Version == "" || result.Length <= 0 ||
		result.Offset < 0 || result.Offset >= result.Length ||
		result.Reserved <= 0 || result.Reserved > result.Length ||
		len(items)-6 != result.Reserved {
		return Reservation{}, fmt.Errorf(
			"short recommend reserve returned invalid bounds: version=%q length=%d offset=%d reserved=%d ids=%d",
			result.Version, result.Length, result.Offset, result.Reserved, len(items)-6,
		)
	}
	for _, item := range items[6:] {
		result.IDs = append(result.IDs, fmt.Sprint(item))
	}
	return result, nil
}

// CommitReservationContext 原子校验预留并推进实际扫描的条数。
func CommitReservationContext(
	ctx context.Context,
	client *redis.Client,
	uid uint64,
	reservation Reservation,
	consumed int,
) error {
	if client == nil {
		return fmt.Errorf("short recommend commit Redis client must not be nil")
	}
	return commitReservationContext(ctx, client, uid, reservation, consumed)
}

func commitReservationContext(
	ctx context.Context,
	client scriptClient,
	uid uint64,
	reservation Reservation,
	consumed int,
) error {
	if ctx == nil {
		return fmt.Errorf("short recommend commit context must not be nil")
	}
	if client == nil {
		return fmt.Errorf("short recommend commit Redis client must not be nil")
	}
	if err := validateReservationCommit(uid, reservation, consumed); err != nil {
		return err
	}
	if err := ctx.Err(); err != nil {
		return err
	}
	raw, err := client.RunScriptContext(
		ctx,
		commitReservationRedisScript,
		[]string{CurrentKey},
		"recommend:short:queue:",
		"recommend:short:offset:",
		"recommend:short:meta:",
		"recommend:short:reservation:",
		strconv.FormatUint(uid, 10),
		reservation.Version,
		reservation.LeaseToken,
		reservation.ReceiptID,
		reservation.Length,
		reservation.Offset,
		reservation.Reserved,
		consumed,
		int64(takeIdempotencyTTL/time.Second),
		"recommend:short:commit-receipt:",
	)
	if err != nil {
		return err
	}
	switch fmt.Sprint(raw) {
	case "1", "2":
		return nil
	case "-1":
		return ErrVersionChanged
	case "-2":
		return ErrQueueUnhealthy
	case "-3":
		return ErrReservationExpired
	case "-4", "-5":
		return ErrReservationConflict
	default:
		return fmt.Errorf("short recommend commit returned unexpected result %v", raw)
	}
}

func validateReservationCommit(uid uint64, reservation Reservation, consumed int) error {
	if uid == 0 {
		return fmt.Errorf("short recommend commit uid must be positive")
	}
	if reservation.Version == "" || reservation.Length <= 0 ||
		reservation.Offset < 0 || reservation.Offset >= reservation.Length ||
		reservation.Reserved <= 0 || reservation.Reserved > reservation.Length ||
		reservation.LeaseToken == "" || reservation.ReceiptID == "" {
		return fmt.Errorf("short recommend commit reservation is invalid")
	}
	if consumed <= 0 || consumed > reservation.Reserved {
		return fmt.Errorf("short recommend consumed size is out of reservation bounds")
	}
	return nil
}

// AbortReservationContext 仅删除token仍匹配的未提交预留，不改变用户偏移。
func AbortReservationContext(
	ctx context.Context,
	client *redis.Client,
	uid uint64,
	reservation Reservation,
) error {
	if client == nil {
		return fmt.Errorf("short recommend abort Redis client must not be nil")
	}
	return abortReservationContext(ctx, client, uid, reservation)
}

func abortReservationContext(
	ctx context.Context,
	client scriptClient,
	uid uint64,
	reservation Reservation,
) error {
	if ctx == nil {
		return fmt.Errorf("short recommend abort context must not be nil")
	}
	if client == nil {
		return fmt.Errorf("short recommend abort Redis client must not be nil")
	}
	if uid == 0 || reservation.Version == "" ||
		reservation.LeaseToken == "" || reservation.ReceiptID == "" {
		return nil
	}
	if err := ctx.Err(); err != nil {
		return err
	}
	_, err := client.RunScriptContext(
		ctx,
		abortReservationRedisScript,
		nil,
		"recommend:short:reservation:",
		reservation.Version,
		strconv.FormatUint(uid, 10),
		reservation.LeaseToken,
		reservation.ReceiptID,
	)
	return err
}

type QueueHealth struct {
	Healthy        bool
	Status         string
	Version        string
	QueueLength    int
	MetadataLength int
	QueueTTL       time.Duration
	MetadataTTL    time.Duration
}

const healthScriptSource = `
local version = redis.call('GET', KEYS[1])
if not version then
  return {'NO_CURRENT', '', '0', '-1', '-2', '-2'}
end
if ARGV[1] ~= '' then
  local expectedDate = ARGV[1]
  local matchesDate =
    version == expectedDate or
    (
      string.len(version) > string.len(expectedDate) + 1 and
      string.sub(version, 1, string.len(expectedDate) + 1) == expectedDate .. '-'
    )
  if not matchesDate then
    return {'VERSION_MISMATCH', version, '0', '-1', '-2', '-2'}
  end
end
local queueKey = ARGV[2] .. version
local metaKey = ARGV[3] .. version
local queueLength = redis.call('LLEN', queueKey)
local metadataLength = tonumber(redis.call('HGET', metaKey, 'length') or '-1')
local queueTTL = redis.call('PTTL', queueKey)
local metadataTTL = redis.call('PTTL', metaKey)
local status = 'OK'
if queueLength <= 0 then
  status = 'QUEUE_EMPTY'
elseif metadataLength < 0 then
  status = 'META_MISSING'
elseif metadataLength ~= queueLength then
  status = 'LENGTH_MISMATCH'
elseif queueTTL <= 0 or metadataTTL <= 0 then
  status = 'TTL_MISSING'
elseif math.abs(queueTTL - metadataTTL) > tonumber(ARGV[4]) then
  status = 'TTL_MISMATCH'
end
return {
  status,
  version,
  tostring(queueLength),
  tostring(metadataLength),
  tostring(queueTTL),
  tostring(metadataTTL)
}
`

var healthRedisScript = redis.NewScript(healthScriptSource)

// Health 原子检查current、队列长度、metadata.length和两者TTL。
func Health(client *redis.Client, expectedVersion string) (QueueHealth, error) {
	return HealthContext(context.Background(), client, expectedVersion)
}

// HealthContext 与 Health 相同，并将取消和超时传递到Redis。
func HealthContext(
	ctx context.Context,
	client *redis.Client,
	expectedVersion string,
) (QueueHealth, error) {
	return healthContext(ctx, client, expectedVersion)
}

func healthContext(
	ctx context.Context,
	client scriptClient,
	expectedVersion string,
) (QueueHealth, error) {
	if ctx == nil {
		return QueueHealth{}, fmt.Errorf("short recommend health context must not be nil")
	}
	if client == nil {
		return QueueHealth{}, fmt.Errorf("short recommend health Redis client must not be nil")
	}
	if err := ctx.Err(); err != nil {
		return QueueHealth{}, err
	}
	raw, err := client.RunScriptContext(ctx, healthRedisScript, []string{CurrentKey},
		expectedVersion,
		"recommend:short:queue:",
		"recommend:short:meta:",
		int64(healthTTLDriftTolerance/time.Millisecond))
	if err != nil {
		return QueueHealth{}, err
	}
	items, ok := raw.([]interface{})
	if !ok || len(items) != 6 {
		return QueueHealth{}, fmt.Errorf("short recommend health returned malformed result")
	}
	health := QueueHealth{
		Status:  fmt.Sprint(items[0]),
		Version: fmt.Sprint(items[1]),
	}
	health.QueueLength, err = strconv.Atoi(fmt.Sprint(items[2]))
	if err != nil {
		return QueueHealth{}, fmt.Errorf("short recommend health returned invalid queue length: %w", err)
	}
	health.MetadataLength, err = strconv.Atoi(fmt.Sprint(items[3]))
	if err != nil {
		return QueueHealth{}, fmt.Errorf("short recommend health returned invalid metadata length: %w", err)
	}
	queueTTLMillis, err := strconv.ParseInt(fmt.Sprint(items[4]), 10, 64)
	if err != nil {
		return QueueHealth{}, fmt.Errorf("short recommend health returned invalid queue TTL: %w", err)
	}
	metadataTTLMillis, err := strconv.ParseInt(fmt.Sprint(items[5]), 10, 64)
	if err != nil {
		return QueueHealth{}, fmt.Errorf("short recommend health returned invalid metadata TTL: %w", err)
	}
	health.QueueTTL = time.Duration(queueTTLMillis) * time.Millisecond
	health.MetadataTTL = time.Duration(metadataTTLMillis) * time.Millisecond
	health.Healthy = health.Status == "OK"
	return health, nil
}