import json, hashlib, os, re, time, logging
from pathlib import Path
import httpx

logger = logging.getLogger(__name__)


class LLMClient:
    def __init__(self, model=None, timeout=15, cache_dir=".cache/llm"):
        # 模型解析链: 显式实参 → LLM_MODEL 环境变量 → 项目默认（与 config 层一致）
        self.model = model or os.environ.get("LLM_MODEL", "deepseek-v4-flash")
        self.timeout = timeout
        self.dir = Path(cache_dir)
        self.dir.mkdir(parents=True, exist_ok=True)
        self._last_call_time = 0
        self._min_interval = 0.1  # 100ms 最小调用间隔

    def _key(self, msgs):
        return hashlib.sha256(json.dumps(msgs, sort_keys=True).encode()).hexdigest()

    def _get(self, k):
        p = self.dir / f"{k}.json"
        if not p.exists():
            return None
        try:
            return json.loads(p.read_text())["response"]
        except (json.JSONDecodeError, KeyError):
            return None

    def _set(self, k, v):
        (self.dir / f"{k}.json").write_text(json.dumps({"response": v}))

    def _sanitize_input(self, text: str) -> str:
        """输入过滤：移除潜在的prompt注入模式，限制输入长度"""
        if not isinstance(text, str):
            text = str(text)
        
        # 1. 限制输入长度（防止过长输入）
        max_length = 50000
        if len(text) > max_length:
            logger.warning(f"Input truncated from {len(text)} to {max_length} chars")
            text = text[:max_length]
        
        # 2. 移除潜在的prompt注入模式
        injection_patterns = [
            r'(?i)ignore\s+previous\s+instructions',
            r'(?i)disregard\s+all\s+prior',
            r'(?i)you\s+are\s+now\s+',
            r'(?i)system\s*:\s*',
            r'(?i)assistant\s*:\s*',
            r'(?i)\[INST\]',
            r'(?i)\[/INST\]',
            r'(?i)<\|im_start\|>',
            r'(?i)<\|im_end\|>',
        ]
        
        for pattern in injection_patterns:
            if re.search(pattern, text):
                logger.warning(f"Potential prompt injection detected: {pattern}")
                text = re.sub(pattern, '[FILTERED]', text)
        
        return text

    def _validate_output(self, output: str) -> bool:
        """输出验证：检查LLM返回内容的基本有效性"""
        if not output or not isinstance(output, str):
            return False
        
        # 检查长度是否合理
        if len(output) > 100000:
            logger.warning(f"Output too long: {len(output)} chars")
            return False
        
        return True

    def _rate_limit(self):
        """速率限制：确保最小调用间隔"""
        current_time = time.time()
        elapsed = current_time - self._last_call_time
        if elapsed < self._min_interval:
            time.sleep(self._min_interval - elapsed)
        self._last_call_time = time.time()

    def call(self, messages, retries=1):
        k = self._key(messages)
        c = self._get(k)
        if c:
            logger.debug(f"Cache hit for key: {k[:8]}...")
            return c
        
        # 对用户消息进行输入过滤
        sanitized_messages = []
        for msg in messages:
            if msg.get("role") == "user":
                sanitized_content = self._sanitize_input(msg.get("content", ""))
                sanitized_messages.append({**msg, "content": sanitized_content})
            else:
                sanitized_messages.append(msg)
        
        key = os.environ.get("LLM_API_KEY", os.environ.get("OPENAI_API_KEY", ""))
        base = os.environ.get("LLM_API_BASE", "https://api.openai.com/v1")
        
        for a in range(retries + 1):
            try:
                # 速率限制
                self._rate_limit()
                
                logger.info(f"LLM call: model={self.model}, attempt={a+1}/{retries+1}")
                r = httpx.post(f"{base}/chat/completions", json={"model": self.model, "messages": sanitized_messages},
                               headers={"Authorization": f"Bearer {key}"}, timeout=self.timeout)
                r.raise_for_status()
                v = r.json()["choices"][0]["message"]["content"]
                
                # 输出验证
                if not self._validate_output(v):
                    logger.warning("Invalid output from LLM, using fallback")
                    v = ""
                
                self._set(k, v)
                logger.info(f"LLM call success: response length={len(v)}")
                return v
            except Exception as e:
                logger.error(f"LLM call failed: {e}")
                if a == retries:
                    raise
                time.sleep(1)  # 失败后等待1秒重试
        
        return ""