#!/usr/bin/env python3
"""
统一对话引擎 — 整合四层路由 + AreaNetwork推理 + LLM兜底

路由优先级:
  1. qa_cache    — 精确匹配(0ms)
  2. qa_fuzzy    — 模糊匹配(0.79阈值)
  3. brain       — AreaNetwork推理(需训练后有效)
  4. llm         — 外部LLM API(NIM/GLM)
  5. fallback    — vecs75语义联想(离线兜底)
"""
import os
import json
import time
import pickle
import numpy as np
from typing import Dict, Optional, Tuple, List


class ChatEngine:
    """虫群统一对话引擎 — 产品版"""

    def __init__(self, brain=None, api_key: str = None,
                 data_dir: str = None, model_dir: str = None):
        # 目录配置
        self.data_dir = data_dir or os.path.expanduser('~/.swarm/data')
        self.model_dir = model_dir or os.path.expanduser('~/.swarm/models')
        os.makedirs(self.data_dir, exist_ok=True)
        os.makedirs(self.model_dir, exist_ok=True)

        # 核心引擎
        self.brain = brain  # SwarmBrain实例(可选)

        # 加载vecs75语义索引
        self.words = []
        self.vecs75 = None
        self.vecs_n = None
        self.w2i = {}
        self._load_vecs75()

        # QA缓存
        self.qa_cache: Dict[str, str] = {}
        self._qa_vecs: List[Tuple[str, np.ndarray]] = []
        self._load_qa_cache()

        # 对话历史
        self.history: List[Dict] = []
        self.max_history = 20

        # LLM配置
        self.api_key = api_key
        self.api_provider = 'nim'

        # 统计
        self._mode_hits: Dict[str, int] = {}
        self._total_ms = 0
        self._learned = 0

    # ========== 加载 ==========

    def _load_vecs75(self):
        """加载vecs75语义索引"""
        pkl_path = os.path.join(self.model_dir, 'vocab75_index.pkl')
        if not os.path.exists(pkl_path):
            # 尝试旧路径
            pkl_path = os.path.expanduser('~/meta_model/models/vocab75_index.pkl')
        if os.path.exists(pkl_path):
            with open(pkl_path, 'rb') as f:
                data = pickle.load(f)
            self.words = data['words']
            self.vecs75 = data['vecs75']
            norms = np.linalg.norm(self.vecs75, axis=1, keepdims=True)
            norms[norms < 1e-8] = 1
            self.vecs_n = self.vecs75 / norms
            self.w2i = {w: i for i, w in enumerate(self.words)}
            print(f'[Chat] vecs75加载: {len(self.words)}词')
        else:
            print('[Chat] vecs75未找到，离线兜底不可用')

    def _load_qa_cache(self):
        """加载QA缓存"""
        cache_path = os.path.join(self.data_dir, 'qa_cache.json')
        if os.path.exists(cache_path):
            with open(cache_path, 'r') as f:
                self.qa_cache = json.load(f)
            for key in self.qa_cache:
                self._add_qa_vec(key)
            print(f'[Chat] QA缓存: {len(self.qa_cache)}条')

    # ========== 向量工具 ==========

    def _text_to_vec(self, text: str) -> Optional[np.ndarray]:
        """文本→vecs75平均向量"""
        if self.vecs75 is None:
            return None
        chars = [c for c in text if c in self.w2i]
        if not chars:
            return None
        idxs = [self.w2i[c] for c in chars]
        vec = self.vecs75[idxs].mean(axis=0)
        norm = np.linalg.norm(vec)
        return vec / norm if norm > 1e-8 else None

    def _add_qa_vec(self, key: str):
        """为QA key添加向量索引"""
        vec = self._text_to_vec(key)
        if vec is not None:
            self._qa_vecs.append((key, vec))

    # ========== 四层路由 ==========

    def _route_qa_cache(self, text: str) -> Optional[str]:
        """第1层: 精确匹配"""
        if text in self.qa_cache:
            return self.qa_cache[text]
        # 去空格/小写匹配
        normalized = text.strip().lower()
        for k, v in self.qa_cache.items():
            if k.strip().lower() == normalized:
                return v
        return None

    def _route_qa_fuzzy(self, text: str, threshold: float = 0.79) -> Optional[str]:
        """第2层: 模糊匹配"""
        vec = self._text_to_vec(text)
        if vec is None or not self._qa_vecs:
            return None
        best_key, best_sim = None, 0
        for key, qvec in self._qa_vecs:
            sim = float(np.dot(vec, qvec))
            if sim > best_sim:
                best_sim = sim
                best_key = key
        if best_sim >= threshold and best_key in self.qa_cache:
            return self.qa_cache[best_key]
        return None

    def _route_brain(self, text: str) -> Optional[Dict]:
        """第3层: AreaNetwork推理"""
        if self.brain is None:
            return None
        try:
            result = self.brain.chat(text)
            if result and result.get('confidence', 0) > 0.3:
                return result  # 返回完整结果(含decoded_words)
        except Exception as e:
            print(f'[Chat] brain路由异常: {e}')
        return None

    def _route_llm(self, text: str) -> Optional[str]:
        """第4层: LLM API"""
        if not self.api_key:
            return None
        try:
            return self._call_nim(text)
        except Exception as e:
            print(f'[Chat] LLM路由异常: {e}')
        return None

    def _fallback(self, text: str) -> str:
        """离线兜底: vecs75语义联想"""
        if self.vecs75 is None:
            return f'[未知] {text}'
        vec = self._text_to_vec(text)
        if vec is None:
            return f'[未知] {text}'
        sims = self.vecs_n @ vec
        top3 = np.argsort(sims)[-3:][::-1]
        words = [self.words[i] for i in top3 if sims[i] > 0.5]
        return ' '.join(words) if words else f'[联想] {text}'

    # ========== LLM调用 ==========

    def _call_nim(self, text: str) -> Optional[str]:
        """调用NIM API"""
        import urllib.request
        url = 'https://integrate.api.nvidia.com/v1/chat/completions'
        headers = {
            'Authorization': f'Bearer {self.api_key}',
            'Content-Type': 'application/json',
        }
        body = json.dumps({
            'model': 'meta/llama-3.1-8b-instruct',
            'messages': [{'role': 'user', 'content': text}],
            'max_tokens': 256, 'temperature': 0.7,
        }).encode()
        req = urllib.request.Request(url, data=body, headers=headers)
        with urllib.request.urlopen(req, timeout=15) as resp:
            data = json.loads(resp.read())
        reply = data['choices'][0]['message']['content']
        # 回存学习
        self._learn_from_llm(text, reply)
        return reply

    # ========== 学习 ==========

    def _learn_from_llm(self, question: str, answer: str):
        """LLM回答回存到QA缓存"""
        self.qa_cache[question] = answer
        self._add_qa_vec(question)
        self._learned += 1
        self._save_qa_cache()

    def teach(self, question: str, answer: str):
        """教学: 手动添加QA对"""
        self.qa_cache[question] = answer
        self._add_qa_vec(question)
        self._learned += 1
        self._save_qa_cache()

    def _save_qa_cache(self):
        """持久化QA缓存"""
        cache_path = os.path.join(self.data_dir, 'qa_cache.json')
        with open(cache_path, 'w') as f:
            json.dump(self.qa_cache, f, ensure_ascii=False, indent=2)

    # ========== 主入口 ==========

    def chat(self, text: str) -> Dict:
        """
        统一对话入口
        
        Returns:
            {'text': 回复, 'mode': 路由层, 'ms': 耗时,
             'confidence': 置信度, 'complexity': 复杂度}
        """
        t0 = time.time()
        
        # 四层路由
        reply = None
        mode = 'fallback'
        confidence = 0.0
        decoded_words = []

        # 第1层: 精确匹配
        reply = self._route_qa_cache(text)
        if reply:
            mode, confidence = 'qa_cache', 1.0
        else:
            # 第2层: 模糊匹配
            reply = self._route_qa_fuzzy(text)
            if reply:
                mode, confidence = 'qa_fuzzy', 0.85
            else:
                # 第3层: brain推理(返回Dict)
                brain_result = self._route_brain(text)
                if brain_result:
                    reply = brain_result.get('text', '')
                    decoded_words = brain_result.get('decoded_words', [])
                    mode, confidence = 'brain', brain_result.get('confidence', 0.6)
                else:
                    # 第4层: LLM
                    reply = self._route_llm(text)
                    if reply:
                        mode, confidence = 'llm', 0.9

        # 兜底
        if not reply:
            reply = self._fallback(text)
            mode, confidence = 'fallback', 0.2

        ms = int((time.time() - t0) * 1000)
        self._mode_hits[mode] = self._mode_hits.get(mode, 0) + 1
        self._total_ms += ms

        # 记录历史
        self.history.append({
            'input': text, 'reply': reply, 'mode': mode,
            'ms': ms, 'confidence': confidence,
        })
        if len(self.history) > self.max_history:
            self.history = self.history[-self.max_history:]

        return {
            'text': reply, 'mode': mode, 'ms': ms,
            'confidence': confidence,
            'decoded_words': decoded_words,
        }

    def stats(self) -> Dict:
        """路由统计"""
        total = sum(self._mode_hits.values())
        return {
            'total_queries': total,
            'mode_hits': dict(self._mode_hits),
            'avg_ms': self._total_ms / max(total, 1),
            'learned': self._learned,
            'qa_cache_size': len(self.qa_cache),
            'history_len': len(self.history),
        }