# corpus_builder.py
import chromadb
from chromadb.config import Settings
import hashlib
import uuid
from typing import List, Dict
import numpy as np
from sentence_transformers import SentenceTransformer
import pickle
import os


class LiteratureCorpus:
    """文献语料库系统"""

    def __init__(self, persist_directory="./corpus_db"):
        self.persist_directory = persist_directory

        # 初始化ChromaDB
        self.client = chromadb.Client(Settings(
            chroma_db_impl="duckdb+parquet",
            persist_directory=persist_directory
        ))

        # 加载嵌入模型（可以选择小一点的模型加速）
        print("加载嵌入模型...")
        self.embedding_model = SentenceTransformer('all-MiniLM-L6-v2')

        # 创建或获取集合
        self.collection = self.client.get_or_create_collection(
            name="literature_corpus",
            metadata={"hnsw:space": "cosine"}
        )

    def add_document(self, article_id: str, metadata: Dict, segments: List[Dict]):
        """添加文档到语料库"""

        documents = []
        metadatas = []
        ids = []

        for seg_idx, segment in enumerate(segments):
            # 为每个段落创建唯一ID
            seg_id = f"{article_id}_seg_{seg_idx}"

            # 准备文档文本
            doc_text = segment['text']
            if len(doc_text) > 10000:  # 太长则截断
                doc_text = doc_text[:10000]

            documents.append(doc_text)

            # 准备元数据
            seg_metadata = {
                'article_id': article_id,
                'section_type': segment['section_type'],
                'section_title': segment['section_title'],
                'page_number': segment['page_number'],
                'article_title': metadata.get('title', 'Unknown'),
                'authors': metadata.get('authors', 'Unknown'),
                'year': metadata.get('year', 'Unknown'),
                'journal': metadata.get('publicationTitle', 'Unknown'),
                'word_count': segment['word_count'],
                'segment_index': seg_idx,
                'full_text_preview': doc_text[:500]  # 预览文本
            }
            seg_metadata.update(metadata)
            metadatas.append(seg_metadata)
            ids.append(seg_id)

        # 添加到向量数据库
        if documents:
            embeddings = self.embedding_model.encode(documents).tolist()
            self.collection.add(
                embeddings=embeddings,
                documents=documents,
                metadatas=metadatas,
                ids=ids
            )

    def search(self, query: str, n_results: int = 10,
               filter_conditions: Dict = None) -> List[Dict]:
        """搜索语料库"""

        # 生成查询嵌入
        query_embedding = self.embedding_model.encode([query]).tolist()

        # 执行搜索
        results = self.collection.query(
            query_embeddings=query_embedding,
            n_results=n_results,
            where=filter_conditions,
            include=["documents", "metadatas", "distances"]
        )

        # 整理结果
        formatted_results = []
        if results['documents']:
            for i in range(len(results['documents'][0])):
                result = {
                    'document': results['documents'][0][i],
                    'metadata': results['metadatas'][0][i],
                    'distance': results['distances'][0][i],
                    'relevance_score': 1 - results['distances'][0][i]  # 转换为相关性分数
                }
                formatted_results.append(result)

        return formatted_results

    def search_by_section(self, query: str, section_type: str = None):
        """按章节搜索"""
        filter_condition = None
        if section_type:
            filter_condition = {"section_type": section_type}

        return self.search(query, filter_conditions=filter_condition)

    def get_statistics(self):
        """获取语料库统计信息"""
        count = self.collection.count()

        # 获取所有文档的元数据
        all_docs = self.collection.get()

        # 按章节统计
        section_stats = {}
        if all_docs['metadatas']:
            for metadata in all_docs['metadatas']:
                section = metadata.get('section_type', 'unknown')
                if section not in section_stats:
                    section_stats[section] = 0
                section_stats[section] += 1

        return {
            'total_segments': count,
            'section_distribution': section_stats,
            'unique_articles': len(set([m.get('article_id') for m in all_docs['metadatas']]))
        }