import os
import sqlite3

from pprint import pprint

from langchain_community.vectorstores import Chroma
from langchain_community.embeddings import LlamaCppEmbeddings

import chromadb

def read_aqlite3():
    # 连接到 SQLite 数据库（如果数据库不存在，会自动创建）
    # 请替换为你的数据库文件路径
    db_path = 'src/Model/AI_Agent/scripts/ai_agent/memory/knowledge_base/map/vector_store/osm_map4/chroma.sqlite3'  
    conn = sqlite3.connect(db_path)

    # 创建一个游标对象，用于执行 SQL 查询
    cursor = conn.cursor()

    # 执行查询，获取所有表名
    cursor.execute("SELECT name FROM sqlite_master WHERE type='table';")

    # 获取所有表的名称
    tables = cursor.fetchall()
    print("数据库中的所有表:")
    for table in tables:
        print(table[0])  # 打印每个表的名称

    # 假设我们选择某个表进行查询，获取表的前 10 行数据
    table_name = 'your_table_name'  # 替换为你实际的表名
    cursor.execute(f"SELECT * FROM {table_name} LIMIT 10;")

    # 获取查询结果
    rows = cursor.fetchall()
    print(f"{table_name} 表的前 10 行数据:")
    for row in rows:
        print(row)

    # 完成后关闭游标和连接
    cursor.close()
    conn.close()

def read_chromdb():

    # 加载现有的 Chroma 数据库
    # 替换为实际的路径
    db_dir = "/home/ubuntu/Workspace/Projects/VLM_VLA/UAV_AI/src/Model/AI_Agent/scripts/ai_agent/memory/knowledge_base/map/vector_store/osm_map1"  
    # vector_db = Chroma(persist_directory=db_dir)

    # # 获取所有文档
    # documents = vector_db.get_documents()

    # # 输出前 5 个文档的信息
    # print(f"共有 {len(documents)} 个文档存储在 Chroma 数据库中。")
    # for i, doc in enumerate(documents[:5]):
    #     print(f"文档 {i + 1}: {doc.metadata['source']}")  # 输出每个文档的源信息（文件名等）
    #     print(f"内容: {doc.page_content[:300]}...")  # 打印文档的前 300 个字符
    #     print("-" * 50)

    vector_db = Chroma(persist_directory=db_dir)

    # # Retrieve all document IDs
    document_ids = vector_db.get_all_ids()  # Get all document IDs

    # Retrieve all documents based on the IDs
    documents = vector_db.get_by_ids(document_ids)

    # Print out all documents' metadata and content
    for i, doc in enumerate(documents, 1):
        print(f"\nDocument {i} (Source: {doc.metadata.get('source', 'N/A')}):")
        print(f"  Content: {doc.page_content[:500]}...")  # Show first 500 characters to avoid overflow
        print(f"  Metadata: {doc.metadata}")


class ChromaDataViewer:
    def __init__(
        self,
        persist_directory: str,
        embedding_model_path: str,
        collection_name: str = "langchain",
        n_ctx: int = 1024,
        n_threads: int = 4
    ):
        """初始化Chroma向量数据库查看器"""
        self.persist_directory = persist_directory
        self.collection_name = collection_name
        
        # 初始化与向量库匹配的嵌入模型（必须与入库时一致）
        self.embeddings = LlamaCppEmbeddings(
            model_path=embedding_model_path,
            n_ctx=n_ctx,
            n_threads=n_threads
            # embedding=True  # 启用嵌入模式
        )
        
        # 连接到现有向量库
        self.vector_db = self._connect_to_chroma()

    def _connect_to_chroma(self) -> Chroma:
        """连接到Chroma向量数据库"""
        if not os.path.exists(self.persist_directory):
            raise FileNotFoundError(f"向量库路径不存在：{self.persist_directory}")
        
        try:
            return Chroma(
                persist_directory=self.persist_directory,
                embedding_function=self.embeddings,
                collection_name=self.collection_name
            )
        except Exception as e:
            raise RuntimeError(f"连接向量库失败：{str(e)}")

    def get_collection_info(self) -> dict:
        """获取向量库集合基本信息"""
        # 通过底层Chroma客户端获取集合信息
        collection = self.vector_db._collection
        return {
            "collection_name": self.collection_name,
            "document_count": collection.count(),  # 文档总数
            "vector_dimension": collection.get(limit=1)["embeddings"][0].shape[0] if collection.count() > 0 else 0,  # 向量维度
            "persist_directory": self.persist_directory
        }

    def list_all_document_ids(self) -> list[str]:
        """获取所有文档的ID列表"""
        return self.vector_db._collection.get()["ids"]

    def get_documents_by_ids(self, ids: list[str] = None, limit: int = 5) -> list[dict]:
        """
        获取指定ID的文档内容和元数据
        :param ids: 文档ID列表（None则获取前limit个）
        :param limit: 最大返回数量（ids为None时生效）
        """
        if ids is None:
            all_ids = self.list_all_document_ids()
            ids = all_ids[:limit]  # 取前limit个
        
        result = self.vector_db._collection.get(ids=ids)
        documents = []
        for i in range(len(ids)):
            # 匹配ID对应的文档内容和元数据
            doc_idx = result["ids"].index(ids[i]) if ids[i] in result["ids"] else -1
            if doc_idx == -1:
                continue
            documents.append({
                "id": ids[i],
                "content": result["documents"][doc_idx],  # 文档文本内容
                "metadata": result["metadatas"][doc_idx],  # 元数据（如来源、类型）
                "embedding_sample": result["embeddings"][doc_idx][:5]  # 向量前5维（示例）
            })
        return documents

    def search_similar_documents(self, query: str, k: int = 3) -> list[dict]:
        """检索与查询相似的文档（带相似度分数）"""
        results = self.vector_db.similarity_search_with_score(query, k=k)
        similar_docs = []
        for doc, score in results:
            similar_docs.append({
                "content": doc.page_content,
                "metadata": doc.metadata,
                "similarity_score": score,  # 相似度分数（越小越相似，取决于距离函数）
                "id": doc.metadata.get("id", "未知ID")
            })
        return similar_docs

def read_chromdb2():
    # 配置向量库路径和嵌入模型路径（替换为你的实际路径）
    # 向量库存储目录
    PERSIST_DIRECTORY = "/home/ubuntu/Workspace/Projects/VLM_VLA/UAV_AI/src/Model/AI_Agent/scripts/ai_agent/memory/knowledge_base/map/vector_store/osm_map1/"  
    EMBEDDING_MODEL_PATH = "/home/ubuntu/Workspace/Projects/VLM_VLA/Models/Embedding-GGUF/all-MiniLM-L6-v2-GGUF/all-MiniLM-L6-v2.F16.gguf"  # 与入库时相同的嵌入模型
    COLLECTION_NAME = "drone_docs"  # 集合名（与入库时一致）

    # 初始化查看器
    viewer = ChromaDataViewer(
        persist_directory=PERSIST_DIRECTORY,
        embedding_model_path=EMBEDDING_MODEL_PATH,
        collection_name=COLLECTION_NAME,
        n_threads=4
    )

    # 1. 查看集合基本信息
    print("=== 向量库集合信息 ===")
    pprint(viewer.get_collection_info())
    print("\n" + "="*50 + "\n")

    # 2. 查看前3个文档的详细信息（内容、元数据、向量示例）
    print("=== 前3个文档详情 ===")
    sample_docs = viewer.get_documents_by_ids(limit=3)
    for i, doc in enumerate(sample_docs, 1):
        print(f"\n--- 文档 {i} ---")
        print(f"ID: {doc['id']}")
        print(f"来源: {doc['metadata'].get('source', '未知')}")
        print(f"内容预览: {doc['content'][:200]}...")  # 显示前200字符
        print(f"向量前5维: {doc['embedding_sample']}")
    print("\n" + "="*50 + "\n")

    # # 3. 检索相似文档（示例）
    # test_query = "你的查询关键词"  # 替换为你的查询
    # print(f"=== 与 '{test_query}' 相似的文档 ===")
    # similar_docs = viewer.search_similar_documents(query=test_query, k=2)
    # for i, doc in enumerate(similar_docs, 1):
    #     print(f"\n--- 相似文档 {i} (分数: {doc['similarity_score']:.4f}) ---")
    #     print(f"来源: {doc['metadata'].get('source', '未知')}")
    #     print(f"内容预览: {doc['content'][:200]}...")



def read_chromadb_data(data_path,collection_name,show_data_count):

    # ---------------------- 1. 连接ChromaDB并提取原数据 ----------------------
    # 1.1 初始化Chroma客户端（根据你的实际存储方式选择，这里以持久化存储为例）
    client = chromadb.PersistentClient(path=data_path)  # 替换为你的ChromaDB存储路径
    #client = chromadb.PersistentClient(path="./osm_map1")  # 替换为你的ChromaDB存储路径


    # 1.2 获取目标集合（Collection）
    #collection = client.get_collection(name="drone_docs")  # 替换为你的集合名
    collection = client.get_collection(name=collection_name)  # 替换为你的集合名
    print("all data len:",collection.count())

    # breakpoint()
    # cnt=10

    data = collection.get(include=['embeddings', 'documents', 'metadatas'])
    all_info={}
    for i in range(show_data_count):
        ids=data["ids"][i]
        docs=data["documents"][i]
        embedings=data["embeddings"][i]
        metadata=data["metadatas"][i]
        
        all_info["ids"]=ids
        all_info["docs"]=docs
        all_info["embeddings"]=embedings
        all_info["metadata"]=metadata
        
        print(all_info)

def main():
    vector_data_path = "/home/ubuntu/Workspace/Projects/VLM_VLA/UAV_AI/src/Model/AI_Agent/scripts/ai_agent/memory/knowledge_base/map/vector_store/osm_map_4/"
    read_chromadb_data(data_path=vector_data_path,
                       collection_name="osm_map_docs",
                       show_data_count=10)
    print("main finished!")
    

if __name__ == "__main__":
    main()
    