tencent-baidu-tracking/src/main/java/com/adx/tencent/storage/RedisHotStore.java

package com.adx.tencent.storage;

import com.adx.tencent.storage.model.BidRecord;
import com.adx.tencent.storage.model.QueuedEvent;
import com.adx.tencent.storage.model.TrackingRecord;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.data.redis.connection.stream.*;
import org.springframework.data.redis.core.StringRedisTemplate;

import java.time.Duration;
import java.time.Instant;
import java.nio.charset.StandardCharsets;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 对应 Go internal/storage/redis.go 的 RedisHotStore。
 * 职责：
 *   1. bid 记录缓存（adx:bid:{qk} / adx:media:{media}:{traceId}）
 *   2. Redis Stream 事件队列（adx:events）
 *   3. BidFinder 接口
 */
public class RedisHotStore {

    public static final String BID_NOT_FOUND_MSG = "bid not found";
    private static final Duration DEFAULT_CONVERSION_SEEN_TTL = Duration.ofDays(7);
    private static final Duration DEFAULT_DEDUCTION_COUNTER_TTL = Duration.ofDays(7);

    private static final String DEFAULT_PREFIX = "adx:";
    private static final String DEFAULT_STREAM  = "adx:events";

    private final StringRedisTemplate redis;
    private final ObjectMapper objectMapper;
    private final TiDBColdStore coldStore;
    private final String prefix;
    private final String stream;
    private final Duration bidTtl;
    private final long streamMaxLen;
    private final DefaultRedisScript<List> deductionDecisionScript;
    private final Set<String> initializedGroups = ConcurrentHashMap.newKeySet();

    public RedisHotStore(StringRedisTemplate redis, ObjectMapper objectMapper,
                         TiDBColdStore coldStore, String prefix, String stream, Duration bidTtl,
                         long streamMaxLen) {
        this.redis = redis;
        this.objectMapper = objectMapper;
        this.coldStore = coldStore;
        this.prefix = (prefix == null || prefix.isBlank()) ? DEFAULT_PREFIX : prefix;
        this.stream = (stream == null || stream.isBlank()) ? DEFAULT_STREAM : stream;
        this.bidTtl = (bidTtl == null || bidTtl.isZero()) ? Duration.ofHours(24) : bidTtl;
        this.streamMaxLen = streamMaxLen;
        this.deductionDecisionScript = buildDeductionDecisionScript();
    }

    // ─── EventRecorder ───────────────────────────────────────────────────────

    /**
     * 对应 Go RecordBid：
     *   Pipeline: SET bid:{qk} + SET media:{media}:{traceId} + XADD
     */
    public void recordBid(BidRecord record) {
        if (record.getCreatedAt() == null) {
            record.setCreatedAt(Instant.now());
        }
        try {
            // Redis KEY 只存关键字段（瘦身），降低内存占用
            String hotPayload = objectMapper.writeValueAsString(record.toHotRecord());
            // Stream 存全量数据，供 ColdWorker 持久化到 TiDB
            String fullPayload = objectMapper.writeValueAsString(record);

            redis.executePipelined((org.springframework.data.redis.connection.RedisConnection conn) -> {
                trimStreamIfNeeded(conn);
                byte[] hotBytes = hotPayload.getBytes(StandardCharsets.UTF_8);
                long ttlMillis = bidTtl.toMillis();
                if (record.getQk() != null && !record.getQk().isEmpty()) {
                    conn.stringCommands().set(
                            bidKey(record.getQk()).getBytes(StandardCharsets.UTF_8),
                            hotBytes,
                            org.springframework.data.redis.core.types.Expiration.milliseconds(ttlMillis),
                            org.springframework.data.redis.connection.RedisStringCommands.SetOption.UPSERT
                    );
                }
                if (record.getMedia() != null && !record.getMedia().isEmpty()
                        && record.getMediaTraceId() != null && !record.getMediaTraceId().isEmpty()) {
                    conn.stringCommands().set(
                            mediaTraceKey(record.getMedia(), record.getMediaTraceId()).getBytes(StandardCharsets.UTF_8),
                            hotBytes,
                            org.springframework.data.redis.core.types.Expiration.milliseconds(ttlMillis),
                            org.springframework.data.redis.connection.RedisStringCommands.SetOption.UPSERT
                    );
                }
                // XADD 全量数据到 Stream
                byte[] streamKey = stream.getBytes(StandardCharsets.UTF_8);
                Map<byte[], byte[]> body = new LinkedHashMap<>();
                body.put("type".getBytes(StandardCharsets.UTF_8), "bid".getBytes(StandardCharsets.UTF_8));
                body.put("payload".getBytes(StandardCharsets.UTF_8), fullPayload.getBytes(StandardCharsets.UTF_8));
                conn.streamCommands().xAdd(MapRecord.create(streamKey, body));
                return null;
            });
        } catch (Exception e) {
            throw new RuntimeException("recordBid failed", e);
        }
    }

    /**
     * 对应 Go RecordTracking：XADD 到 Stream。
     */
    public void recordTracking(TrackingRecord record) {
        if (record.getCreatedAt() == null) {
            record.setCreatedAt(Instant.now());
        }
        try {
            String payload = objectMapper.writeValueAsString(record);
            trim(streamMaxLen);
            redis.opsForStream().add(
                    MapRecord.create(stream, Map.of("type", "tracking", "payload", payload))
                            .withStreamKey(stream)
            );
        } catch (Exception e) {
            throw new RuntimeException("recordTracking failed", e);
        }
    }

    // ─── BidFinder ───────────────────────────────────────────────────────────

    /**
     * 点击处理完成后，将 Redis 中的 bid 记录缩减为仅保留转化匹配所需字段。
     * 清除 clickUrls、landingPage、appStoreLink 以释放内存。
     */
    public void shrinkAfterClick(BidRecord record) {
        try {
            String slimPayload = objectMapper.writeValueAsString(record.toPostClickRecord());
            byte[] slimBytes = slimPayload.getBytes(StandardCharsets.UTF_8);
            // 获取剩余 TTL，保持原有过期时间
            Long remainingTtl = redis.getExpire(bidKey(record.getQk()), java.util.concurrent.TimeUnit.MILLISECONDS);
            if (remainingTtl == null || remainingTtl <= 0) remainingTtl = bidTtl.toMillis();
            long ttlMs = remainingTtl;

            redis.executePipelined((org.springframework.data.redis.connection.RedisConnection conn) -> {
                if (record.getQk() != null && !record.getQk().isEmpty()) {
                    conn.stringCommands().set(
                            bidKey(record.getQk()).getBytes(StandardCharsets.UTF_8),
                            slimBytes,
                            org.springframework.data.redis.core.types.Expiration.milliseconds(ttlMs),
                            org.springframework.data.redis.connection.RedisStringCommands.SetOption.UPSERT
                    );
                }
                if (record.getMedia() != null && !record.getMedia().isEmpty()
                        && record.getMediaTraceId() != null && !record.getMediaTraceId().isEmpty()) {
                    conn.stringCommands().set(
                            mediaTraceKey(record.getMedia(), record.getMediaTraceId()).getBytes(StandardCharsets.UTF_8),
                            slimBytes,
                            org.springframework.data.redis.core.types.Expiration.milliseconds(ttlMs),
                            org.springframework.data.redis.connection.RedisStringCommands.SetOption.UPSERT
                    );
                }
                return null;
            });
        } catch (Exception e) {
            // 缩减失败不影响主流程，仅记录日志
            throw new RuntimeException("shrinkAfterClick failed", e);
        }
    }

    public BidRecord findBidByQk(String qk) {
        String payload = redis.opsForValue().get(bidKey(qk));
        if (payload != null) {
            return parseBidRecord(payload);
        }
        if (coldStore != null) {
            return coldStore.getBidByQk(qk);
        }
        return null;
    }

    public Map<String, BidRecord> findBidsByQks(List<String> qks) {
        if (qks == null || qks.isEmpty()) return Collections.emptyMap();
        Map<String, BidRecord> result = new HashMap<>();
        List<String> keys = new ArrayList<>(qks.size());
        Map<String, String> keyToQk = new HashMap<>();
        for (String qk : qks) {
            if (qk == null || qk.isBlank() || keyToQk.containsValue(qk)) continue;
            String key = bidKey(qk);
            keys.add(key);
            keyToQk.put(key, qk);
        }
        List<String> payloads = redis.opsForValue().multiGet(keys);
        Set<String> missing = new LinkedHashSet<>();
        for (int i = 0; i < keys.size(); i++) {
            String payload = payloads != null ? payloads.get(i) : null;
            String qk = keyToQk.get(keys.get(i));
            if (payload != null) {
                result.put(qk, parseBidRecord(payload));
            } else if (qk != null) {
                missing.add(qk);
            }
        }
        if (!missing.isEmpty() && coldStore != null) {
            for (BidRecord bid : coldStore.getBidsByQks(new ArrayList<>(missing))) {
                if (bid != null && bid.getQk() != null) {
                    result.put(bid.getQk(), bid);
                }
            }
        }
        return result;
    }

    public BidRecord findBidByMediaTrace(String media, String traceId) {
        String payload = redis.opsForValue().get(mediaTraceKey(media, traceId));
        if (payload == null) return null;
        return parseBidRecord(payload);
    }

    /**
     * 首次见到某条转化时返回 true，并写入带 TTL 的判重标记。
     * 使用 SETNX 保证多实例并发下只会有一个消费者把它视为“新转化”。
     */
    public boolean markConversionSeen(String dedupeKey) {
        if (dedupeKey == null || dedupeKey.isBlank()) {
            return false;
        }
        Boolean created = redis.opsForValue().setIfAbsent(
                conversionSeenKey(dedupeKey),
                "1",
                DEFAULT_CONVERSION_SEEN_TTL
        );
        return Boolean.TRUE.equals(created);
    }

    /**
     * 原子计算一条新转化是否应被扣量。
     * 规则：任意时点都尽量逼近目标比例，例如 50 条时扣 25，100 条时扣 50。
     */
    public DeductionDecision evaluateDeduction(String accountId, int act, String trackingVersion, String date, int deductionRate) {
        String key = deductionCounterKey(accountId, act, trackingVersion, date);
        List result = redis.execute(
                deductionDecisionScript,
                Collections.singletonList(key),
                String.valueOf(Math.max(deductionRate, 0)),
                String.valueOf(DEFAULT_DEDUCTION_COUNTER_TTL.getSeconds())
        );
        if (result == null || result.size() < 3) {
            throw new RuntimeException("evaluate deduction failed");
        }
        long total = toLong(result.get(0));
        long deducted = toLong(result.get(1));
        boolean shouldDeduct = toLong(result.get(2)) == 1L;
        return new DeductionDecision(total, deducted, shouldDeduct);
    }

    // ─── EventQueue ──────────────────────────────────────────────────────────

    /**
     * 对应 Go EnsureGroup：XGROUP CREATE MKSTREAM（忽略 BUSYGROUP 错误）。
     */
    public void ensureGroup(String group) {
        if (group == null || group.isBlank()) return;
        if (initializedGroups.contains(group)) return;
        if (groupExists(group)) {
            initializedGroups.add(group);
            return;
        }
        try {
            redis.opsForStream().createGroup(stream, ReadOffset.from("0"), group);
            initializedGroups.add(group);
        } catch (Exception e) {
            if (isBusyGroupError(e)) {
                initializedGroups.add(group);
                return; // 组已存在，忽略
            }
            // Stream 可能不存在，先创建 Stream 再重试
            try {
                Map<String, String> initMap = java.util.Collections.singletonMap("_init", "1");
                org.springframework.data.redis.connection.stream.RecordId id =
                        redis.opsForStream().add(org.springframework.data.redis.connection.stream.StreamRecords
                                .newRecord().ofMap(initMap).withStreamKey(stream));
                if (id != null) {
                    redis.opsForStream().delete(stream, id);
                }
                if (groupExists(group)) {
                    initializedGroups.add(group);
                    return;
                }
                redis.opsForStream().createGroup(stream, ReadOffset.from("0"), group);
                initializedGroups.add(group);
            } catch (Exception retryEx) {
                if (isBusyGroupError(retryEx)) {
                    initializedGroups.add(group);
                    return;
                }
                throw new RuntimeException("ensureGroup failed", retryEx);
            }
        }
    }

    private boolean groupExists(String group) {
        try {
            var groups = redis.opsForStream().groups(stream);
            return groups != null && groups.stream().anyMatch(info -> group.equals(info.groupName()));
        } catch (Exception e) {
            return false;
        }
    }

    private boolean isBusyGroupError(Throwable e) {
        while (e != null) {
            if (e.getMessage() != null && e.getMessage().contains("BUSYGROUP")) {
                return true;
            }
            e = e.getCause();
        }
        return false;
    }

    /**
     * 对应 Go Read：XREADGROUP GROUP {group} {consumer} COUNT {count} BLOCK 1000 STREAMS {stream} >
     */
    public List<QueuedEvent> read(String group, String consumer, int count) {
        if (count <= 0) count = 100;
        try {
            List<MapRecord<String, Object, Object>> records = redis.opsForStream().read(
                    Consumer.from(group, consumer),
                    StreamReadOptions.empty().count(count).block(Duration.ofSeconds(1)),
                    StreamOffset.create(stream, ReadOffset.lastConsumed())
            );
            return toQueuedEvents(records);
        } catch (Exception e) {
            if (e.getMessage() != null && e.getMessage().contains("NOGROUP")) {
                return Collections.emptyList();
            }
            throw new RuntimeException("read stream failed", e);
        }
    }

    /**
     * 对应 Go ClaimStale：XAUTOCLAIM（转交超时未 ACK 的消息）。
     * 注意：StreamOperations.autoClaim 在 Spring Data Redis 3.2.x 中不可用，
     * 使用 XCLAIM + XPENDING 替代。
     */
    public List<QueuedEvent> claimStale(String group, String consumer, Duration minIdle, int count) {
        if (count <= 0) count = 100;
        if (minIdle == null || minIdle.isZero()) minIdle = Duration.ofMinutes(2);
        try {
            // 先获取 pending 消息
            var pendingMessages = redis.opsForStream().pending(
                    stream, group, org.springframework.data.domain.Range.unbounded(), count);
            if (pendingMessages == null || !pendingMessages.iterator().hasNext()) {
                return Collections.emptyList();
            }
            // 筛选超时的消息 ID
            Duration threshold = minIdle;
            List<String> staleIds = new ArrayList<>();
            pendingMessages.forEach(msg -> {
                if (msg.getElapsedTimeSinceLastDelivery().compareTo(threshold) >= 0) {
                    staleIds.add(msg.getIdAsString());
                }
            });
            if (staleIds.isEmpty()) return Collections.emptyList();
            // XCLAIM
            List<MapRecord<String, Object, Object>> claimed = redis.opsForStream().claim(
                    stream, group, consumer, threshold,
                    staleIds.stream().map(RecordId::of).toArray(RecordId[]::new));
            return toQueuedEvents(claimed);
        } catch (Exception e) {
            // 退化到空列表
            return Collections.emptyList();
        }
    }

    /**
     * 对应 Go Ack：XACK。
     */
    public void ack(String group, List<String> ids) {
        if (ids == null || ids.isEmpty()) return;
        String[] idArr = ids.toArray(new String[0]);
        redis.opsForStream().acknowledge(stream, group, idArr);
    }

    public void delete(List<String> ids) {
        if (ids == null || ids.isEmpty()) return;
        redis.opsForStream().delete(stream, ids.toArray(String[]::new));
    }

    /**
     * 对应 Go Trim：XTRIMAPPROX。
     */
    public void trim(long maxLen) {
        if (maxLen <= 0) return;
        redis.opsForStream().trim(stream, maxLen, false);
    }

    private void trimStreamIfNeeded(org.springframework.data.redis.connection.RedisConnection conn) {
        if (streamMaxLen <= 0) {
            return;
        }
        conn.streamCommands().xTrim(stream.getBytes(StandardCharsets.UTF_8), streamMaxLen);
    }

    // ─── 分布式锁支持（供 RedisLock 使用）────────────────────────────────────

    public StringRedisTemplate getRedisTemplate() {
        return redis;
    }

    // ─── 私有工具方法 ─────────────────────────────────────────────────────────

    private String bidKey(String qk) {
        return prefix + "bid:" + qk;
    }

    private String mediaTraceKey(String media, String traceId) {
        return prefix + "media:" + media + ":" + traceId;
    }

    private String conversionSeenKey(String dedupeKey) {
        return prefix + "conv:seen:" + dedupeKey;
    }

    private String deductionCounterKey(String accountId, int act, String trackingVersion, String date) {
        return String.format("%sdeduct:state:tencent:acct:%s:act:%d:ver:%s:date:%s",
                prefix,
                sanitizeCounterPart(accountId, "unknown"),
                act,
                sanitizeCounterPart(trackingVersion, "v1"),
                sanitizeCounterPart(date, "unknown"));
    }

    private static String sanitizeCounterPart(String value, String fallback) {
        if (value == null || value.isBlank()) {
            return fallback;
        }
        return value.replace(':', '_').trim();
    }

    private static long toLong(Object value) {
        if (value instanceof Number number) {
            return number.longValue();
        }
        return Long.parseLong(String.valueOf(value));
    }

    private DefaultRedisScript<List> buildDeductionDecisionScript() {
        DefaultRedisScript<List> script = new DefaultRedisScript<>();
        script.setResultType(List.class);
        script.setScriptText("""
                local key = KEYS[1]
                local rate = tonumber(ARGV[1]) or 0
                local ttl = tonumber(ARGV[2]) or 0
                local total = redis.call('HINCRBY', key, 'total', 1)
                local deducted = tonumber(redis.call('HGET', key, 'deducted') or '0')
                local target = math.floor((total * rate) / 100)
                local should_deduct = 0
                if deducted < target then
                    deducted = redis.call('HINCRBY', key, 'deducted', 1)
                    should_deduct = 1
                end
                if ttl > 0 then
                    redis.call('EXPIRE', key, ttl)
                end
                return { total, deducted, should_deduct }
                """);
        return script;
    }

    public static class DeductionDecision {
        private final long totalSeen;
        private final long totalDeducted;
        private final boolean shouldDeduct;

        public DeductionDecision(long totalSeen, long totalDeducted, boolean shouldDeduct) {
            this.totalSeen = totalSeen;
            this.totalDeducted = totalDeducted;
            this.shouldDeduct = shouldDeduct;
        }

        public long getTotalSeen() {
            return totalSeen;
        }

        public long getTotalDeducted() {
            return totalDeducted;
        }

        public boolean isShouldDeduct() {
            return shouldDeduct;
        }
    }

    private BidRecord parseBidRecord(String payload) {
        try {
            return objectMapper.readValue(payload, BidRecord.class);
        } catch (Exception e) {
            throw new RuntimeException("parse BidRecord failed", e);
        }
    }

    @SuppressWarnings("unchecked")
    private List<QueuedEvent> toQueuedEvents(List<MapRecord<String, Object, Object>> records) {
        if (records == null) return Collections.emptyList();
        List<QueuedEvent> events = new ArrayList<>();
        for (MapRecord<String, Object, Object> rec : records) {
            Map<Object, Object> values = rec.getValue();
            QueuedEvent event = new QueuedEvent();
            event.setId(rec.getId().getValue());
            event.setType(String.valueOf(values.getOrDefault("type", "")));
            event.setPayload(String.valueOf(values.getOrDefault("payload", "")));
            events.add(event);
        }
        return events;
    }
}