#!/usr/bin/env python3
"""对话记录器 - 说话人识别 + 面部识别 + 持久化存储"""
import asyncio
import json
import time
import hashlib
import sqlite3
import os
from pathlib import Path
from typing import Optional, List, Dict, Any
import logging

logger = logging.getLogger(__name__)


class DialogueLogger:
    """
    对话记录器：记录谁在什么时候说了什么，关联视觉场景。
    
    功能：
    1. 说话人识别 - 通过 LiveKit participant identity
    2. 面部快照 - 关键帧时记录人脸特征（可选）
    3. JSONL 持久化 - 每次对话追加写入文件
    4. 时间对齐 - 语音和视觉事件关联
    5. 分层存储 - 热数据内存 + 冷数据 SQLite
    6. FTS5 全文搜索
    7. 消息压缩与归档
    """
    
    def __init__(self, db_path: str = None):
        self.db_path = db_path or str(Path.home() / ".hermes" / "workspace" / "livekit-agents" / "dialogue_history.db")
        self._conn = None
        self._current_speaker = None
        self._speech_start = None
        self._conversation_buffer = []
        
        # 分层存储配置
        self.hot_limit = 500      # 热数据内存上限
        self.cold_limit = 10000   # 冷数据 SQLite 上限
        self.compact_threshold = 8000  # token 压缩阈值
        self.archive_days = 7     # 归档天数
        self._hot_dialogues = []  # 内存热数据
        self._hot_speakers = {}   # 内存说话人统计
        
    def initialize(self):
        """初始化 SQLite 数据库并加载热数据"""
        self._conn = sqlite3.connect(self.db_path)
        c = self._conn.cursor()
        
        # 对话记录表
        c.execute('''
            CREATE TABLE IF NOT EXISTS dialogues (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                timestamp REAL NOT NULL,
                speaker_id TEXT,
                participant_identity TEXT,
                text TEXT,
                sentiment TEXT,
                scene_hash TEXT,
                embedding_id INTEGER,
                metadata TEXT,
                token_count INTEGER DEFAULT 0,
                archived INTEGER DEFAULT 0,
                compacted INTEGER DEFAULT 0
            )
        ''')
        
        # FTS5 全文搜索索引 - 创建在外部分离的表
        c.execute('''
            CREATE VIRTUAL TABLE IF NOT EXISTS dialogues_fts USING fts5(
                text, participant_identity
            )
        ''')
        
        # 视觉记忆表（与 VisualMonitor 共享）
        c.execute('''
            CREATE TABLE IF NOT EXISTS visual_memories (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                timestamp REAL NOT NULL,
                embedding BLOB,
                frame_hash TEXT,
                scene_description TEXT,
                face_snapshots TEXT
            )
        ''')
        
        # 说话人索引
        c.execute('''
            CREATE TABLE IF NOT EXISTS speakers (
                identity TEXT PRIMARY KEY,
                first_seen REAL,
                last_seen REAL,
                total_turns INTEGER DEFAULT 0,
                face_embedding BLOB
            )
        ''')
        
        # 创建 FTS5 触发器
        c.execute('''
            CREATE TRIGGER IF NOT EXISTS dialogues_ai AFTER INSERT ON dialogues BEGIN
                INSERT INTO dialogues_fts(rowid, text, participant_identity)
                VALUES (new.id, new.text, new.participant_identity);
            END
        ''')
        
        c.execute('''
            CREATE TRIGGER IF NOT EXISTS dialogues_ad AFTER DELETE ON dialogues BEGIN
                INSERT INTO dialogues_fts(dialogues_fts, rowid)
                VALUES ('delete', old.id);
            END
        ''')
        
        c.execute('''
            CREATE TRIGGER IF NOT EXISTS dialogues_au AFTER UPDATE ON dialogues BEGIN
                INSERT INTO dialogues_fts(dialogues_fts, rowid)
                VALUES ('delete', old.id);
                INSERT INTO dialogues_fts(rowid, text, participant_identity)
                VALUES (new.id, new.text, new.participant_identity);
            END
        ''')
        
        self._conn.commit()
        
        # 加载热数据（最近 hot_limit 条）
        self._load_hot_data()
        
        # 初始化 FTS5 索引
        self._init_fts_index()
        
        logger.info(f"📝 对话数据库: {self.db_path} (热数据: {self.hot_limit}, 冷数据上限: {self.cold_limit})")
        
    def record_speech(self, speaker_id: str, participant_identity: str, text: str, 
                      scene_hash: str = None, embedding_id: int = None) -> int:
        """记录一次对话发言，返回记录 ID"""
        if not self._conn:
            self.initialize()
            
        timestamp = time.time()
        token_count = len(text) // 2  # 简单估算
        
        # 插入对话记录
        cursor = self._conn.execute('''
            INSERT INTO dialogues (timestamp, speaker_id, participant_identity, text, 
                                  scene_hash, embedding_id, token_count, metadata)
            VALUES (?, ?, ?, ?, ?, ?, ?, ?)
        ''', (timestamp, speaker_id, participant_identity, text, 
              scene_hash, embedding_id, token_count, 
              json.dumps({'scene': scene_hash} if scene_hash else None)))
        
        record_id = cursor.lastrowid
        
        # 更新说话人统计
        self._conn.execute('''
            INSERT INTO speakers (identity, first_seen, last_seen, total_turns)
            VALUES (?, ?, ?, 1)
            ON CONFLICT(identity) DO UPDATE SET
                last_seen = excluded.last_seen,
                total_turns = speakers.total_turns + 1
        ''', (participant_identity, timestamp, timestamp))
        
        self._conn.commit()
        logger.info(f"📝 记录对话: {participant_identity} - {text[:30]}... (ID: {record_id})")
        
        # 同步 FTS5 索引
        self._sync_fts_index()
        
        # 写入热数据内存
        self._add_to_hot_data(timestamp, speaker_id, participant_identity, text, 
                             scene_hash, embedding_id, record_id, token_count)
        
        # 检查是否需要压缩
        self._check_compaction()
        
        return record_id
        
    def _init_fts_index(self):
        """初始化 FTS5 索引"""
        if not self._conn:
            return
        try:
            # 清空并重建索引
            self._conn.execute('DELETE FROM dialogues_fts')
            self._conn.commit()
            # 重新插入所有现有数据
            c = self._conn.cursor()
            c.execute('SELECT id, text, participant_identity FROM dialogues WHERE compacted = 0')
            for row in c.fetchall():
                self._conn.execute(
                    'INSERT INTO dialogues_fts(rowid, text, participant_identity) VALUES (?, ?, ?)',
                    (row[0], row[1], row[2])
                )
            self._conn.commit()
            logger.info(f"🔍 FTS5 全文索引已初始化 ({self._conn.execute('SELECT COUNT(*) FROM dialogues_fts').fetchone()[0]} 条)")
        except Exception as e:
            logger.warning(f"⚠️ FTS5 索引初始化失败: {e}")
    
    def _sync_fts_index(self):
        """同步 FTS5 索引"""
        if not self._conn:
            return
        try:
            # 重建索引（简单但可靠）
            self._conn.execute('DELETE FROM dialogues_fts')
            self._conn.commit()
            
            c = self._conn.cursor()
            c.execute('SELECT id, text, participant_identity FROM dialogues WHERE compacted = 0')
            count = 0
            for row in c.fetchall():
                self._conn.execute(
                    'INSERT INTO dialogues_fts(rowid, text, participant_identity) VALUES (?, ?, ?)',
                    (row[0], row[1], row[2])
                )
                count += 1
            self._conn.commit()
            logger.debug(f"🔍 FTS5 索引已同步 ({count} 条)")
        except Exception as e:
            logger.warning(f"⚠️ FTS5 索引同步失败: {e}")
    
    def record_face_snapshot(self, frame_hash: str, face_embedding: List[float], 
                             scene_description: str = None):
        """记录面部快照"""
        if not self._conn:
            self.initialize()
            
        timestamp = time.time()
        
        # 存储 embedding 为 BLOB
        embedding_blob = json.dumps(face_embedding).encode('utf-8')
        
        self._conn.execute('''
            INSERT INTO visual_memories (timestamp, frame_hash, embedding, face_snapshots, scene_description)
            VALUES (?, ?, ?, ?, ?)
        ''', (timestamp, frame_hash, embedding_blob, json.dumps([{'time': timestamp}]), scene_description))
        
        self._conn.commit()
        
    def _add_to_hot_data(self, timestamp: float, speaker_id: str, 
                         participant_identity: str, text: str,
                         scene_hash: str = None, embedding_id: int = None,
                         record_id: int = None, token_count: int = 0):
        """将对话添加到热数据内存"""
        record = {
            'id': record_id,
            'timestamp': timestamp,
            'speaker_id': speaker_id,
            'participant': participant_identity,
            'text': text,
            'scene_hash': scene_hash,
            'embedding_id': embedding_id,
            'token_count': token_count,
            'archived': 0,
            'compacted': 0
        }
        self._hot_dialogues.append(record)
        
        # 保持热数据大小在限制内
        if len(self._hot_dialogues) > self.hot_limit:
            self._hot_dialogues = self._hot_dialogues[-self.hot_limit:]
        
        # 更新说话人统计
        if participant_identity not in self._hot_speakers:
            self._hot_speakers[participant_identity] = {
                'total_turns': 0,
                'first_seen': timestamp,
                'last_seen': timestamp
            }
        self._hot_speakers[participant_identity]['total_turns'] += 1
        self._hot_speakers[participant_identity]['last_seen'] = timestamp
        
    def search_dialogues(self, query: str, top_k: int = 5, speaker_id: str = None) -> List[Dict]:
        """搜索历史对话（优先热数据，fallback 到冷数据，支持 FTS5）"""
        results = []
        
        # 先从热数据搜索
        if speaker_id:
            filtered = [d for d in self._hot_dialogues 
                       if d['participant'] == speaker_id and d.get('compacted') != 1]
            results = filtered[-top_k:]
        else:
            # 简单关键词过滤
            query_lower = query.lower()
            filtered = [d for d in self._hot_dialogues 
                       if query_lower in d['text'].lower() and d.get('compacted') != 1]
            results = filtered[-top_k:]
        
        # 如果热数据不够，从冷数据补充
        if len(results) < top_k and self._conn:
            cold_results = self._search_cold(query, top_k - len(results), speaker_id)
            results.extend(cold_results)
        
        return results
    
    def search_with_fts(self, query: str, top_k: int = 10, speaker_id: str = None) -> List[Dict]:
        """使用 FTS5 全文搜索"""
        if not self._conn:
            return []
            
        c = self._conn.cursor()
        
        try:
            # 先尝试 FTS5 搜索
            if speaker_id:
                c.execute('''
                    SELECT d.timestamp, d.speaker_id, d.participant_identity, d.text, d.scene_hash
                    FROM dialogues d
                    JOIN dialogues_fts fts ON d.id = fts.rowid
                    WHERE dialogues_fts MATCH ?
                    AND d.participant_identity = ?
                    ORDER BY rank
                    LIMIT ?
                ''', (query, speaker_id, top_k))
            else:
                c.execute('''
                    SELECT d.timestamp, d.speaker_id, d.participant_identity, d.text, d.scene_hash
                    FROM dialogues d
                    JOIN dialogues_fts fts ON d.id = fts.rowid
                    WHERE dialogues_fts MATCH ?
                    ORDER BY rank
                    LIMIT ?
                ''', (query, top_k))
                
            results = []
            for row in c.fetchall():
                results.append({
                    'timestamp': row[0],
                    'speaker_id': row[1],
                    'participant': row[2],
                    'text': row[3],
                    'scene_hash': row[4]
                })
            
            if results:
                return results
        except Exception as e:
            logger.warning(f"⚠️ FTS5 搜索失败，回退到普通搜索: {e}")
        
        # 回退到普通搜索
        return self.search_dialogues(query, top_k, speaker_id)
    
    def _search_cold(self, query: str, limit: int, speaker_id: str = None) -> List[Dict]:
        """从 SQLite 冷数据搜索"""
        if not self._conn:
            return []
            
        c = self._conn.cursor()
        
        if speaker_id:
            c.execute('''
                SELECT timestamp, speaker_id, participant_identity, text, scene_hash
                FROM dialogues
                WHERE participant_identity = ? AND compacted = 0
                ORDER BY timestamp DESC
                LIMIT ?
            ''', (speaker_id, limit))
        else:
            c.execute('''
                SELECT timestamp, speaker_id, participant_identity, text, scene_hash
                FROM dialogues
                WHERE compacted = 0
                ORDER BY timestamp DESC
                LIMIT ?
            ''', (limit,))
            
        results = []
        for row in c.fetchall():
            results.append({
                'timestamp': row[0],
                'speaker_id': row[1],
                'participant': row[2],
                'text': row[3],
                'scene_hash': row[4]
            })
        return results
        
    def get_recent_summary(self, last_minutes: int = 10) -> str:
        """获取最近 N 分钟的对话摘要（优先热数据）"""
        cutoff = time.time() - (last_minutes * 60)
        
        # 先从热数据获取
        recent = [d for d in self._hot_dialogues 
                  if d['timestamp'] > cutoff and d.get('compacted') != 1]
        
        # 如果不够，从冷数据补充
        if len(recent) < 10 and self._conn:
            c = self._conn.cursor()
            c.execute('''
                SELECT participant_identity, text, timestamp
                FROM dialogues
                WHERE timestamp > ? AND compacted = 0
                ORDER BY timestamp DESC
            ''', (cutoff,))
            
            for row in c.fetchall():
                recent.append({
                    'timestamp': row[2],
                    'participant': row[0],
                    'text': row[1]
                })
        
        if not recent:
            return f"过去{last_minutes}分钟没有对话记录"
        
        # 按时间排序
        recent.sort(key=lambda x: x['timestamp'])
        
        lines = []
        for record in recent[-10:]:  # 最近10条
            from datetime import datetime
            time_str = datetime.fromtimestamp(record['timestamp']).strftime("%H:%M:%S")
            lines.append(f"[{time_str}] {record['participant']}: {record['text']}")
            
        return "\n".join(lines)
        
    def get_speaker_stats(self) -> Dict:
        """获取说话人统计（优先热数据，fallback 到冷数据）"""
        # 先返回热数据统计
        if self._hot_speakers:
            return self._hot_speakers
        
        # fallback 到 SQLite
        if not self._conn:
            return {}
            
        c = self._conn.cursor()
        c.execute('SELECT identity, total_turns, first_seen, last_seen FROM speakers ORDER BY total_turns DESC')
        
        stats = {}
        for row in c.fetchall():
            stats[row[0]] = {
                'total_turns': row[1],
                'first_seen': row[2],
                'last_seen': row[3]
            }
        return stats
    
    def _load_hot_data(self):
        """从 SQLite 加载热数据到内存"""
        if not self._conn:
            return
            
        c = self._conn.cursor()
        c.execute('''
            SELECT id, timestamp, speaker_id, participant_identity, text, 
                   scene_hash, embedding_id, token_count, archived, compacted
            FROM dialogues
            ORDER BY timestamp DESC
            LIMIT ?
        ''', (self.hot_limit,))
        
        self._hot_dialogues = []
        for row in c.fetchall():
            self._hot_dialogues.append({
                'id': row[0],
                'timestamp': row[1],
                'speaker_id': row[2],
                'participant': row[3],
                'text': row[4],
                'scene_hash': row[5],
                'embedding_id': row[6],
                'token_count': row[7] or 0,
                'archived': row[8] or 0,
                'compacted': row[9] or 0
            })
        
        # 恢复说话人统计
        c.execute('SELECT identity, total_turns, first_seen, last_seen FROM speakers')
        self._hot_speakers = {}
        for row in c.fetchall():
            self._hot_speakers[row[0]] = {
                'total_turns': row[1],
                'first_seen': row[2],
                'last_seen': row[3]
            }
        
        logger.info(f"📊 已加载 {len(self._hot_dialogues)} 条热数据到内存")
    
    def _check_compaction(self):
        """检查是否需要压缩旧对话"""
        # 计算热数据总 token 数
        total_tokens = sum(d.get('token_count', 0) for d in self._hot_dialogues)
        
        if total_tokens > self.compact_threshold and len(self._hot_dialogues) > 10:
            # 压缩最旧的一半
            self._hot_dialogues.sort(key=lambda x: x['timestamp'])
            half = len(self._hot_dialogues) // 2
            
            for record in self._hot_dialogues[:half]:
                if record.get('id'):
                    self._conn.execute(
                        'UPDATE dialogues SET compacted = 1 WHERE id = ?',
                        (record['id'],)
                    )
            
            # 保留后半部分和所有新数据
            self._hot_dialogues = self._hot_dialogues[half:]
            self._conn.commit()
            logger.info(f"💾 已压缩 {half} 条旧对话")
        
    def archive_old_dialogues(self, days: int = None):
        """归档 N 天前的对话"""
        if days is None:
            days = self.archive_days
            
        if not self._conn:
            return 0
            
        cutoff = time.time() - (days * 86400)
        c = self._conn.cursor()
        
        c.execute('''
            UPDATE dialogues 
            SET archived = 1 
            WHERE timestamp < ? AND compacted = 0
        ''', (cutoff,))
        
        archived_count = c.rowcount
        self._conn.commit()
        
        if archived_count > 0:
            logger.info(f"📦 已归档 {archived_count} 条旧对话")
        
        return archived_count
    
    def get_stats(self) -> Dict:
        """获取详细统计信息"""
        if not self._conn:
            return {}
            
        c = self._conn.cursor()
        
        # 总数统计
        c.execute('SELECT COUNT(*) FROM dialogues')
        total = c.fetchone()[0]
        
        c.execute('SELECT COUNT(*) FROM dialogues WHERE archived = 0')
        active = c.fetchone()[0]
        
        c.execute('SELECT COUNT(*) FROM dialogues WHERE compacted = 1')
        compacted = c.fetchone()[0]
        
        c.execute('SELECT COUNT(*) FROM dialogues WHERE archived = 1')
        archived = c.fetchone()[0]
        
        # Token 统计
        c.execute('SELECT COALESCE(SUM(token_count), 0) FROM dialogues')
        total_tokens = c.fetchone()[0]
        
        return {
            'total': total,
            'active': active,
            'compacted': compacted,
            'archived': archived,
            'total_tokens': total_tokens,
            'hot_data_count': len(self._hot_dialogues),
            'speakers': self.get_speaker_stats()
        }
    
    def close(self):
        """关闭数据库连接"""
        if self._conn:
            self._conn.close()
            self._conn = None
        logger.info(f"📊 关闭对话记录器 (热数据: {len(self._hot_dialogues)} 条)")


class FaceMatcher:
    """
    简单的人脸匹配器 - 基于 WeMM embedding 匹配
    """
    
    def __init__(self, threshold: float = 0.7):
        self.threshold = threshold
        self.known_faces: Dict[str, List[float]] = {}  # identity -> embedding
        
    def register_face(self, identity: str, embedding: List[float]):
        """注册人脸"""
        self.known_faces[identity] = embedding
        logger.info(f"👤 注册人脸: {identity}")
        
    def identify_face(self, embedding: List[float]) -> Optional[str]:
        """识别人脸"""
        if not self.known_faces:
            return None
            
        best_match = None
        best_score = 0
        
        for identity, known_emb in self.known_faces.items():
            score = self._cosine_similarity(embedding, known_emb)
            if score > best_score:
                best_score = score
                best_match = identity
                
        if best_score >= self.threshold:
            return best_match
        return None
        
    def _cosine_similarity(self, a: List[float], b: List[float]) -> float:
        """计算余弦相似度"""
        import numpy as np
        va = np.array(a)
        vb = np.array(b)
        return float(np.dot(va, vb) / (np.linalg.norm(va) * np.linalg.norm(vb)))