"""
Embedding 服务 - 处理语义搜索和文本分块

当前没启用，启用需要修改环境变量 EMBEDDINGS_AVAILABLE 为 True
"""
import logging
import numpy as np
from sentence_transformers import SentenceTransformer
from sklearn.metrics.pairwise import cosine_similarity
from sqlalchemy.orm import joinedload
from src.clients import db
from src.models import Recording, TranscriptChunk, RecordingTag

logger = logging.getLogger(__name__)

# 全局 Embedding 模型实例
_embedding_model = None

# 检查 Embedding 功能是否可用
EMBEDDINGS_AVAILABLE = False
try:
    import numpy as np
    from sentence_transformers import SentenceTransformer
    from sklearn.metrics.pairwise import cosine_similarity
    EMBEDDINGS_AVAILABLE = True
except ImportError:
    pass


def get_embedding_model():
    """获取或初始化 Sentence Transformer 模型"""
    global _embedding_model
    
    if not EMBEDDINGS_AVAILABLE:
        return None
        
    if _embedding_model is None:
        try:
            _embedding_model = SentenceTransformer('all-MiniLM-L6-v2')
            logger.info("Embedding 模型加载成功")
        except Exception as e:
            logger.error(f"Embedding 模型加载失败：{e}")
            return None
    return _embedding_model


def chunk_transcription(transcription, max_chunk_length=500, overlap=50):
    """
    将转录文本分割为带重叠的 chunk，以便更好地进行上下文检索
    
    Args:
        transcription (str): 完整的转录文本
        max_chunk_length (int): 每个 chunk 的最大字符数
        overlap (int): chunk 之间的字符重叠量
    
    Returns:
        list: 文本 chunk 列表
    """
    if not transcription or len(transcription) <= max_chunk_length:
        return [transcription] if transcription else []
    
    chunks = []
    start = 0
    
    while start < len(transcription):
        end = start + max_chunk_length
        
        # 尝试在句子边界处分割
        if end < len(transcription):
            # 在最后 100 个字符内查找句子结尾
            sentence_end = -1
            for i in range(max(0, end - 100), end):
                if transcription[i] in '.!?':
                    # 检查是否不是缩写
                    if i + 1 < len(transcription) and transcription[i + 1].isspace():
                        sentence_end = i + 1
            
            if sentence_end > start:
                end = sentence_end
        
        chunk = transcription[start:end].strip()
        if chunk:
            chunks.append(chunk)
        
        # 移动起始位置，包含重叠部分
        start = max(start + 1, end - overlap)
        
        # 防止无限循环
        if start >= len(transcription):
            break
    
    return chunks


def generate_embeddings(texts):
    """
    为文本列表生成 embedding 向量
    
    Args:
        texts (list): 文本字符串列表
    
    Returns:
        list: embedding 向量列表（numpy 数组），如果 embedding 不可用则返回空列表
    """
    if not EMBEDDINGS_AVAILABLE:
        logger.warning("Embedding 功能不可用 - 跳过 embedding 生成")
        return []
        
    model = get_embedding_model()
    if not model or not texts:
        return []
    
    try:
        embeddings = model.encode(texts)
        return [embedding.astype(np.float32) for embedding in embeddings]
    except Exception as e:
        logger.error(f"生成 embedding 向量时出错：{e}")
        return []


def serialize_embedding(embedding):
    """将 numpy 数组转换为二进制以便数据库存储"""
    if embedding is None or not EMBEDDINGS_AVAILABLE:
        return None
    return embedding.tobytes()


def deserialize_embedding(binary_data):
    """将二进制数据转换回 numpy 数组"""
    if binary_data is None or not EMBEDDINGS_AVAILABLE:
        return None
    return np.frombuffer(binary_data, dtype=np.float32)


def process_recording_chunks(recording_id):
    """
    处理录音：创建 chunk 并生成 embedding
    应在录音转录完成后调用此函数
    """
    try:
        recording = db.session.get(Recording, recording_id)
        if not recording or not recording.transcription:
            return False
        
        # 删除此录音的现有 chunk
        TranscriptChunk.query.filter_by(recording_id=recording_id).delete()
        
        # 创建 chunk
        chunks = chunk_transcription(recording.transcription)
        
        if not chunks:
            return True
        
        # 生成 embedding
        embeddings = generate_embeddings(chunks)
        
        # 将 chunk 存储到数据库
        for i, (chunk_text, embedding) in enumerate(zip(chunks, embeddings)):
            chunk = TranscriptChunk(
                recording_id=recording_id,
                user_id=recording.user_id,
                chunk_index=i,
                content=chunk_text,
                embedding=serialize_embedding(embedding) if embedding is not None else None
            )
            db.session.add(chunk)
        
        db.session.commit()
        logger.info(f"为录音 {recording_id} 创建了 {len(chunks)} 个 chunk")
        return True
        
    except Exception as e:
        logger.error(f"处理录音 {recording_id} 的 chunk 时出错：{e}")
        db.session.rollback()
        return False


def basic_text_search_chunks(user_id, query, filters=None, top_k=5):
    """
    当 embedding 不可用时使用的基本文本搜索回退方案
    使用简单文本匹配代替语义搜索
    """
    try:
        # 构建 chunk 的基础查询
        chunks_query = TranscriptChunk.query.filter_by(user_id=user_id)
        
        # 如果提供了过滤器则应用
        if filters:
            if filters.get('tag_ids'):
                chunks_query = chunks_query.join(Recording).join(
                    RecordingTag, Recording.id == RecordingTag.recording_id
                ).filter(RecordingTag.tag_id.in_(filters['tag_ids']))
            
            if filters.get('speaker_names'):
                # 通过录音的 participants 字段过滤，而不是 chunk 的 speaker_name
                if not any(hasattr(desc, 'name') and desc.name == 'recording' for desc in chunks_query.column_descriptions):
                    chunks_query = chunks_query.join(Recording)
                
                # 为每个说话人名称构建 OR 条件
                speaker_conditions = []
                for speaker_name in filters['speaker_names']:
                    speaker_conditions.append(
                        Recording.participants.ilike(f'%{speaker_name}%')
                    )
                
                chunks_query = chunks_query.filter(db.or_(*speaker_conditions))
                logger.info(f"已应用说话人过滤器：{filters['speaker_names']}")
            
            if filters.get('recording_ids'):
                chunks_query = chunks_query.filter(
                    TranscriptChunk.recording_id.in_(filters['recording_ids'])
                )
            
            if filters.get('date_from') or filters.get('date_to'):
                chunks_query = chunks_query.join(Recording)
                if filters.get('date_from'):
                    chunks_query = chunks_query.filter(Recording.meeting_date >= filters['date_from'])
                if filters.get('date_to'):
                    chunks_query = chunks_query.filter(Recording.meeting_date <= filters['date_to'])
        
        # 简单文本搜索 - 将查询拆分为单词并搜索
        query_words = query.lower().split()
        if query_words:
            # 创建一个过滤器，匹配内容中的任意查询单词
            text_conditions = []
            for word in query_words:
                text_conditions.append(TranscriptChunk.content.ilike(f'%{word}%'))
            
            # 使用 OR 组合条件
            from sqlalchemy import or_
            chunks_query = chunks_query.filter(or_(*text_conditions))
        
        # 获取 chunk 并返回（带有虚拟相似度分数）
        chunks = chunks_query.limit(top_k).all()
        
        # 返回 chunk 及虚拟相似度分数（找到的 chunk 为 1.0）
        return [(chunk, 1.0) for chunk in chunks]
        
    except Exception as e:
        logger.error(f"基本文本搜索出错：{e}")
        return []


def semantic_search_chunks(user_id, query, filters=None, top_k=5):
    """
    在转录 chunk 上执行带过滤的语义搜索
    
    Args:
        user_id (int): 用于权限过滤的用户 ID
        query (str): 搜索查询
        filters (dict): 可选的标签、说话人、日期、录音 ID 过滤器
        top_k (int): 返回的顶级 chunk 数量
    
    Returns:
        list: 相关 chunk 及相似度分数列表
    """
    try:
        # 如果 embedding 不可用，回退到基本文本搜索
        if not EMBEDDINGS_AVAILABLE:
            logger.info("Embedding 功能不可用 - 使用基本文本搜索作为回退")
            return basic_text_search_chunks(user_id, query, filters, top_k)
        
        # 为查询生成 embedding
        model = get_embedding_model()
        if not model:
            return basic_text_search_chunks(user_id, query, filters, top_k)
        
        query_embedding = model.encode([query])[0]
        
        # 构建 chunk 的基础查询，预先加载录音关系
        chunks_query = TranscriptChunk.query.options(joinedload(TranscriptChunk.recording)).filter_by(user_id=user_id)
        
        # 如果提供了过滤器则应用
        if filters:
            if filters.get('tag_ids'):
                # 与具有指定标签的录音连接
                chunks_query = chunks_query.join(Recording).join(
                    RecordingTag, Recording.id == RecordingTag.recording_id
                ).filter(RecordingTag.tag_id.in_(filters['tag_ids']))
            
            if filters.get('speaker_names'):
                # 通过录音的 participants 字段过滤
                if not any(hasattr(desc, 'name') and desc.name == 'recording' for desc in chunks_query.column_descriptions):
                    chunks_query = chunks_query.join(Recording)
                
                # 为每个说话人名称构建 OR 条件
                speaker_conditions = []
                for speaker_name in filters['speaker_names']:
                    speaker_conditions.append(
                        Recording.participants.ilike(f'%{speaker_name}%')
                    )
                
                chunks_query = chunks_query.filter(db.or_(*speaker_conditions))
                logger.info(f"已应用说话人过滤器：{filters['speaker_names']}")
            
            if filters.get('recording_ids'):
                chunks_query = chunks_query.filter(
                    TranscriptChunk.recording_id.in_(filters['recording_ids'])
                )
            
            if filters.get('date_from') or filters.get('date_to'):
                chunks_query = chunks_query.join(Recording)
                if filters.get('date_from'):
                    chunks_query = chunks_query.filter(Recording.meeting_date >= filters['date_from'])
                if filters.get('date_to'):
                    chunks_query = chunks_query.filter(Recording.meeting_date <= filters['date_to'])
        
        # 获取具有 embedding 的 chunk
        chunks = chunks_query.filter(TranscriptChunk.embedding.isnot(None)).all()
        
        if not chunks:
            return []
        
        # 计算相似度
        chunk_similarities = []
        for chunk in chunks:
            try:
                chunk_embedding = deserialize_embedding(chunk.embedding)
                if chunk_embedding is not None:
                    similarity = cosine_similarity(
                        query_embedding.reshape(1, -1),
                        chunk_embedding.reshape(1, -1)
                    )[0][0]
                    chunk_similarities.append((chunk, float(similarity)))
            except Exception as e:
                logger.warning(f"计算 chunk {chunk.id} 的相似度时出错：{e}")
                continue
        
        # 按相似度排序并返回前 k 个
        chunk_similarities.sort(key=lambda x: x[1], reverse=True)
        return chunk_similarities[:top_k]
        
    except Exception as e:
        logger.error(f"语义搜索出错：{e}")
        return []
