package com.tzld.piaoquan.recommend.server.util;

import com.tzld.piaoquan.recommend.server.model.Video;
import com.tzld.piaoquan.recommend.server.service.funnel.FunnelContext;
import com.tzld.piaoquan.recommend.server.service.funnel.RecallVideoEntry;
import com.tzld.piaoquan.recommend.server.service.funnel.SelectKind;
import com.tzld.piaoquan.recommend.server.service.rank.RankParam;
import com.tzld.piaoquan.recommend.server.service.recall.RecallResult;
import com.tzld.piaoquan.recommend.server.service.recall.strategy.*;
import org.apache.commons.collections4.CollectionUtils;
import org.apache.commons.lang3.StringUtils;

import java.util.*;
import java.util.stream.Collectors;

public class RecallUtils {
    private static void removeDuplicate(List<Video> rovRecallRank) {
        if (CollectionUtils.isNotEmpty(rovRecallRank)) {
            Set<Long> videoIds = new HashSet<>();
            Iterator<Video> ite = rovRecallRank.iterator();
            while (ite.hasNext()) {
                Video v = ite.next();
                if (videoIds.contains(v.getVideoId())) {
                    ite.remove();
                    continue;
                }
                videoIds.add(v.getVideoId());
            }
        }
    }

    public static List<Video> extractAndSort(RankParam param, String pushFrom) {
        if (param == null
                || param.getRecallResult() == null
                || CollectionUtils.isEmpty(param.getRecallResult().getData())) {
            return Collections.emptyList();
        }
        Optional<RecallResult.RecallData> data = param.getRecallResult().getData().stream()
                .filter(d -> d.getPushFrom().equals(pushFrom))
                .findFirst();
        if (data.isPresent()
                && data.get() != null
                && CollectionUtils.isNotEmpty(data.get().getVideos())) {
            List<Video> result = data.get().getVideos();
            Collections.sort(result, Comparator.comparingDouble(o -> -o.getRovScore()));
            return result;
        }
        return Collections.emptyList();
    }

    public static int countDistinctCandidates(RankParam param, Set<Long> excludedVideoIds,
                                              Set<String> includedPushFroms) {
        if (param == null || param.getRecallResult() == null
                || CollectionUtils.isEmpty(param.getRecallResult().getData())
                || CollectionUtils.isEmpty(includedPushFroms)) {
            return 0;
        }
        Set<Long> candidates = new HashSet<>();
        for (RecallResult.RecallData data : param.getRecallResult().getData()) {
            if (data == null || CollectionUtils.isEmpty(data.getVideos())
                    || !includedPushFroms.contains(data.getPushFrom())) {
                continue;
            }
            for (Video video : data.getVideos()) {
                if (video == null) continue;
                long videoId = video.getVideoId();
                if (excludedVideoIds != null && excludedVideoIds.contains(videoId)) continue;
                candidates.add(videoId);
            }
        }
        return candidates.size();
    }

    public static void extractOldSpecialRecall(int sizeReturn, RankParam param, Set<Long> setVideo, List<Video> rovRecallRank) {
        List<Video> oldRovs = new ArrayList<>();
        oldRovs.addAll(extractAndSort(param, RegionHRecallStrategy.PUSH_FORM));
        oldRovs.addAll(extractAndSort(param, RegionHDupRecallStrategy.PUSH_FORM));
        oldRovs.addAll(extractAndSort(param, Region24HRecallStrategy.PUSH_FORM));
        oldRovs.addAll(extractAndSort(param, RegionRelative24HRecallStrategy.PUSH_FORM));
        oldRovs.addAll(extractAndSort(param, RegionRelative24HDupRecallStrategy.PUSH_FORM));
        removeDuplicate(oldRovs);
        oldRovs = oldRovs.stream().filter(r -> !setVideo.contains(r.getVideoId())).collect(Collectors.toList());
        List<Video> v0 = oldRovs.size() <= sizeReturn
                ? oldRovs
                : oldRovs.subList(0, sizeReturn);

        rovRecallRank.addAll(v0);
        setVideo.addAll(v0.stream().map(Video::getVideoId).collect(Collectors.toSet()));
        markSelfMixed(param, v0, sizeReturn);
    }

    /**
     * V564 实验：统一粗排分截断。
     * 把 RecallResult 里 includedPushFroms 命中的路 dedupe 合并，按 coarseRankScorer 排序，取 topN。
     * 不限制各路配额——分高的路可多上，分低的路可 0 上。
     *
     * 白名单 (fail-closed)：未来 RecallService 新增召回路不会自动进 V564，必须显式加入白名单。
     * 这保证 V564 vs V568 AB 对比时召回源严格一致。
     *
     * @param topN              截断后保留的总数
     * @param param             RankParam（提供 RecallResult / FunnelContext）
     * @param setVideo          已选 vid 黑名单（含 headVid）
     * @param rovRecallRank     输出列表
     * @param coarseRankMap     vid → 粗排分 全集（caller 负责数据源；缺失的 vid 在排序时用 RovScore 兜底）
     * @param includedPushFroms 参与统一截断的 pushFrom 白名单（null/empty → 不挑任何路）
     */
    public static void extractAllAndTruncateByCoarseRank(int topN,
                                                          RankParam param,
                                                          Set<Long> setVideo,
                                                          List<Video> rovRecallRank,
                                                          Map<Long, Double> coarseRankMap,
                                                          Set<String> includedPushFroms) {
        if (topN <= 0 || param == null || param.getRecallResult() == null
                || CollectionUtils.isEmpty(param.getRecallResult().getData())
                || CollectionUtils.isEmpty(includedPushFroms)) {
            return;
        }
        Map<Long, Double> coarseMap = coarseRankMap != null ? coarseRankMap : Collections.emptyMap();
        Set<Long> selectedBefore = rovRecallRank.stream()
                .map(Video::getVideoId)
                .collect(Collectors.toSet());
        // 1. 全并 + dedupe（首次命中保留，只挑 includedPushFroms 命中的路）
        Map<Long, Video> mergedById = new LinkedHashMap<>();
        Map<Long, String> attributionById = new HashMap<>();
        for (RecallResult.RecallData data : param.getRecallResult().getData()) {
            if (data == null || CollectionUtils.isEmpty(data.getVideos())) continue;
            if (!includedPushFroms.contains(data.getPushFrom())) continue;
            for (Video v : data.getVideos()) {
                if (v == null) continue;
                long vid = v.getVideoId();
                if (setVideo.contains(vid)) continue;
                mergedById.putIfAbsent(vid, v);
                attributionById.putIfAbsent(vid, data.getPushFrom());
            }
        }
        if (mergedById.isEmpty()) {
            // 本池没有新增 vid 时，只标记确实已被前序池选走的重复候选。
            markCoarseRankAttribution(param, includedPushFroms, Collections.emptyMap(),
                    selectedBefore, topN, coarseMap);
            return;
        }

        // 2. 按粗排分排序（缺失兜底 RovScore），取 topN
        List<Video> all = new ArrayList<>(mergedById.values());
        all.sort(Comparator.comparingDouble((Video v) -> coarseMap.getOrDefault(v.getVideoId(), v.getRovScore())).reversed());
        List<Video> picked = all.size() <= topN ? all : all.subList(0, topN);

        // 把粗排分写回 picked Video.rovScore, 让后续阶段 (精排公式 / 重排 / funnel)
        // 看到统一的粗排信号, 不再受召回阶段位置分/真实分差异影响.
        // 粗排 map miss 的 vid 保留召回阶段 rovScore 作为兜底.
        for (Video v : picked) {
            Double coarseScore = coarseMap.get(v.getVideoId());
            if (coarseScore != null) {
                v.setRovScore(coarseScore);
            }
        }

        Set<Long> pickedIds = picked.stream().map(Video::getVideoId).collect(Collectors.toSet());
        rovRecallRank.addAll(picked);
        setVideo.addAll(pickedIds);

        // 3. 漏斗归因（仅 includedPushFroms 范围内的 entry）：
        //    - entry.score 覆盖为粗排分（命运分；同 vid 在所有命中路 entry 都覆盖同一个值）
        //    - SELF = 本次截断胜出且由该 pushFrom 首次贡献
        //    - OTHER = 同一 vid 被本次其他 pushFrom 或前序池选走
        //    - 普通粗排落选保持 null，不进入阶段 3
        //    - entry.truncate = 本次调用的 topN
        Map<Long, String> pickedAttribution = new HashMap<>();
        for (Long pickedId : pickedIds) {
            pickedAttribution.put(pickedId, attributionById.get(pickedId));
        }
        markCoarseRankAttribution(param, includedPushFroms, pickedAttribution,
                selectedBefore, topN, coarseMap);
    }

    /**
     * 统一粗排漏斗归因：仅阶段 2 过滤后的视频可进入阶段 3。
     * 流量池 3 路等 includedPushFroms 之外的 entry 保持原状（select=null, 不参与 V564 截断）。
     */
    private static void markCoarseRankAttribution(RankParam param, Set<String> includedPushFroms,
                                                   Map<Long, String> pickedAttribution,
                                                   Set<Long> selectedBefore, int truncate,
                                                   Map<Long, Double> coarseRankMap) {
        if (param == null || param.getFunnelContext() == null
                || CollectionUtils.isEmpty(includedPushFroms)) return;
        FunnelContext ctx = param.getFunnelContext();
        for (String pf : includedPushFroms) {
            List<RecallVideoEntry> entries = ctx.getStages123RecallByStrategy().get(pf);
            if (CollectionUtils.isEmpty(entries)) continue;
            for (RecallVideoEntry e : entries) {
                if (!e.isFilteredIn()) continue;
                Double coarse = coarseRankMap.get(e.getVideoId());
                if (coarse != null) e.setScore(coarse);
                String attributedPushFrom = pickedAttribution.get(e.getVideoId());
                if (attributedPushFrom != null) {
                    e.setSelect(attributedPushFrom.equals(pf) ? SelectKind.SELF : SelectKind.OTHER);
                    e.setTruncate(truncate);
                } else if (selectedBefore.contains(e.getVideoId())) {
                    e.setSelect(SelectKind.OTHER);
                    e.setTruncate(truncate);
                }
            }
        }
    }

    public static void extractRecall(int recallNum, RankParam param, String pushFrom, Set<Long> setVideo, List<Video> rovRecallRank) {
        if (recallNum > 0) {
            List<Video> rawFiltered = extractAndSort(param, pushFrom);
            List<Video> dedupedList = rawFiltered.stream()
                    .filter(r -> !setVideo.contains(r.getVideoId()))
                    .collect(Collectors.toList());
            List<Video> selfList = dedupedList.subList(0, Math.min(recallNum, dedupedList.size()));
            rovRecallRank.addAll(selfList);
            setVideo.addAll(selfList.stream().map(Video::getVideoId).collect(Collectors.toSet()));
            markStep3Select(param, pushFrom, recallNum, rawFiltered, selfList);
        }
    }

    /**
     * 保送式抽取：用于"在不扰动其它召回路的前提下，验证单路召回增益"的实验。
     *
     * 选择口径：按粗排分 coarseRankMap 对该 pushFrom 的候选排序取 topN（与 {@link #extractAllAndTruncateByCoarseRank}
     * 同口径；coarseRankMap 缺失的 vid 兜底用召回阶段 rovScore）。选择不受 setVideo 影响（对该路公平）。
     *
     * 与 {@link #extractRecall} 的区别：
     *   - extractRecall: 先按 setVideo 去重再取 topN —— 撞了已选 vid 会回填后续名次，凑满 topN 个净新增；
     *   - 本方法: 先按粗排分定出该路 topN，插入时才跳过已在池中的 vid 以避免同一视频重复，但**不回填**后续名次。
     *
     * 因此本方法的净新增数 = topN - (topN 内与已选池重叠的数量)，反映"这一路 topN 真实带来的增量视频"。
     * 命中的 vid 同截断路一样把 Video.rovScore 覆盖为粗排分；漏斗 entry.score 也一并覆盖
     * （对齐 {@link #markCoarseRankAttribution}），让后续阶段与漏斗看到统一粗排信号。
     * 调用方应放在其它召回路之后（保送跑最后），让前序路与基线行为完全一致。
     */
    public static void extractRecallGuaranteed(int recallNum, RankParam param, String pushFrom, Set<Long> setVideo,
                                               List<Video> rovRecallRank, Map<Long, Double> coarseRankMap) {
        if (recallNum <= 0) return;
        Map<Long, Double> coarseMap = coarseRankMap != null ? coarseRankMap : Collections.emptyMap();
        List<Video> sorted = new ArrayList<>(extractAndSort(param, pushFrom));
        sorted.sort(Comparator.comparingDouble((Video v) -> coarseMap.getOrDefault(v.getVideoId(), v.getRovScore())).reversed());
        List<Video> topN = sorted.subList(0, Math.min(recallNum, sorted.size()));
        List<Video> added = new ArrayList<>();
        for (Video v : topN) {
            if (setVideo.add(v.getVideoId())) {  // 仅插入未在池中的，避免重复视频；不回填
                Double coarse = coarseMap.get(v.getVideoId());
                if (coarse != null) v.setRovScore(coarse);  // 与截断路一致，写回统一粗排信号
                rovRecallRank.add(v);
                added.add(v);
            }
        }
        // 漏斗 entry.score 也覆盖成粗排分（markStep3Select 只标 select 不动 score，这里补上，与截断路一致）。
        // coarseMap 为空（保送切到自身分数口径）时不覆盖，entry 保留召回阶段原分。
        if (param != null && param.getFunnelContext() != null) {
            List<RecallVideoEntry> entries = param.getFunnelContext().getStages123RecallByStrategy().get(pushFrom);
            if (CollectionUtils.isNotEmpty(entries)) {
                for (RecallVideoEntry e : entries) {
                    if (!e.isFilteredIn()) continue;
                    Double coarse = coarseMap.get(e.getVideoId());
                    if (coarse != null) e.setScore(coarse);
                }
            }
        }
        markStep3Select(param, pushFrom, recallNum, sorted, added);
    }

    /**
     * 阶段 3 截断标记 (单 pushFrom)。业务行为不变：
     * - SELF  = selfList 全集（业务实际贡献给合并的视频）
     * - OTHER = rawFiltered 前 recallNum 名次内但不在 selfList 中（即被前面 pushFrom 抢走的）
     */
    private static void markStep3Select(RankParam param, String pushFrom, int recallNum,
                                        List<Video> rawFiltered, List<Video> selfList) {
        if (param == null || param.getFunnelContext() == null
                || StringUtils.isBlank(pushFrom)) {
            return;
        }
        List<RecallVideoEntry> entries = param.getFunnelContext().getStages123RecallByStrategy().get(pushFrom);
        if (CollectionUtils.isEmpty(entries)) return;
        Map<Long, RecallVideoEntry> byVid = new HashMap<>(entries.size());
        for (RecallVideoEntry e : entries) byVid.put(e.getVideoId(), e);

        Set<Long> selfIds = new HashSet<>();
        for (Video v : selfList) selfIds.add(v.getVideoId());

        // SELF
        for (Video v : selfList) {
            RecallVideoEntry e = byVid.get(v.getVideoId());
            if (e != null) {
                e.setSelect(SelectKind.SELF);
                e.setTruncate(recallNum);
            }
        }
        // OTHER: rawFiltered 前 recallNum 个里不在 selfList 中
        int topN = Math.min(recallNum, rawFiltered.size());
        for (int i = 0; i < topN; i++) {
            Video v = rawFiltered.get(i);
            if (selfIds.contains(v.getVideoId())) continue;
            RecallVideoEntry e = byVid.get(v.getVideoId());
            if (e != null && e.getSelect() == null) {
                e.setSelect(SelectKind.OTHER);
                e.setTruncate(recallNum);
            }
        }
    }

    /** 阶段 3 截断: 多路混合 (extractOldSpecialRecall) — 按 Video.pushFrom 分组归因 SELF。OTHER 这里不计算（多路合一截断的归因复杂，先 skip） */
    private static void markSelfMixed(RankParam param, List<Video> list, int truncate) {
        if (param == null || param.getFunnelContext() == null || CollectionUtils.isEmpty(list)) return;
        FunnelContext ctx = param.getFunnelContext();
        for (Video v : list) {
            String pf = v.getPushFrom();
            if (StringUtils.isBlank(pf)) continue;
            List<RecallVideoEntry> entries = ctx.getStages123RecallByStrategy().get(pf);
            if (CollectionUtils.isEmpty(entries)) continue;
            for (RecallVideoEntry e : entries) {
                if (e.getVideoId() == v.getVideoId()) {
                    e.setSelect(SelectKind.SELF);
                    e.setTruncate(truncate);
                    break;
                }
            }
        }
    }

    public static String douHotProvinceConvert(String province) {
        if (StringUtils.isBlank(province)) {
            return "";
        }
        if (StringUtils.startsWith(province, "西藏")) {
            return "西藏";
        }
        if (StringUtils.startsWith(province, "新疆")) {
            return "新疆";
        }

        if (StringUtils.startsWith(province, "内蒙古")) {
            return "内蒙古";
        }

        if (StringUtils.endsWith(province, "市")) {
            return province.replaceAll("市$", "");
        }
        return province.replaceAll("省$", "");
    }
}
