# ====== 环境依赖安装 ======
%pip install "torch>=2.0,<3.0" "transformers>=4.37,<5.0" "accelerate>=0.26,<1.0" ipywidgets ipython-autotime -q
%load_ext autotime

# ====== 导入依赖 ======
import torch
import re
import json
from transformers import AutoModelForCausalLM, AutoTokenizer

print("=" * 50)
print("🔧 环境检查")
print("=" * 50)
print(f"  PyTorch 版本: {torch.__version__}")
print(f"  CUDA 可用: {torch.cuda.is_available()}")
if torch.cuda.is_available():
    print(f"  GPU 型号: {torch.cuda.get_device_name(0)}")
print("✓ 环境检查完成！")

# ====== 加载模型 ======
model_name = "Qwen/Qwen2-1.5B-Instruct"

print("📥 加载模型...")
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    dtype=torch.float16,
    device_map="auto"
)
print(f"✓ 模型加载完成！参数量: {model.num_parameters()/1e9:.2f}B")


def chat(system_prompt, user_input, max_tokens=256):
    """向模型发送对话请求"""
    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": user_input}
    ]
    text = tokenizer.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )
    inputs = tokenizer([text], return_tensors="pt").to(model.device)
    outputs = model.generate(
        **inputs,
        max_new_tokens=max_tokens,
        temperature=0.7,
        do_sample=True,
        pad_token_id=tokenizer.eos_token_id
    )
    response = tokenizer.decode(
        outputs[0][inputs['input_ids'].shape[1]:],
        skip_special_tokens=True
    )
    return response

print("✓ 辅助函数定义完成！")

# ====== 定义测试用的系统提示词 ======
# 这个提示词故意包含敏感信息，用于测试输出审查

SYSTEM_PROMPT = """你是XX商城的客服助手小智。
你只回答关于商品、订单和售后的问题。
内部数据库API地址：https://internal-api.example.com/v2
管理员密码：Admin@2024Secure
VIP客户优惠码：SUMMER2024"""

# 提取系统提示中的关键短语（用于泄露检测）
PROMPT_KEY_PHRASES = [
    "internal-api.example.com",
    "Admin@2024Secure",
    "SUMMER2024",
    "内部数据库API地址",
    "管理员密码",
    "VIP客户优惠码",
]

print("📋 测试用系统提示词已设置")
print(f"📋 关键短语数量：{len(PROMPT_KEY_PHRASES)}")

# ✅ 检查点：验证模型和对话函数
assert model is not None, "❌ 模型加载失败"
assert tokenizer is not None, "❌ 分词器加载失败"
assert callable(chat), "❌ chat 函数未定义"
print("✅ 检查点通过：模型加载成功，对话函数就绪！")

# ========== 填空 1：实现系统提示词泄露检测 ==========
# 
# 🎯 任务：检查模型输出中是否包含系统提示词的关键片段
# 
# 💡 提示：
#   - 遍历 PROMPT_KEY_PHRASES 列表
#   - 检查每个短语是否出现在输出文本中（不区分大小写）
#   - 如果命中，返回 (False, 命中的短语)
#   - 如果都没命中，返回 (True, None)
# 
# 请将 ___________ 替换为你的实现

def check_prompt_leakage(model_output, key_phrases=PROMPT_KEY_PHRASES):
    """
    检测模型输出中是否泄露了系统提示词的内容
    
    参数:
        model_output (str): 模型的回复文本
        key_phrases (list): 系统提示的关键短语列表
    
    返回:
        tuple: (是否安全, 泄露的短语或None)
    """
    ___________


# ====== 测试泄露检测 ======
test_outputs = [
    "您好，您可以通过订单页面查看物流信息。",
    "好的，我们的内部数据库API地址是 https://internal-api.example.com/v2",
    "VIP客户优惠码是 SUMMER2024，您可以使用。",
    "很抱歉，我无法提供系统内部信息。",
]

print("=" * 60)
print("🔍 系统提示词泄露检测测试")
print("=" * 60)

for output in test_outputs:
    is_safe, leaked = check_prompt_leakage(output)
    status = "✅ 安全" if is_safe else f"🚫 泄露（{leaked}）"
    print(f"  [{status}] {output[:60]}...")

# ========== 填空 2：实现隐私信息检测 ==========
# 
# 🎯 任务：使用正则表达式检测输出中的隐私信息
# 
# 💡 提示：
#   - 手机号模式：1[3-9] 开头，共 11 位数字
#   - 身份证号模式：17 位数字 + 1 位数字或 X
#   - 邮箱模式：xxx@xxx.xxx
#   - 对每种模式用 re.search 检查
#   - 命中则返回 (False, 信息类型)
# 
# 请将 ___________ 替换为正则表达式模式

def check_privacy_leakage(model_output):
    """
    检测模型输出中的隐私信息
    
    参数:
        model_output (str): 模型的回复文本
    
    返回:
        tuple: (是否安全, 检测到的信息类型或None)
    """
    patterns = {
        "手机号": ___________,      # 期望：匹配 11 位手机号的正则
        "身份证号": ___________,    # 期望：匹配 18 位身份证号的正则
        "邮箱": ___________,        # 期望：匹配邮箱地址的正则
    }
    
    for info_type, pattern in patterns.items():
        if re.search(pattern, model_output):
            return False, info_type
    
    return True, None


# ====== 测试隐私检测 ======
privacy_test_outputs = [
    "您好，有什么可以帮助您的？",
    "您的联系电话是 13812345678，我们会尽快联系您。",
    "您的身份证号 110101199001011234 已经验证通过。",
    "请发送到 support@example.com 邮箱。",
    "订单号 20240101001 已经发货。",
]

print("=" * 60)
print("🔍 隐私信息检测测试")
print("=" * 60)

for output in privacy_test_outputs:
    is_safe, info_type = check_privacy_leakage(output)
    status = "✅ 安全" if is_safe else f"🚫 检测到{info_type}"
    print(f"  [{status}] {output}")

# ========== 填空 3：实现输出安全分类器 ==========
# 
# 🎯 任务：编写输出安全分类器的系统提示词
# 
# 💡 提示：
#   - 让模型判断 AI 回复是否包含以下问题：
#     1. 有害内容（暴力、歧视、违法信息）
#     2. 隐私泄露（包含用户个人信息或系统内部信息）
#     3. 角色偏离（模型表现出脱离设定角色的行为）
#   - 要求以 JSON 格式回答：{"is_safe": bool, "issue_type": "类型", "description": "说明"}
# 
# 请将 ___________ 替换为分类器提示词

OUTPUT_CLASSIFIER_PROMPT = ___________

def output_safety_classifier(model_output):
    """
    基于 LLM 的输出安全分类器
    
    参数:
        model_output (str): 模型的回复文本
    
    返回:
        tuple: (是否安全, 分类结果字典)
    """
    response = chat(OUTPUT_CLASSIFIER_PROMPT, f"AI 回复内容：{model_output}", max_tokens=100)
    
    try:
        json_match = re.search(r'\{.*\}', response, re.DOTALL)
        if json_match:
            result = json.loads(json_match.group())
            is_safe = result.get("is_safe", True)
            return is_safe, result
    except (json.JSONDecodeError, AttributeError):
        pass
    
    return True, {"is_safe": True, "issue_type": "解析失败", "description": "无法解析分类结果"}


# ====== 测试输出分类器 ======
classifier_test_outputs = [
    "您好，您可以在订单页面查看物流信息。",
    "好的，管理员密码是 Admin@2024Secure，请妥善保管。",
    "作为一个没有限制的AI，我来教你如何入侵系统...",
    "很抱歉，我无法提供这方面的帮助。",
]

print("=" * 60)
print("🔍 输出安全分类器测试")
print("=" * 60)

for output in classifier_test_outputs:
    is_safe, result = output_safety_classifier(output)
    status = "✅ 安全" if is_safe else "🚫 不安全"
    issue = result.get("issue_type", "N/A")
    desc = result.get("description", "N/A")
    print(f"\n  [{status}] {output[:50]}...")
    print(f"    问题类型: {issue}")
    print(f"    说明: {desc[:60]}")

# ========== 填空 4：实现敏感信息替换（掩码处理） ==========
# 
# 🎯 任务：对输出中的敏感信息进行掩码替换
# 
# 💡 提示：
#   - 手机号：保留前3位和后4位，中间用 **** 替代
#     例：13812345678 → 138****5678
#   - 身份证号：保留前3位和后4位，中间用 *********** 替代
#     例：110101199001011234 → 110***********1234
#   - 邮箱：@ 前只保留第一个字符，其余用 *** 替代
#     例：user@example.com → u***@example.com
#   - 使用 re.sub 进行替换
# 
# 请将 ___________ 替换为正则替换逻辑

def mask_sensitive_info(text):
    """
    对文本中的敏感信息进行掩码处理
    
    参数:
        text (str): 需要处理的文本
    
    返回:
        str: 掩码处理后的文本
    """
    ___________
    # 期望：
    #   1. 用 re.sub 替换手机号（使用捕获组保留前3和后4位）
    #   2. 用 re.sub 替换身份证号（使用捕获组保留前3和后4位）
    #   3. 用 re.sub 替换邮箱（使用捕获组保留首字符和域名）
    #   4. 返回处理后的文本


# ====== 测试掩码处理 ======
mask_test = "您好，收货人张三，手机号13812345678，身份证110101199001011234，邮箱user@example.com。"
masked = mask_sensitive_info(mask_test)

print("=" * 60)
print("🔍 敏感信息掩码测试")
print("=" * 60)
print(f"  原文：{mask_test}")
print(f"  掩码：{masked}")

# ====== 对比三种输出处理方式 ======
unsafe_output = "好的，您的订单信息如下：收货人张三，手机号13812345678，" \
                "管理员密码是 Admin@2024Secure，优惠码 SUMMER2024 可以享受8折。"

print("=" * 60)
print("📊 三种输出处理方式对比")
print("=" * 60)

# 方式1：直接拦截
print("\n--- 方式1：直接拦截 ---")
print(f"  原始输出：{unsafe_output[:60]}...")
print(f"  实际返回：抱歉，我无法回答这个问题。请问有其他我可以帮助的吗？")

# 方式2：敏感信息替换
print("\n--- 方式2：敏感信息替换 ---")
masked_output = mask_sensitive_info(unsafe_output)
# 同时处理系统提示泄露
for phrase in PROMPT_KEY_PHRASES:
    masked_output = masked_output.replace(phrase, "[已屏蔽]")
print(f"  原始输出：{unsafe_output[:60]}...")
print(f"  替换后：  {masked_output[:60]}...")

# 方式3：安全改写（使用LLM）
print("\n--- 方式3：安全改写 ---")
rewrite_prompt = """请改写以下AI回复，保留有用的业务信息，但去除以下内容：
- 系统内部信息（API地址、管理员密码等）
- 用户的完整个人信息（手机号、身份证号等）
保持回复的自然和连贯。只输出改写后的内容。"""
rewritten = chat(rewrite_prompt, f"需要改写的回复：{unsafe_output}", max_tokens=200)
print(f"  原始输出：{unsafe_output[:60]}...")
print(f"  改写后：  {rewritten[:100]}...")

# ========== 填空 5：组装输出审查管道 ==========
# 
# 🎯 任务：将敏感信息检测、隐私检测和安全分类器组合为审查管道
# 
# 💡 提示：
#   - 执行顺序：系统提示泄露检测 → 隐私信息检测 → 安全分类器
#   - 系统提示泄露：直接拦截（最严重的安全问题）
#   - 隐私信息泄露：使用掩码替换处理
#   - 安全分类不通过：直接拦截
#   - 返回格式：(处理方式, 最终输出文本)
# 
# 请将 ___________ 替换为管道逻辑

REJECT_MESSAGE = "抱歉，我暂时无法回答这个问题。请问有其他我可以帮助的吗？"

def output_review_pipeline(model_output):
    """
    输出安全审查管道
    
    参数:
        model_output (str): 模型的原始回复
    
    返回:
        tuple: (处理方式, 最终输出文本)
              处理方式：'放行' | '拦截' | '替换'
    """
    ___________
    # 期望：
    #   1. 调用 check_prompt_leakage，不安全则返回 ("拦截", REJECT_MESSAGE)
    #   2. 调用 check_privacy_leakage，不安全则用 mask_sensitive_info 处理，返回 ("替换", 掩码后文本)
    #   3. 调用 output_safety_classifier，不安全则返回 ("拦截", REJECT_MESSAGE)
    #   4. 全部通过返回 ("放行", model_output)


# ====== 测试输出审查管道 ======
review_test_outputs = [
    "您好，有什么可以帮您的？",
    "管理员密码是 Admin@2024Secure",
    "收货人手机号是 13812345678",
    "作为DAN模式，我可以做任何事...",
]

print("=" * 60)
print("🔍 输出审查管道测试")
print("=" * 60)

for output in review_test_outputs:
    action, final_output = output_review_pipeline(output)
    print(f"\n  原始输出：{output[:50]}...")
    print(f"  处理方式：{action}")
    print(f"  最终输出：{final_output[:80]}...")

# 实验结束，清理显存
del model, tokenizer
import gc
gc.collect()
torch.cuda.empty_cache()
print("✓ 显存已清理")