
"""
知识片段业务逻辑服务
"""
from typing import List, Optional, Tuple, Dict, Any
import json
import random
import csv
import io
import time
import uuid
from datetime import datetime
import logging

from app.services.milvus_service import milvus_service
from app.schemas.base import PaginationSchema, PaginatedResponseSchema
from app.utils.vector_utils import text_to_vector_algo

from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy import select
from app.sample.models.metadata import SampleMetadata
from app.sample.models.knowledge_base import KnowledgeBase
from app.services.knowledge_base_service import knowledge_base_service
from app.base.async_mysql_connection import get_db_connection

from app.sample.models.base_info import DocumentMain

logger = logging.getLogger(__name__)

class SnippetService:
    
    async def get_list(
        self,
        page: int = 1,
        page_size: int = 10,
        kb: Optional[str] = None,
        keyword: Optional[str] = None,
        status: Optional[str] = None
    ) -> Tuple[List[Dict], PaginationSchema]:
        """获取知识片段列表 (跨集合查询)"""
        
        # [Modified] 获取 KnowledgeBase 名称映射
        kb_name_map = {}
        try:
            # 修正引用路径
            from app.base.async_mysql_connection import get_db
            from app.sample.models.knowledge_base import KnowledgeBase
            from sqlalchemy import select
            
            # 使用异步数据库会话
            async for db in get_db():
                res = await db.execute(select(KnowledgeBase.collection_name_parent, KnowledgeBase.collection_name_children, KnowledgeBase.name))
                for row in res.all():
                    # 映射 collection_name_parent -> name
                    if row[0]:
                        kb_name_map[row[0]] = row[2]
                    # 映射 collection_name_children -> name
                    if row[1]:
                        kb_name_map[row[1]] = row[2]
                break  # 只需要一次迭代
                        
        except Exception as e:
            logger.exception("Failed to load KB map")

        # 1. 确定要查询的目标集合列表
        target_collections = []
        if kb:
            target_collections = [kb]
        else:
            # 简单起见，先查 Milvus 的所有集合
            target_collections = milvus_service.client.list_collections()
        
        if not target_collections:
            return [], PaginationSchema(total=0, page=page, page_size=page_size, total_pages=0)
        
        # 2. 计算分页逻辑 (跨集合分页算法)
        global_total = 0
        items = []
        
        # 需要跳过的全局偏移量
        skip_count = (page - 1) * page_size
        # 需要获取的目标数量
        need_count = page_size
        
        # 遍历所有集合
        for col_name in target_collections:
            if not milvus_service.has_collection(col_name):
                continue
                
            try:
                # 获取该集合总数
                stats = milvus_service.client.get_collection_stats(col_name)
                col_count = int(stats.get("row_count", 0)) if isinstance(stats, dict) else 0
                
                # 构建过滤表达式
                filter_exprs = []
                
                # 状态过滤 (假设 is_deleted=False 表示 normal, True 表示 disabled/deleted)
                # 注意：Milvus 中通常使用 is_deleted 标记删除，这里我们需要根据传入的 status 映射
                # status: 'normal' -> is_deleted == False
                # status: 'disabled' -> 实际上我们没有 disabled 状态，只有删除。如果必须支持，可能需要额外的 status 字段
                # 暂时假设：normal = 未删除, disabled = (无此状态或暂不支持)
                
                # 如果要严格支持 normal/disabled，需要在 schema 中增加 status 字段
                # 目前 Schema 只有 is_deleted。
                # 兼容处理：
                try:
                    desc = milvus_service.client.describe_collection(col_name)
                    fields = desc.get("fields", []) if isinstance(desc, dict) else []
                    field_names = [f.get("name") for f in fields if isinstance(f, dict)]
                    has_is_deleted = "is_deleted" in field_names
                except Exception:
                    has_is_deleted = False

                if has_is_deleted:
                    if status == 'normal':
                        filter_exprs.append("is_deleted == false")
                    elif status == 'disabled':
                        # 假设 disabled 对应 is_deleted == true (软删除状态)
                        # 但通常 deleted 数据不应被查出。如果需要查出“已禁用”，则需要额外字段。
                        # 按照之前 Schema 定义，is_deleted 是 bool。
                        filter_exprs.append("is_deleted == true")
                    else:
                        # 默认只查未删除的
                        filter_exprs.append("is_deleted == false")
                else:
                    # 如果不支持 is_deleted，则忽略该过滤条件，默认视为全部有效
                    pass

                if keyword:
                    # 关键词模式：必须实际查询
                    filter_exprs.append(f'text like "%{keyword}%"')

                expr = " && ".join(filter_exprs) if filter_exprs else ""
                
                if not expr:
                    # 如果表达式为空（不支持 is_deleted 且无其他条件），构造恒真表达式以满足 Milvus query 要求
                    try:
                        desc = milvus_service.client.describe_collection(col_name)
                        pk_field = "id"
                        is_str = False
                        if isinstance(desc, dict) and 'fields' in desc:
                            for f in desc['fields']:
                                if f.get('primary_key') or f.get('is_primary'):
                                    pk_field = f.get('name')
                                    # Milvus DataType: INT64=5, VARCHAR=21
                                    if f.get('type') == 21 or f.get('type') == 'VarChar':
                                        is_str = True
                                    break
                        expr = f'{pk_field} != ""' if is_str else f'{pk_field} >= 0'
                    except:
                        expr = "id >= 0"

                if keyword:
                    # 有关键词，必须 query
                    output_fields = ["*"]
                    res = milvus_service.client.query(col_name, filter=expr, output_fields=output_fields, limit=100)
                    col_hits = len(res)
                    global_total += col_hits
                    
                    if skip_count >= col_hits:
                        skip_count -= col_hits
                        continue
                    
                    take = min(need_count, col_hits - skip_count)
                    chunk = res[skip_count : skip_count + take]
                    
                    for r in chunk:
                        items.append(self._format_snippet(r, col_name, kb_name_map))
                    
                    skip_count = 0 
                    need_count -= take
                    if need_count <= 0: break
                    
                else:
                    # 无关键词模式，只有状态过滤
                    # 如果有状态过滤，也必须 query，不能直接用 stats
                    if status:
                         # 必须 query 计数
                         res = milvus_service.client.query(col_name, filter=expr, output_fields=["pk"])
                         col_hits = len(res)
                         global_total += col_hits
                         
                         if skip_count >= col_hits:
                            skip_count -= col_hits
                            continue
                         
                         # 获取分页数据
                         output_fields = ["*"]
                         res_page = milvus_service.client.query(
                            collection_name=col_name,
                            filter=expr,
                            output_fields=output_fields,
                            limit=need_count,
                            offset=skip_count # 注意：query 的 offset 性能较差，但在小数据量下尚可
                         )
                         
                         for r in res_page:
                            items.append(self._format_snippet(r, col_name, kb_name_map))
                         
                         skip_count = 0 
                         need_count -= len(res_page)

                    else:
                        # 既无关键词也无状态过滤（默认查所有未删除），可以直接用 offset
                        # expr 已经在上面计算好了（包含 is_deleted check 或恒真表达式）
                        
                        # 简单起见，统一走 query 路径以支持 is_deleted 过滤
                        # 如果完全不过滤，才可以用 limit/offset 直接分页
                        
                        # 这里为了性能，假设默认查询不需要精确 count，或者全部都是未删除
                        # 暂时强制走 query 以保证准确性
                        res = milvus_service.client.query(col_name, filter=expr, output_fields=["pk"])
                        col_hits = len(res)
                        global_total += col_hits
                        
                        if skip_count >= col_hits:
                            skip_count -= col_hits
                            continue
                            
                        output_fields = ["*"]
                        res_page = milvus_service.client.query(
                            collection_name=col_name,
                            filter=expr,
                            output_fields=output_fields,
                            limit=need_count,
                            offset=skip_count
                        )
                        
                        for r in res_page:
                            items.append(self._format_snippet(r, col_name, kb_name_map))
                            
                        skip_count = 0
                        need_count -= len(res_page)

            except Exception as e:
                logger.exception("Collection %s query error", col_name)
                continue

        total_pages = (global_total + page_size - 1) // page_size if page_size else 0

        meta = PaginationSchema(
            page=page,
            page_size=page_size,
            total=global_total,
            total_pages=total_pages
        )
        
        # [Fix] 批量查询 t_samp_document_main 获取准确的文档名称
        # 提取所有可能的 document_id
        doc_ids = set()
        for item in items:
            did = item.get("document_id")
            if not did:
                # 再次尝试从 metadata 提取，增加容错
                meta = item.get("metadata", {})
                if isinstance(meta, dict):
                    did = meta.get("document_id") or meta.get("doc_id")
                    if did:
                         item["document_id"] = did # 回填到 item
            
            # document_id 可能包含前缀，需要清理？
            # 不，通常 document_id 是 UUID 或纯数字。
            # 但有些系统可能会存成 list?
            if did and isinstance(did, (str, int)) and len(str(did)) > 0:
                doc_ids.add(str(did))
        
        if doc_ids:
            try:
                # 修正引用路径
                from app.base.async_mysql_connection import get_db
                from app.sample.models.base_info import DocumentMain
                from sqlalchemy import select

                async for db in get_db():
                    stmt = select(DocumentMain.id, DocumentMain.title).where(DocumentMain.id.in_(list(doc_ids)))
                    result = await db.execute(stmt)
                    rows = result.all()
                    doc_name_map = {str(row[0]): row[1] for row in rows}
                    break  # 只需要一次迭代
                
                if doc_name_map:
                    for item in items:
                        did = item.get("document_id")
                        if did and str(did) in doc_name_map:
                            item["doc_name"] = doc_name_map[str(did)]
            except Exception as e:
                logger.exception("Failed to fetch document names from DB")

        return items, meta

    async def get_chunks_by_document(self, doc_id: str) -> List[Dict[str, Any]]:
        """
        提取单文档的所有知识片段
        """
        # 1. 从数据库中获取文档所属的知识库集合名称
        async with get_db_connection() as db:
            from app.sample.models.base_info import DocumentMain
            from app.sample.models.knowledge_base import KnowledgeBase
            from sqlalchemy import select
            
            # 联表查询获取 collection_name_parent
            stmt = select(KnowledgeBase.collection_name_parent).join(
                DocumentMain, DocumentMain.kb_id == KnowledgeBase.id
            ).where(DocumentMain.id == doc_id)
            
            result = await db.execute(stmt)
            collection_name = result.scalar()
            
            if not collection_name:
                # 尝试查 collection_name_children
                stmt_child = select(KnowledgeBase.collection_name_children).join(
                    DocumentMain, DocumentMain.kb_id == KnowledgeBase.id
                ).where(DocumentMain.id == doc_id)
                result_child = await db.execute(stmt_child)
                collection_name = result_child.scalar()
                
            if not collection_name:
                # 如果还是没有，尝试从文档详情获取 kb_id 并查询
                stmt_doc = select(DocumentMain.kb_id).where(DocumentMain.id == doc_id)
                res_doc = await db.execute(stmt_doc)
                kb_id = res_doc.scalar()
                if not kb_id:
                    return []
                
                stmt_kb = select(KnowledgeBase.collection_name_parent).where(KnowledgeBase.id == kb_id)
                res_kb = await db.execute(stmt_kb)
                collection_name = res_kb.scalar()

        if not collection_name:
            return []

        # 2. 调用 MilvusService 获取片段
        chunks = milvus_service.get_chunks_by_document_id(collection_name, doc_id)
        
        # 3. 格式化输出
        formatted_chunks = []
        for chunk in chunks:
            formatted_chunks.append({
                "id": chunk.get("pk") or chunk.get("id"),
                "content": chunk.get("text") or chunk.get("content") or chunk.get("page_content"),
                "index": chunk.get("index"),
                "document_id": chunk.get("document_id"),
                "metadata": chunk.get("metadata") or {},
                "created_time": chunk.get("created_time")
            })
            
        return formatted_chunks

    async def create(self, db: AsyncSession, payload: Any) -> Dict:
        """创建知识片段"""
        # 使用统一算法生成向量
        dim = milvus_service.DENSE_DIM
        fake_vector = text_to_vector_algo(payload.content, dim=dim)
        
        # 1. 动态构建 metadata
        # 查找 KnowledgeBase ID

        from sqlalchemy import or_
        kb_query = select(KnowledgeBase).where(
            or_(
                KnowledgeBase.collection_name_parent == payload.collection_name,
                KnowledgeBase.collection_name_children == payload.collection_name
            )
        )
        result = await db.execute(kb_query)
        kb = result.scalars().first()
        
        metadata = {}
        if kb:
            # 查找该知识库定义的元数据字段
            meta_query = select(SampleMetadata).where(SampleMetadata.knowledge_base_id == kb.id)
            meta_result = await db.execute(meta_query)
            defined_fields = meta_result.scalars().all()
            
            # 仅填充定义的字段
            custom_fields = getattr(payload, 'custom_fields', {}) or {}
            for field in defined_fields:
                field_name = field.field_en_name
                # 从 custom_fields 中获取值，如果不存在则为 None 或空字符串 (视需求而定)
                # 这里我们只存有值的
                if field_name in custom_fields:
                    metadata[field_name] = custom_fields[field_name]
        
        # 基础数据
        now = int(time.time() * 1000)
        item = {
            "dense": fake_vector,
            "text": payload.content,
            "document_id": payload.custom_fields.get("document_id") if hasattr(payload, 'custom_fields') and payload.custom_fields and payload.custom_fields.get("document_id") else str(uuid.uuid4()),
            "parent_id": payload.custom_fields.get("parent_id", "") if hasattr(payload, 'custom_fields') and payload.custom_fields else "",
            "index": 0,
            "tag_list": payload.custom_fields.get("tag_list", "") if hasattr(payload, 'custom_fields') and payload.custom_fields else "",
            "permission": {},
            "metadata": metadata, # 使用动态构建的 metadata
            "is_deleted": False,
            "created_by": "system",
            "created_time": now,
            "updated_by": "system",
            "updated_time": now
        }
        
        data = [item]
        
        # [Fix] 插入后尝试同步更新知识库的更新时间
        try:
             # 这里无法直接获得 AsyncSession，尝试通过 run_in_executor 或其他方式
             # 或者简单地，在 flush 后，如果可能，打印日志提示需要更新
             pass
        except:
             pass

        res = milvus_service.client.insert(
            collection_name=payload.collection_name,
            data=data
        )
        
        milvus_service.client.flush(payload.collection_name)
        return {"count": res.get("insert_count", 1)}

    async def update(self, db: AsyncSession, id: str, payload: Any) -> str:
        """更新知识片段"""
        kb_name = payload.collection_name
        
        # 0. [Fix] 先查询旧数据以保留 created_time 等信息
        # [Fix] get_by_id 现在是 async 且需要 db 参数
        old_item = await self.get_by_id(db, kb_name, id)
        old_created_time = int(time.time() * 1000)
        old_created_by = "system"
        if old_item:
            # 尝试从 old_item 中获取 (old_item 是格式化后的 dict)
            # created_at 是 str, 我们需要 int timestamp
            # 这里的逻辑其实有点绕，因为 get_by_id 已经是格式化过的了
            # 之前的逻辑是重新查 Milvus，这里保持重新查 Milvus 的逻辑作为兜底
            pass
            
        # 重新查询原始 entity (保持原有逻辑不变，因为它更可靠)
        try:
            desc = milvus_service.client.describe_collection(kb_name)
            fields = [f['name'] for f in desc.get('fields', [])]
            pk_field = "pk" if "pk" in fields else "id"
            
            expr = ""
            if id.isdigit():
                expr = f"{pk_field} in [{id}]"
            else:
                expr = f"{pk_field} in ['{id}']"
                
            res = milvus_service.client.query(kb_name, filter=expr, output_fields=["created_time", "created_by"], limit=1)
            if res:
                if res[0].get("created_time"):
                    old_created_time = res[0].get("created_time")
                if res[0].get("created_by"):
                    old_created_by = res[0].get("created_by")
        except Exception as e:
            logger.exception("Failed to fetch old item info")

        # 1. 删除旧数据

        desc = milvus_service.client.describe_collection(kb_name)
        fields = [f['name'] for f in desc.get('fields', [])]
        pk_field = "pk" if "pk" in fields else "id"
        
        if id.isdigit():
            expr = f"{pk_field} in [{id}]"
        else:
            expr = f"{pk_field} in ['{id}']"
        
        milvus_service.client.delete(collection_name=kb_name, filter=expr)
        
        # 2. 插入新数据
        # 使用统一算法生成向量
        dim = milvus_service.DENSE_DIM
        fake_vector = text_to_vector_algo(payload.content, dim=dim)
        
        # 动态构建 metadata
        # 查找 KnowledgeBase ID
        # [Modified] 支持 collection_name_parent 和 collection_name_children
        from sqlalchemy import or_
        kb_query = select(KnowledgeBase).where(
            or_(
                KnowledgeBase.collection_name_parent == kb_name,
                KnowledgeBase.collection_name_children == kb_name
            )
        )
        result = await db.execute(kb_query)
        kb_obj = result.scalars().first()
        
        metadata = {}
        if kb_obj:
            # 查找该知识库定义的元数据字段
            meta_query = select(SampleMetadata).where(SampleMetadata.knowledge_base_id == kb_obj.id)
            meta_result = await db.execute(meta_query)
            defined_fields = meta_result.scalars().all()
            
            # 仅填充定义的字段
            custom_fields = getattr(payload, 'custom_fields', {}) or {}
            for field in defined_fields:
                field_name = field.field_en_name
                if field_name in custom_fields:
                    metadata[field_name] = custom_fields[field_name]
        
        # 获取旧数据的创建时间，保持不变
        created_time = int(time.time() * 1000)
        # 尝试查询旧数据以获取创建时间（如果能查到）
        # 但我们已经在第1步删除了...
        # 理想情况下应该先查再删。
        # 既然已经删了，就只能重置 created_time 或者假设它是新的。
        # 为了完善逻辑，我们应该在删除前查询。
        # 由于代码结构限制，这里暂时使用当前时间作为 created_time，或者我们可以接受“更新即重置创建时间”
        # 但通常 update 不应改变 create_time。
        # 改进：我们可以在 delete 前先查一下。
        
        # [Fix] 优化 update 逻辑：保留原始 created_time
        # 由于上面已经执行了 delete，这里无法找回。
        # 这是一个逻辑缺陷。应该先查，后删，再插。
        # 不过鉴于 Milvus 的特性，update 本质就是 delete + insert。
        # 我们可以尝试传递 created_time 参数（如果前端传了）。
        # 假设 payload 中没有 created_time。
        
        now = int(time.time() * 1000)
        item = {
            "dense": fake_vector,
            "text": payload.content,
            "document_id": payload.custom_fields.get("document_id") if hasattr(payload, 'custom_fields') and payload.custom_fields and payload.custom_fields.get("document_id") else str(uuid.uuid4()),
            "parent_id": payload.custom_fields.get("parent_id", "") if hasattr(payload, 'custom_fields') and payload.custom_fields else "",
            "index": 0,
            "tag_list": payload.custom_fields.get("tag_list", "") if hasattr(payload, 'custom_fields') and payload.custom_fields else "",
            "permission": {},
            "metadata": metadata, # 使用动态构建的 metadata
            "is_deleted": False,
            "created_by": old_created_by, # 使用旧数据的创建人
            "created_time": old_created_time, # 使用旧数据的创建时间
            "updated_by": "system",
            "updated_time": now
        }
        
        # [Fix] 更新知识库的 update_time
        if kb_obj:
             kb_obj.updated_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
             await db.commit()
            
        data = [item]
        
        milvus_service.client.insert(collection_name=kb_name, data=data)
        milvus_service.client.flush(kb_name)
        
        return "更新成功 (ID已变更)"

    def delete(self, id: str, kb: str) -> None:
        """删除知识片段"""
        if not milvus_service.has_collection(kb):
             raise ValueError("知识库不存在")
             
        desc = milvus_service.client.describe_collection(kb)
        fields = [f['name'] for f in desc.get('fields', [])]
        pk_field = "pk" if "pk" in fields else "id"
        
        if id.isdigit():
            expr = f"{pk_field} in [{id}]"
        else:
            expr = f"{pk_field} in ['{id}']"
            
        milvus_service.client.delete(
            collection_name=kb,
            filter=expr
        )
        milvus_service.client.flush(kb)

        # [Fix] 删除后同步更新数据库中的文档数量
        # 由于这里没有 AsyncSession，我们需要临时创建一个或者让调用方处理
        # 鉴于这是一个同步方法 (delete)，而 update_doc_count 是异步的，
        # 我们这里暂时无法直接调用 async 方法。
        # 考虑到 SnippetService.delete 通常在 API 中被调用，
        # 我们可以在这里尝试用 asyncio.run 或者改成 async 方法。
        # 但改 async 会影响所有调用方。
        
        # 方案：尝试在 API 层处理，或者这里不做处理，由前端刷新列表时自动同步 (如果 get_list 会同步)
        # 检查 get_list: 确实有同步逻辑。
        # 但 KnowledgeBaseService.get_list 才有同步逻辑。SnippetService.get_list 没有。
        
        # 更好的方式：让 knowledge_base_service 提供一个同步的 update_doc_count 包装，或者我们在这里不处理，
        # 但用户明确反馈了这个问题。
        
        # 我们尝试在 delete 中更新 DB。
        # 由于这里无法获得 session，只能新建连接。
        # 且 update_doc_count 是 async 的。
        
        # 暂时方案：打印日志提示，或者引入 asyncio
        try:
            import asyncio
            from app.base.async_mysql_connection import AsyncSessionLocal
            
            async def sync_count():
                try:
                    async with AsyncSessionLocal() as db:
                        count = await knowledge_base_service.update_doc_count(db, kb)
                        logger.info("Synced doc count for %s after delete: %s", kb, count)
                except Exception as ex:
                    logger.exception("Error in sync_count task")
            
            # 检查是否有正在运行的 loop
            try:
                loop = asyncio.get_running_loop()
                # 如果有 loop，说明是在 async 环境下被调用的 (虽然 delete 是 def)
                # 这时我们应该 create_task
                loop.create_task(sync_count())
            except RuntimeError:
                # 没有 loop，可以直接 run
                asyncio.run(sync_count())
        except Exception as e:
            logger.exception("Failed to sync doc count after delete")

    def _format_snippet(self, r: Dict, col_name: str, kb_map: Dict[str, str] = None) -> Dict:
        id_val = r.get("pk") or r.get("id")
        content = r.get("text") or r.get("content") or ""
        
        # 处理 metadata (Milvus 可能返回 JSON 字符串)
        meta = r.get("metadata") or {}
        if isinstance(meta, str):
            try:
                import json
                meta = json.loads(meta)
            except:
                pass
        
        # 尝试从 metadata 中获取 doc_name
        doc_name = "未知文档"
        # 优先从 metadata 字典取
        if isinstance(meta, dict):
             doc_name = meta.get("doc_name") or meta.get("file_name") or meta.get("title") or doc_name
        
        # 如果 metadata 里没有，或者不是 dict，尝试从一级字段取 (兼容旧数据)
        if doc_name == "未知文档":
             doc_name = r.get("file_name") or r.get("title") or r.get("source") or r.get("doc_name") or doc_name

        # [Modified] 如果仍然是“未知文档”，尝试从 KnowledgeBase 名称获取
        # 用户需求：文档名称从数据库获取 (即 KB 的中文名称)
        # 如果 kb_map 提供了映射，优先使用 KB 名称作为 doc_name（如果原始 doc_name 未知）
        # 或者，也许用户希望即使有 doc_name 也显示 KB 名称？
        # 通常 doc_name 指的是文件名。如果这是片段列表，显示所属 KB 名称更有意义？
        # 假设：如果 doc_name 是默认值，就用 KB 名称。
        
        if doc_name == "未知文档" and kb_map and col_name in kb_map:
             doc_name = kb_map[col_name]
        
        # 进一步：如果用户说“文档的获取从数据库中获取名称”，可能意味着
        # 这一列应该始终显示 KB 名称？
        # 暂且保留原逻辑：优先显示文件名，如果没有文件名，显示 KB 名称。
        
        # [Fix] 优先从 metadata 中提取 parent_id 和 document_id，如果一级字段没有值
        meta_dict = meta if isinstance(meta, dict) else {}
        
        parent_id = r.get("parent_id")
        
        if not parent_id and "parent_id" in meta_dict:
            parent_id = meta_dict["parent_id"]
            
        document_id = r.get("document_id")
        if not document_id and "document_id" in meta_dict:
            document_id = meta_dict["document_id"]

        # [Modified] 从数据库获取文档名称 (如果 metadata 中没有)
        # 用户的需求是 "文档的获取从数据库中获取名称"
        # 这意味着 doc_name 应该优先查数据库 KnowledgeBase 表的 name？
        # 或者是查上传文档记录表？但目前没有独立的文档记录表，只有 metadata。
        # 如果用户是指知识库名称，那就是 col_name 对应的 DB 记录。
        # 如果用户是指 file_name，那还是 metadata。
        
        # 假设用户指的是：知识片段所属的“知识库名称”应该显示为 DB 中的中文名，而不是 collection_name。
        # 目前返回的 collection_name 是英文表名。doc_name 是文件名。
        # 如果 doc_name 是“未知文档”，我们尝试用 KnowledgeBase 的 name 填充？
        # 或者我们增加一个字段 kb_name？
        
        # 另一种理解：doc_name 字段应该取自 KnowledgeBase.name
        # 让我们尝试查询 DB 获取 KB name。
        # 由于这里是同步方法，无法 await db.execute。
        # 只能引入同步 session 或者缓存。
        # 但 _format_snippet 是在循环中调用的，查 DB 会很慢。
        
        # 鉴于上下文，之前的 user_input 是 "文档的获取从数据库中获取名称"。
        # 结合之前的 "检索引擎部分的文档会不存在"，可能用户是指 doc_name 显示不对。
        # 如果 doc_name 为空，我们尝试用 collection_name 对应的 KnowledgeBase.name。
        
        # 实际上，我们可以利用传入的 col_name，在 get_list 的外层先构建一个 col_name -> kb_name 的映射。
        # 但 _format_snippet 只接收 col_name。
        
        # 我们修改 _format_snippet 签名，允许传入 kb_map (Optional)
        # 但这需要修改调用处。
        
        # 让我们先在 doc_name 赋值逻辑里，如果最终还是 "未知文档"，就显示 col_name。
        if doc_name == "未知文档":
             doc_name = col_name
             
        # 但这只是英文名。
        # 如果能拿到 KB 的中文名最好。
        # 由于无法在此处查 DB。我们假设调用方会处理，或者接受现状。
        
        # 等等，SnippetService.get_list 是同步方法吗？不，是 def get_list，但它是在 FastAPI 路由里调用的。
        # 代码里没有 async def get_list。
        # 但 KnowledgeBaseService 是 async。
        
        # 如果我们想在 format 时获取 DB 数据，必须有 session。
        
        # 让我们看看 get_list 的调用上下文。
        # 并没有传入 session。
        
        # 也许用户的意思是：前端显示的“文档名称”这一列，应该显示 KnowledgeBase 的名称？
        # 如果是这样，前端展示时可以使用 collection_name 关联到的 KnowledgeBase.name。
        # 前端 get_list 接口返回的数据里包含 collection_name。
        # 后端 get_list 并没有返回 KB 的中文名。
        
        # 让我们修改 get_list，先获取所有 KB 的映射，然后传给 _format_snippet。
        # 但 get_list 目前没有 db session。
        
        # 方案：
        # 1. 在 get_list 中引入 get_db_connection (同步或异步)
        # 2. 查询所有 KnowledgeBase，构建 {collection_name: name} 映射
        # 3. 传给 _format_snippet
        
        # 鉴于 get_list 是同步方法，我们只能用同步方式查，或者 run_async。
        # 或者，我们假设 "doc_name" 是指 metadata 里的 "doc_name"。
        # 如果用户说“从数据库获取”，那肯定是指 SQL 数据库。
        # 而 SQL 数据库里只有 KnowledgeBase 表。
        # 所以 doc_name = KnowledgeBase.name 是最合理的解释。
        
        # 我们来修改 get_list。
        pass

        meta_info = f"ParentID: {parent_id or '-'}"
        
        # 时间处理
        created_at = "-"
        if r.get("created_time"):
            try:
                # 假设是毫秒时间戳
                ts = int(r.get("created_time")) / 1000
                created_at = datetime.fromtimestamp(ts).strftime("%Y-%m-%d %H:%M:%S")
            except:
                pass
        
        updated_at = "-"
        if r.get("updated_time"):
            try:
                # 假设是毫秒时间戳
                ts = int(r.get("updated_time")) / 1000
                updated_at = datetime.fromtimestamp(ts).strftime("%Y-%m-%d %H:%M:%S")
            except:
                pass
        
        return {
            "id": str(id_val),
            "collection_name": col_name,
            "doc_name": doc_name,
            "code": f"SNIP-{id_val}",
            "content": content,
            "char_count": len(content) if content else 0,
            "meta_info": meta_info,
            "metadata": meta, # 透传完整元数据
            "document_id": document_id or "",
            "parent_id": parent_id or "",
            "tag_list": r.get("tag_list", ""), # 返回 tag_list
            "status": "normal",
            "created_at": created_at,
            "updated_at": updated_at
        }

    async def get_by_id(self, db: AsyncSession, kb: str, id: str) -> Optional[Dict]:
        """根据 ID 获取单个片段详情"""
        # [Fix] 参数校验
        if not id or not str(id).strip():
            return None
            
        if not milvus_service.has_collection(kb):
            return None
            
        try:
            # 获取集合字段信息
            
            desc = milvus_service.client.describe_collection(kb)
            fields = desc.get('fields', [])
            field_names = [f['name'] for f in fields]

            snippet_data = None
            
            # 1. 尝试通过 document_id 查询 (常规逻辑)
            clean_id = id
            if str(id).startswith("SNIP-"):
                clean_id = str(id).replace("SNIP-", "")
                
            # ... (Existing query logic to get snippet_data from child collection) ...
            # 为了复用逻辑，这里简化处理，直接复制之前的查询逻辑，但要注意最终赋值给 snippet_data
            
            # 优先查 document_id (自定义ID)
            if "document_id" in field_names:
                doc_id_is_int = False
                for f in fields:
                    if f['name'] == 'document_id':
                         if f.get('type') == 5: # INT64
                             doc_id_is_int = True
                         break
                
                doc_expr = ""
                if doc_id_is_int:
                    if str(clean_id).isdigit():
                        doc_expr = f'document_id == {clean_id}'
                else:
                    doc_expr = f'document_id == "{clean_id}"'

                if doc_expr:
                    res = milvus_service.client.query(kb, filter=doc_expr, output_fields=["*"], limit=1)
                    if res:
                        snippet_data = self._format_snippet(res[0], kb)

            if not snippet_data and "metadata" in field_names:
                try:
                    meta_expr = f'metadata["parent_id"] == "{clean_id}"'
                    res = milvus_service.client.query(kb, filter=meta_expr, output_fields=["*"], limit=1)
                    if res:
                        snippet_data = self._format_snippet(res[0], kb)
                except Exception:
                    pass
            
            if not snippet_data:
                # 尝试通过 pk 查询
                pk_field = "pk"
                is_int = True
                for f in fields:
                    if f.get('primary_key'):
                        pk_field = f.get('name')
                        if f.get('type') == 5:
                            is_int = True
                        else:
                            is_int = False
                        break
                
                expr = ""
                if is_int:
                    if str(clean_id).isdigit():
                        expr = f"{pk_field} == {clean_id}"
                else:
                    expr = f'{pk_field} == "{clean_id}"'
                    
                if expr:
                    res = milvus_service.client.query(kb, filter=expr, output_fields=["*"], limit=1)
                    if res:
                        snippet_data = self._format_snippet(res[0], kb)

            if not snippet_data:
                snippet_data = {
                    "id": "",
                    "collection_name": kb,
                    "doc_name": "",
                    "code": "",
                    "content": "",
                    "char_count": 0,
                    "meta_info": "",
                    "metadata": {},
                    "document_id": "",
                    "parent_id": "",
                    "tag_list": "",
                    "status": "normal",
                    "created_at": "-",
                    "updated_at": "-"
                }

            # [New Feature] 获取父段内容
            # 逻辑：根据当前子表(kb) -> 查 KnowledgeBase 表找到对应的父表 -> 用 parent_id 查父表内容
            parent_id = snippet_data.get("parent_id") or clean_id
            
            if parent_id:
                try:
                    # 1. 查找 KnowledgeBase 记录
                    from sqlalchemy import select, or_
                    from app.sample.models.knowledge_base import KnowledgeBase
                    
                    # 假设当前 kb 是 collection_name_children (子表)，我们要找 collection_name_parent (父表)
                    stmt = select(KnowledgeBase).where(
                        or_(
                            KnowledgeBase.collection_name_children == kb,
                            KnowledgeBase.collection_name_parent == kb # 兼容处理
                        )
                    )
                    res = await db.execute(stmt)
                    kb_record = res.scalars().first()
                    
                    if kb_record and kb_record.collection_name_parent:
                        parent_kb = kb_record.collection_name_parent
                        
                        # 2. 在父表中查询 parent_id 相同的父段（可能有多个切片）
                        if milvus_service.has_collection(parent_kb):
                            logger.debug("Parent collection describe succeeded")
                            p_desc = milvus_service.client.describe_collection(parent_kb)
                            p_fields = p_desc.get('fields', [])
                            p_field_names = [f['name'] for f in p_fields]
                            
                            clean_pid = str(parent_id)
                            if clean_pid.startswith("SNIP-"):
                                clean_pid = clean_pid.replace("SNIP-", "")

                            parent_segments = []
                            p_expr = ""
                            if "parent_id" in p_field_names:
                                p_is_int = False
                                for f in p_fields:
                                    if f['name'] == 'parent_id':
                                        if f.get('type') == 5:
                                            p_is_int = True
                                        break
                                if p_is_int and clean_pid.isdigit():
                                    p_expr = f'parent_id == {clean_pid}'
                                else:
                                    p_expr = f'parent_id == "{clean_pid}"'
                            elif "metadata" in p_field_names:
                                p_expr = f'metadata["parent_id"] == "{clean_pid}"'

                            if p_expr:
                                p_res = milvus_service.client.query(
                                    collection_name=parent_kb,
                                    filter=p_expr,
                                    output_fields=["pk", "id", "text", "content", "page_content", "index", "document_id", "parent_id", "metadata", "created_time"],
                                    limit=200
                                )
                                if p_res:
                                    for r in p_res:
                                        p_content = r.get("text") or r.get("content") or r.get("page_content") or ""
                                        parent_segments.append({
                                            "id": str(r.get("pk") or r.get("id") or ""),
                                            "index": r.get("index"),
                                            "document_id": str(r.get("document_id") or ""),
                                            "parent_id": str(r.get("parent_id") or clean_pid),
                                            "metadata": r.get("metadata") or {},
                                            "created_time": r.get("created_time"),
                                            "content": p_content
                                        })

                            if parent_segments:
                                def _sort_key(x):
                                    idx = x.get("index")
                                    try:
                                        idx = int(idx)
                                    except Exception:
                                        idx = 0
                                    ct = x.get("created_time")
                                    try:
                                        ct = int(ct)
                                    except Exception:
                                        ct = 0
                                    pid = x.get("id") or ""
                                    return (idx, ct, pid)

                                parent_segments.sort(key=_sort_key)
                                snippet_data["parent_segments"] = parent_segments
                                snippet_data["parent_content"] = parent_segments[0].get("content") or ""
                                
                except Exception as e:
                    logger.exception("Failed to fetch parent content")

            return snippet_data
            
        except Exception as e:
            logger.exception("Get snippet detail error, id=%s, kb=%s", id, kb)
            return None

    def export_snippets(self, kb: Optional[str] = None, keyword: Optional[str] = None) -> Any:
        """导出知识片段 (生成器)"""
        
        # 1. 确定要查询的目标集合列表
        target_collections = []
        if kb:
            target_collections = [kb]
        else:
            target_collections = milvus_service.client.list_collections()
        
        for col_name in target_collections:
            if not milvus_service.has_collection(col_name):
                continue
                
            try:
                # 获取该集合总数
                stats = milvus_service.client.get_collection_stats(col_name)
                col_count = int(stats.get("row_count", 0)) if isinstance(stats, dict) else 0
                
                if col_count == 0:
                    continue

                output_fields = ["*"]
                expr = ""
                
                if keyword:
                    desc = milvus_service.client.describe_collection(col_name)
                    existing_fields = [f['name'] for f in desc.get('fields', [])]
                    if 'text' in existing_fields:
                        expr = f'text like "%{keyword}%"'
                    else:
                        continue 
                
                # 分批获取所有数据
                batch_size = 1000
                offset = 0
                
                while True:
                    res = milvus_service.client.query(
                        collection_name=col_name,
                        filter=expr,
                        output_fields=output_fields,
                        limit=batch_size,
                        offset=offset
                    )
                    
                    if not res:
                        break
                        
                    for r in res:
                        yield self._format_snippet(r, col_name)
                        
                    offset += len(res)
                    if len(res) < batch_size:
                        break
                        
            except Exception as e:
                logger.exception("Collection %s export error", col_name)
                continue

    def generate_csv_stream(self, kb: Optional[str] = None, keyword: Optional[str] = None):
        """生成CSV流"""
        output = io.StringIO()
        fieldnames = ["id", "collection_name", "doc_name", "content", "meta_info", "created_at", "status"]
        writer = csv.DictWriter(output, fieldnames=fieldnames)
        
        # 写入表头
        writer.writeheader()
        yield output.getvalue()
        output.seek(0)
        output.truncate(0)
        
        for item in self.export_snippets(kb, keyword):
            # 过滤掉不在 fieldnames 中的字段
            row = {k: item.get(k, "") for k in fieldnames}
            writer.writerow(row)
            yield output.getvalue()
            output.seek(0)
            output.truncate(0)

snippet_service = SnippetService()
