import hashlib
import hmac
from typing import List, Dict, Tuple

class TextWatermarkEmbedder:
    """
    基于统计偏移的文本水印嵌入器
    核心思想：利用密钥对词汇表进行动态分区，
    在生成时对"绿名单"token施加概率加成
    """

    def __init__(self, secret_key: bytes, vocab_size: int,
                 green_ratio: float = 0.5, delta: float = 2.0):
        """
        参数：
            secret_key: 水印密钥，用于词汇表分区和检测
            vocab_size: 模型词汇表大小
            green_ratio: 绿名单比例，通常为0.5
            delta: 概率偏移强度，越大越易检测但影响质量
        """
        self.secret_key = secret_key
        self.vocab_size = vocab_size
        self.green_ratio = green_ratio
        self.delta = delta
        self.green_list_size = int(vocab_size * green_ratio)

    def _get_green_list(self, prev_token: int, context_hash: str) -> set:
        """
        根据前一个token和上下文哈希动态生成绿名单
        动态分区使得水印更难被攻击者逆向
        """
        # 使用HMAC确保分区的伪随机性和密钥依赖性
        message = f"{prev_token}:{context_hash}".encode()
        seed = hmac.new(self.secret_key, message, hashlib.sha256).digest()
        
        # 基于种子进行确定性洗牌，选取绿名单
        import random
        rng = random.Random(seed)
        all_tokens = list(range(self.vocab_size))
        rng.shuffle(all_tokens)
        green_list = set(all_tokens[:self.green_list_size])
        return green_list

    def watermark_logits(self, logits: List[float], prev_token: int,
                         context_hash: str) -> List[float]:
        """
        对模型输出的logits施加水印偏移
        这是嵌入水印的核心函数，在推理时被调用
        """
        green_list = self._get_green_list(prev_token, context_hash)
        
        watermarked_logits = logits.copy()
        for token_id in green_list:
            # 对绿名单token的logit施加正向偏移
            watermarked_logits[token_id] += self.delta
        
        return watermarked_logits

    def generate_with_watermark(self, model, prompt: str,
                                 max_tokens: int = 512) -> str:
        """
        带水印的文本生成流程
        """
        tokens = model.tokenize(prompt)
        generated_tokens = []
        
        for step in range(max_tokens):
            # 获取上下文哈希（使用最近N个token作为上下文窗口）
            context_window = tokens[-64:]
            context_hash = hashlib.sha256(
                str(context_window).encode()
            ).hexdigest()
            
            # 模型前向推理，获取原始logits
            logits = model.forward(tokens)
            
            # 施加水印偏移
            prev_token = tokens[-1]
            watermarked = self.watermark_logits(
                logits, prev_token, context_hash
            )
            
            # 从偏移后的logits中采样下一个token
            next_token = model.sample(watermarked)
            generated_tokens.append(next_token)
            tokens.append(next_token)
            
            if next_token == model.eos_token_id:
                break
        
        return model.detokenize(generated_tokens)

import math
from scipy.stats import norm

class TextWatermarkDetector:
    """
    文本水印检测器
    通过统计绿名单token的比例来判断文本是否含水印
    使用假设检验进行统计判定
    """

    def __init__(self, secret_key: bytes, vocab_size: int,
                 green_ratio: float = 0.5,
                 z_threshold: float = 4.0):
        """
        参数：
            z_threshold: Z分数阈值，超过此值判定为含水印
                         4.0对应极低的假阳性率（约1/30000）
        """
        self.secret_key = secret_key
        self.vocab_size = vocab_size
        self.green_ratio = green_ratio
        self.green_list_size = int(vocab_size * green_ratio)
        self.z_threshold = z_threshold

    def _get_green_list(self, prev_token: int, context_hash: str) -> set:
        """与嵌入器相同的绿名单生成逻辑"""
        message = f"{prev_token}:{context_hash}".encode()
        seed = hmac.new(self.secret_key, message, hashlib.sha256).digest()
        import random
        rng = random.Random(seed)
        all_tokens = list(range(self.vocab_size))
        rng.shuffle(all_tokens)
        return set(all_tokens[:self.green_list_size])

    def detect(self, model, text: str) -> Dict:
        """
        检测文本是否包含水印
        返回检测结果和统计信息
        """
        tokens = model.tokenize(text)
        
        if len(tokens) < 20:
            return {
                "watermarked": False,
                "reason": "文本过短，无法可靠检测",
                "z_score": 0.0,
                "green_count": 0,
                "total_count": len(tokens)
            }
        
        green_count = 0
        scored_tokens = 0
        
        for i in range(1, len(tokens)):
            prev_token = tokens[i - 1]
            # 重建上下文哈希
            context_window = tokens[max(0, i-64):i]
            context_hash = hashlib.sha256(
                str(context_window).encode()
            ).hexdigest()
            
            green_list = self._get_green_list(prev_token, context_hash)
            
            if tokens[i] in green_list:
                green_count += 1
            scored_tokens += 1
        
        # 计算Z分数
        # H0: 文本不含水印，绿名单token占比期望为green_ratio
        expected = self.green_ratio
        variance = self.green_ratio * (1 - self.green_ratio)
        z_score = (green_count / scored_tokens - expected) / \
                  math.sqrt(variance / scored_tokens)
        
        is_watermarked = z_score > self.z_threshold
        
        return {
            "watermarked": is_watermarked,
            "z_score": round(z_score, 4),
            "green_count": green_count,
            "total_count": scored_tokens,
            "green_ratio_observed": round(green_count / scored_tokens, 4),
            "confidence": round(
                norm.cdf(z_score) if z_score > 0 else 1 - norm.cdf(z_score), 6
            )
        }

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from typing import Optional, Literal
import aiohttp

app = FastAPI(title="Claude Watermark Detection API", version="1.0")

class DetectionRequest(BaseModel):
    """水印检测请求体"""
    content: str                          # 待检测文本
    content_type: Literal["text", "image_url"] = "text"
    detection_sensitivity: Literal["low", "medium", "high"] = "medium"
    return_details: bool = False          # 是否返回详细统计信息

class DetectionResponse(BaseModel):
    """水印检测响应体"""
    detected: bool                        # 是否检测到水印
    confidence: float                     # 置信度 0-1
    model_source: Optional[str] = None    # 推测的模型来源
    content_length: int                   # 内容长度
    details: Optional[dict] = None        # 详细统计信息

# 敏感度到Z分数阈值的映射
SENSITIVITY_MAP = {
    "low": 6.0,      # 低敏感度：极低假阳性，可能漏检
    "medium": 4.0,   # 中敏感度：平衡假阳性和假阴性
    "high": 2.5,     # 高敏感度：高召回率，假阳性略高
}

@app.post("/v1/detect", response_model=DetectionResponse)
async def detect_watermark(request: DetectionRequest):
    """
    检测文本或图像中是否包含Claude水印
    """
    if request.content_type == "text":
        if len(request.content) < 50:
            raise HTTPException(
                status_code=400,
                detail="文本长度不足50字符，无法可靠检测"
            )
        
        # 初始化检测器，根据敏感度调整阈值
        z_threshold = SENSITIVITY_MAP[request.detection_sensitivity]
        detector = TextWatermarkDetector(
            secret_key=get_api_key(),
            vocab_size=128000,
            z_threshold=z_threshold
        )
        
        result = detector.detect(model, request.content)
        
        response = DetectionResponse(
            detected=result["watermarked"],
            confidence=result["confidence"],
            model_source="Claude" if result["watermarked"] else None,
            content_length=len(request.content),
            details=result if request.return_details else None
        )
        return response
    
    elif request.content_type == "image_url":
        # 图像水印检测逻辑
        image_data = await fetch_image(request.content)
        image_result = detect_image_watermark(image_data)
        return DetectionResponse(
            detected=image_result["detected"],
            confidence=image_result["confidence"],
            model_source="Claude" if image_result["detected"] else None,
            content_length=len(image_data),
            details=image_result if request.return_details else None
        )