#!/usr/bin/env python3
"""
WeMM 视觉-语言融合模块
支持：语义检索、知识记忆、异常检测、上下文理解
"""
import asyncio
import logging
import os
import time
import base64
from typing import Optional, List, Dict, Any
from collections import deque

import numpy as np
import aiohttp

logger = logging.getLogger(__name__)


class WeMMClient:
    """WeMM-Embedding 远程客户端"""
    
    def __init__(self, api_url: str, dimension: int = 2048):
        self.api_url = api_url.rstrip('/')
        self.dimension = dimension
    
    async def encode_text(self, text: str) -> List[float]:
        """编码文本"""
        async with aiohttp.ClientSession() as session:
            async with session.post(
                f"{self.api_url}/embed",
                json={"inputs": text, "dimension": self.dimension}
            ) as resp:
                if resp.status == 200:
                    result = await resp.json()
                    return result["embeddings"][0]
                raise Exception(f"API error: {resp.status}")
    
    async def encode_image(self, img_bytes: bytes) -> List[float]:
        """编码图像"""
        async with aiohttp.ClientSession() as session:
            data = aiohttp.FormData()
            data.add_field('file', img_bytes, filename='frame.jpg', content_type='image/jpeg')
            data.add_field('dimension', str(self.dimension))
            
            async with session.post(f"{self.api_url}/embed", data=data) as resp:
                if resp.status == 200:
                    result = await resp.json()
                    return result["embeddings"][0]
                raise Exception(f"API error: {resp.status}")
    
    async def encode_mixed(self, image_bytes: bytes, text: str) -> tuple:
        """编码图像+文本，返回相似度"""
        async with aiohttp.ClientSession() as session:
            data = aiohttp.FormData()
            data.add_field('file', image_bytes, filename='frame.jpg', content_type='image/jpeg')
            data.add_field('query', text)
            data.add_field('dimension', str(self.dimension))
            
            async with session.post(f"{self.api_url}/embed", data=data) as resp:
                if resp.status == 200:
                    result = await resp.json()
                    embeddings = result.get("embeddings", [])
                    if len(embeddings) >= 2:
                        img_emb = np.array(embeddings[0])
                        text_emb = np.array(embeddings[1])
                        similarity = np.dot(img_emb, text_emb)
                        return similarity, embeddings
                return None, None


class MultimodalMemory:
    """视觉知识记忆库"""
    
    def __init__(self, max_entries: int = 100):
        self.entries: deque = deque(maxlen=max_entries)
    
    def add(self, label: str, embedding: List[float], timestamp: float = None):
        """添加记忆条目"""
        self.entries.append({
            "label": label,
            "embedding": np.array(embedding),
            "timestamp": timestamp or time.time()
        })
    
    def search(self, query_embedding: List[float], top_k: int = 3) -> List[Dict]:
        """相似度搜索"""
        query = np.array(query_embedding)
        results = []
        
        for entry in self.entries:
            similarity = np.dot(query, entry["embedding"])
            results.append({
                "label": entry["label"],
                "similarity": float(similarity),
                "timestamp": entry["timestamp"]
            })
        
        results.sort(key=lambda x: x["similarity"], reverse=True)
        return results[:top_k]
    
    def clear(self):
        """清空记忆"""
        self.entries.clear()


class VisualLanguageFusion:
    """视觉-语言融合引擎"""
    
    def __init__(self, wemm_client: WeMMClient, memory: MultimodalMemory):
        self.wemm = wemm_client
        self.memory = memory
        self.alert_patterns = [
            {"text": "有人跌倒或需要帮助", "threshold": 0.75, "alert": "跌倒检测"},
            {"text": "长时间静止的人", "threshold": 0.70, "alert": "静止检测"},
            {"text": "陌生人出现在画面", "threshold": 0.80, "alert": "陌生人检测"},
            {"text": "紧急求助场景", "threshold": 0.75, "alert": "求助检测"},
        ]
    
    async def learn_scene(self, frame_bytes: bytes, label: str):
        """学习场景（建立视觉记忆）"""
        embedding = await self.wemm.encode_image(frame_bytes)
        self.memory.add(label, embedding)
        logger.info(f"✅ 已学习场景: {label}")
    
    async def answer_video_query(self, frame_bytes: bytes, question: str) -> Dict[str, Any]:
        """回答关于视频的问题"""
        # 编码查询和图像
        query_emb = await self.wemm.encode_text(question)
        img_emb = await self.wemm.encode_image(frame_bytes)
        
        # 计算相似度（简单版本，实际需要更复杂的检索逻辑）
        similarity = np.dot(np.array(query_emb), np.array(img_emb))
        
        # 在记忆中检索
        relevant_scenes = self.memory.search(query_emb, top_k=3)
        
        return {
            "similarity": float(similarity),
            "relevant_scenes": relevant_scenes,
            "question": question
        }
    
    async def detect_anomalies(self, frame_bytes: bytes) -> List[Dict]:
        """检测多种异常模式"""
        img_emb = await self.wemm.encode_image(frame_bytes)
        detections = []
        
        for pattern in self.alert_patterns:
            text_emb = await self.wemm.encode_text(pattern["text"])
            similarity = np.dot(np.array(img_emb), np.array(text_emb))
            
            if similarity > pattern["threshold"]:
                detections.append({
                    "type": pattern["alert"],
                    "similarity": float(similarity),
                    "threshold": pattern["threshold"]
                })
        
        return detections
    
    async def contextual_understanding(self, frame_bytes: bytes, 
                                        conversation_history: str) -> Dict:
        """结合对话历史的上下文理解"""
        # 提取对话中的视觉相关关键词
        visual_keywords = ["看", "看哪里", "检查", "观察", "画面", "视频", "监控"]
        
        # 判断是否需要视觉辅助
        needs_vision = any(kw in conversation_history for kw in visual_keywords)
        
        if needs_vision:
            # 编码当前帧
            img_emb = await self.wemm.encode_image(frame_bytes)
            
            # 结合对话上下文生成回答
            context = {
                "needs_vision": True,
                "embedding": img_emb[:10],  # 只返回前10维作为示例
                "conversation": conversation_history[-200:]
            }
            return context
        
        return {"needs_vision": False}
    
    async def proactive_intervention(self, frame_bytes: bytes, 
                                      user_speech: str) -> Optional[str]:
        """主动介入（结合语音和视觉）"""
        # 检测异常
        anomalies = await self.detect_anomalies(frame_bytes)
        
        # 如果有异常且用户未意识到
        if anomalies and not any(a["type"] in user_speech for a in anomalies):
            alert_type = anomalies[0]["type"]
            similarity = anomalies[0]["similarity"]
            
            return f"⚠️ 检测到{alert_type}（置信度：{similarity:.2f}），需要帮助吗？"
        
        return None


# ========== 使用示例 ==========
async def main():
    """测试 multimodal fusion"""
    
    # 初始化
    wemm = WeMMClient("http://localhost:8765")
    memory = MultimodalMemory()
    fusion = VisualLanguageFusion(wemm, memory)
    
    # 1. 学习场景
    print("\n=== 1. 学习场景 ===")
    # 这里应该有实际的图像数据
    # await fusion.learn_scene(image_bytes, "客厅")
    
    # 2. 异常检测
    print("\n=== 2. 异常检测 ===")
    # 模拟图像数据
    dummy_image = os.urandom(10000)
    # anomalies = await fusion.detect_anomalies(dummy_image)
    # print(f"检测结果: {anomalies}")
    
    # 3. 上下文理解
    print("\n=== 3. 上下文理解 ===")
    conversation = "帮我看看厨房里有没有人"
    # context = await fusion.contextual_understanding(dummy_image, conversation)
    # print(f"上下文: {context}")
    
    print("\n✅ 模块加载成功")


if __name__ == "__main__":
    asyncio.run(main())
