"""
全自动清洗大模型推理污染 —— 三层流水线
  底层 TokenCleaner     ：无效字符 / 冗余短句清洗，保证进入历史的每段文本本身干净
  中层 SessionPurifier  ：四维度脏度打分，把灌水、跑题、违规、模型兜底句挡在上下文外
  上层 ContextReducer   ：话题聚类保留当前话题 + 按时间排序 + token 截断，喂给模型前最后瘦身
入口：LLMDirtyDataCleanPipeline.run(raw_chat_history, latest_user_msg)
"""
import re
import tiktoken
from modelscope import snapshot_download
from sentence_transformers import SentenceTransformer
from sklearn.metrics.pairwise import cosine_similarity
from sklearn.cluster import KMeans
from dataclasses import dataclass
from typing import List, Dict
# ===================== 配置项 =====================
MAX_CONTEXT_TOKEN = 2048       # 上下文窗口上限
DIRTY_SCORE_THRESHOLD = 60     # 脏会话阈值，低于60判定脏会话（逐轮打分，有一轮脏就丢一轮）
TOPIC_SIM_THRESHOLD = 0.2      # 与最新提问的语义相关度低于此值 → 视为跑题扣分
ENCODER = tiktoken.get_encoding("cl100k_base")
cache_dir = "D:\\modelscope\\hub"
embedding_model_dir = snapshot_download(
    model_id="sentence-transformers/all-MiniLM-L6-v2",
    cache_dir=cache_dir,
    revision="master"
)
EMBED_MODEL = SentenceTransformer(embedding_model_dir)
# 会话数据结构
@dataclass
class ChatTurn:
    role: str       # user / assistant
    content: str
    timestamp: float
    dirty_score: float = 100.0
    embedding: List[float] = None
# ===================== 底层：无效Token清洗层 =====================
class TokenCleaner:
    # 清理空白、重复换行、乱码、无效符号
    PATTERN_INVALID_CHAR = re.compile(r"(\s{2,})|([^\u4e00-\u9fa5a-zA-Z0-9\.,;:!?，。；：？！\n])")
    PATTERN_MULTI_NEWLINE = re.compile(r"\n{2,}")
    @staticmethod
    def clean_text(raw_text: str) -> str:
        # 1. 清理特殊无效字符
        text = TokenCleaner.PATTERN_INVALID_CHAR.sub("", raw_text)
        # 2. 合并多个换行
        text = TokenCleaner.PATTERN_MULTI_NEWLINE.sub("\n", text)
        return text.strip()
    @staticmethod
    def count_token(text: str) -> int:
        tokens = ENCODER.encode(text)
        return len(tokens)
    @staticmethod
    def remove_redundant_tokens(raw_text: str) -> str:
        """删除重复冗余短句（简单去重，可替换LLM摘要）"""
        lines = [line.strip() for line in raw_text.split("\n") if len(line.strip())>3]
        unique_lines = []
        seen = set()
        for line in lines:
            if line not in seen:
                seen.add(line)
                unique_lines.append(line)
        return "\n".join(unique_lines)
# ===================== 中层：历史脏会话净化层 =====================
class SessionPurifier:
    # 模型兜底/拒答类句子：内容上没错误，但对后续推理毫无价值，属于"隐性污染"
    FALLBACK_MARKS = ["我不知道", "无法回答", "作为AI", "抱歉，我", "对不起，我不能"]
    # 违规词表（业务可按需扩充，这里只做演示）
    BAD_WORDS = ["乱码测试", "破解", "非法"]
    @staticmethod
    def calc_dirty_score(turn: ChatTurn, latest_query: str) -> float:
        """四维度脏度打分：有效性 / 相关性 / 合规性 / 信息量，满分100，低分即脏"""
        score = 100.0
        content = turn.content
        # ① 有效性扣分：内容太短无业务价值（"好的""嗯"这类）
        if len(content) < 10:
            score -= 30
        # ② 相关性扣分：与最新用户提问语义对立/无关（历史话题已切换，旧话题硬留着只会带偏模型）
        emb_latest = EMBED_MODEL.encode(latest_query)
        emb_turn = EMBED_MODEL.encode(content)
        sim = cosine_similarity([emb_latest], [emb_turn])[0][0]
        if sim < TOPIC_SIM_THRESHOLD:
            score -= 25
        # ③ 合规性扣分：违规关键词，一票重扣
        for w in SessionPurifier.BAD_WORDS:
            if w in content:
                score -= 40
        # ④ 信息量扣分：模型兜底句 / 高度重复（连续多轮"我不知道"会污染后续回复风格）
        if any(m in content for m in SessionPurifier.FALLBACK_MARKS):
            score -= 35
        return max(0, score)
    @staticmethod
    def filter_dirty_chat(chat_list: List[ChatTurn], latest_query: str) -> List[ChatTurn]:
        """逐轮打分过滤：>= 阈值保留。只丢"脏轮"，脏轮前后的好轮不受牵连"""
        valid_turns = []
        for turn in chat_list:
            turn.dirty_score = SessionPurifier.calc_dirty_score(turn, latest_query)
            if turn.dirty_score >= DIRTY_SCORE_THRESHOLD:
                valid_turns.append(turn)
        return valid_turns
# ===================== 上层：上下文冗余智能精简层 =====================
class ContextReducer:
    @staticmethod
    def topic_cluster(chat_list: List[ChatTurn]):
        """话题聚类：KMeans 把历史按语义分组，只保留与"最后一轮"同簇的当前话题。
        中间话题被切换掉后，留在上下文里只会稀释模型对当前任务的注意力。"""
        if len(chat_list) <= 1:
            return chat_list
        embeds = [EMBED_MODEL.encode(t.content) for t in chat_list]
        kmeans = KMeans(n_clusters=min(3, len(chat_list)), random_state=42)
        labels = kmeans.fit_predict(embeds)
        latest_label = labels[-1]
        current_topic_turns = [chat_list[i] for i, lab in enumerate(labels) if lab == latest_label]
        return current_topic_turns
    @staticmethod
    def weight_sort(turns: List[ChatTurn]) -> List[ChatTurn]:
        """按时间权重排序（旧→新）：多轮对话讲究时序，乱序喂模型会答非所问"""
        return sorted(turns, key=lambda x: x.timestamp, reverse=False)
    @staticmethod
    def shrink_context(turns: List[ChatTurn], max_token: int):
        """token 截断：按顺序累加，超预算立即停（后面更晚的轮次宁可整段不要，也不截半句产生残缺上下文）"""
        final_ctx = []
        total_tok = 0
        for t in turns:
            cnt = TokenCleaner.count_token(t.content)
            if total_tok + cnt <= max_token:
                final_ctx.append(t)
                total_tok += cnt
            else:
                break
        return final_ctx
# ===================== 全链路清洗入口 =====================
class LLMDirtyDataCleanPipeline:
    """串起三层：文本清洗 → 脏轮过滤 → 话题聚焦 + token 瘦身。
    输出可直接拼进 prompt 的多轮 messages，附带各层过滤统计方便观测。"""
    def run(self, raw_chat_history: List[ChatTurn], latest_user_msg: str) -> List[ChatTurn]:
        # Step1 Token清洗：逐轮清洗无效字符、去掉重复冗余短句
        cleaned_history = []
        for turn in raw_chat_history:
            clean_content = TokenCleaner.clean_text(turn.content)
            clean_content = TokenCleaner.remove_redundant_tokens(clean_content)
            turn.content = clean_content
            cleaned_history.append(turn)
        # Step2 脏会话过滤：灌水 / 跑题 / 违规 / 兜底句在这里被丢弃
        valid_history = SessionPurifier.filter_dirty_chat(cleaned_history, latest_user_msg)
        # Step3 话题聚类：只保留与最新一轮同话题的历史
        topic_history = ContextReducer.topic_cluster(valid_history)
        # Step4 按时间排序 + Token截断，控制最终上下文不超预算
        sorted_history = ContextReducer.weight_sort(topic_history)
        final_history = ContextReducer.shrink_context(sorted_history, MAX_CONTEXT_TOKEN)
        print(f"[pipeline] 原始 {len(raw_chat_history)} 轮 → 清洗 {len(cleaned_history)} 轮 "
              f"→ 过脏后 {len(valid_history)} 轮 → 聚类后 {len(topic_history)} 轮 "
              f"→ 最终 {len(final_history)} 轮")
        return final_history
# ============ 测试Demo：把 4 步逐层拆开打印 ============
if __name__ == "__main__":
    import time
    from sklearn.cluster import KMeans
    def banner(s):
        print("\n" + "=" * 64 + "\n  " + s + "\n" + "=" * 64)
    latest_input = "现在帮我写一份产品使用手册"
    test_chat = [
        ChatTurn(role="user",      content="  帮我写一个营销文案\n\n\n",    timestamp=time.time()-120),
        ChatTurn(role="assistant", content="下面给你文案。下面给你文案。", timestamp=time.time()-110),
        ChatTurn(role="user",      content="乱码测试，随便灌水闲聊",       timestamp=time.time()-60),
        ChatTurn(role="assistant", content="我不知道你在说什么。",          timestamp=time.time()-50),  # 兜底句
        ChatTurn(role="user",      content="现在帮我写一份产品使用手册",   timestamp=time.time()),
    ]
    # ───── Step 1 ─────
    banner("Step 1 · 文本清洗（去重/去无效字符）")
    cleaned = []
    for t in test_chat:
        before = t.content
        t.content = TokenCleaner.clean_text(t.content)
        t.content = TokenCleaner.remove_redundant_tokens(t.content)
        cleaned.append(t)
        print(f"  [{t.role:<9}]  {before!r:34s} →  {t.content!r}")
    # ───── Step 2 ─────
    banner("Step 2 · 四维度脏度打分（满 100，< 60 丢弃）")
    scored = []
    for t in cleaned:
        score, why = 100.0, []
        # ① 有效性：长度
        if len(t.content) < 10:
            score -= 30; why.append("①短 -30")
        # ② 相关性
        sim = float(cosine_similarity(
            [EMBED_MODEL.encode(latest_input)],
            [EMBED_MODEL.encode(t.content)]
        )[0][0])
        if sim < TOPIC_SIM_THRESHOLD:
            score -= 25; why.append(f"②跑题 -25(sim={sim:.2f})")
        else:
            why.append(f"②相关 sim={sim:.2f}")
        # ③ 合规
        hit_bad = [w for w in SessionPurifier.BAD_WORDS if w in t.content]
        if hit_bad:
            score -= 40; why.append(f"③违规 -40({hit_bad[0]})")
        # ④ 兜底句
        hit_fb = [m for m in SessionPurifier.FALLBACK_MARKS if m in t.content]
        if hit_fb:
            score -= 35; why.append(f"④兜底 -35({hit_fb[0]})")
        t.dirty_score = max(0, score)
        keep = t.dirty_score >= DIRTY_SCORE_THRESHOLD
        scored.append((t, keep))
        flag = "✅ 留" if keep else "❌ 丢"
        print(f"  [{t.role:<9}] score={t.dirty_score:>5.0f} {flag}  {t.content!r:34s}  [{' | '.join(why)}]")
    valid = [t for t, k in scored if k]
    print(f"\n  → 过脏后剩 {len(valid)} 轮 / 丢 {len(scored)-len(valid)} 轮")
    # ───── Step 3 ─────
    banner("Step 3 · 话题聚类（只留与最后一轮同话题）")
    if len(valid) > 1:
        embeds = [EMBED_MODEL.encode(t.content) for t in valid]
        labels = KMeans(n_clusters=min(3, len(valid)), random_state=42, n_init=10).fit_predict(embeds)
        for t, lab in zip(valid, labels):
            print(f"  [{t.role:<9}] cluster={lab}  {t.content!r}")
        current_topic = [valid[i] for i, lab in enumerate(labels) if lab == labels[-1]]
    else:
        labels = [0] * len(valid)
        current_topic = valid
        for t in valid:
            print(f"  [{t.role:<9}] cluster=0  {t.content!r}")
    print(f"\n  → 最后一轮 label={labels[-1]}，同簇保留 {len(current_topic)} 轮")
    # ───── Step 4 ─────
    banner("Step 4 · 时间排序 + token 截断（不超 2048）")
    sorted_h = sorted(current_topic, key=lambda x: x.timestamp)
    total, final = 0, []
    for t in sorted_h:
        cnt = TokenCleaner.count_token(t.content)
        if total + cnt <= MAX_CONTEXT_TOKEN:
            final.append(t); total += cnt
            print(f"  ✅ [{t.role:<9}] +{cnt:>3}tok 累计{total:>4}/{MAX_CONTEXT_TOKEN}  {t.content!r}")
        else:
            print(f"  ❌ [{t.role:<9}] +{cnt:>3}tok 超预算停")
            break
    # ───── 汇总 ─────
    banner("送进 prompt 的最终上下文")
    for t in final:
        print(f"  [{t.role:<9}] score={t.dirty_score:.0f}  {t.content}")
    print(f"\n  漏斗：原始 {len(test_chat)} → 清洗 {len(cleaned)} → 过脏 {len(valid)} "
          f"→ 聚类 {len(current_topic)} → 最终 {len(final)} 轮")
    print(f"  最终 token：{sum(TokenCleaner.count_token(t.content) for t in final)} / {MAX_CONTEXT_TOKEN}")