"""
知识库业务逻辑服务
"""
from math import ceil
from typing import List, Optional, Tuple, Dict, Any
import logging
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy import select, func, or_, delete as sql_delete, update as sql_update
from datetime import datetime
import uuid
import asyncio

from app.sample.models.knowledge_base import KnowledgeBase
from app.sample.models.metadata import SampleMetadata
from app.sample.models.custom_schema import CustomSchema
from app.sample.models.base_info import DocumentMain
from app.sample.schemas.knowledge_base import (
    KnowledgeBaseCreate, 
    KnowledgeBaseUpdate,
    KnowledgeBaseResponse
)
from app.services.milvus_service import milvus_service
from app.schemas.base import PaginationSchema

logger = logging.getLogger(__name__)

class KnowledgeBaseService:
    
    async def _get_collection_row_count(self, collection_name: str) -> int:
        """获取集合行数（优先尝试 count(*) 以获取实时准确值）"""
        try:
            # 确保集合已加载
            state = milvus_service.get_collection_state(collection_name)
            if state != "Loaded":
                # [Fix] 检查是否存在索引，如果没有索引则不能加载，避免抛出 index not found 错误
                try:
                    indexes = milvus_service.client.list_indexes(collection_name)
                    if not indexes:
                         # 无索引无法加载，直接跳过，进入 Fallback 使用 stats
                         raise Exception("Collection has no index, cannot load")
                except Exception:
                    # list_indexes 失败也视为无法加载
                    raise Exception("Failed to check indexes or no index")

                milvus_service.set_collection_state(collection_name, "load")
            
            # 尝试使用 count(*) 获取准确的实时数量
            # 过滤掉已标记删除的数据 (is_deleted == false)
            try:
                # [Fix] 动态检测 is_deleted 类型，自适应 Int64 或 Bool
                desc = milvus_service.client.describe_collection(collection_name)
                fields = desc.get('fields', [])
                is_del_field = next((f for f in fields if f['name'] == 'is_deleted'), None)
                
                filter_expr = ""
                if is_del_field:
                     if is_del_field.get('type') == 1: # Bool=1
                         filter_expr = "is_deleted == false"
                     else:
                         filter_expr = "is_deleted == 0" # 默认 Int
                else:
                    # 字段不存在，使用恒真条件
                    pk_field = "id"
                    for f in fields:
                        if f.get('is_primary') or f.get('primary_key'):
                            pk_field = f.get('name')
                            break
                    # 使用 PK 过滤
                    filter_expr = f"{pk_field} >= 0" # 假设是 Int PK，如果是 VarChar 需调整
                    # 检测 PK 类型
                    pk_type = next((f.get('type') for f in fields if f.get('name') == pk_field), None)
                    # Milvus DataType: INT64=5, VARCHAR=21
                    if pk_type == 21 or pk_type == 'VarChar':
                         filter_expr = f'{pk_field} != ""'

                # 注意：count(*) 聚合查询不支持分页参数 (limit/offset)，也不能用于获取实体
                res = milvus_service.client.query(collection_name, filter=filter_expr, output_fields=["count(*)"])
                if res and isinstance(res, list) and "count(*)" in res[0]:
                    return int(res[0]["count(*)"])
            except Exception as e:
                logger.warning("Query count with filter error for %s: %s", collection_name, e)
                # 再次尝试不过滤 (使用恒真表达式)
                if milvus_service.get_collection_state(collection_name) == "Loaded":
                     # 获取 PK 字段名
                     desc = milvus_service.client.describe_collection(collection_name)
                     fields = desc.get('fields', [])
                     pk_field = "id"
                     pk_type = 5
                     for f in fields:
                         if f.get('is_primary') or f.get('primary_key'):
                             pk_field = f.get('name')
                             pk_type = f.get('type')
                             break
                     
                     filter_expr = f"{pk_field} >= 0"
                     if pk_type == 21 or pk_type == 'VarChar':
                          filter_expr = f'{pk_field} != ""'

                     res = milvus_service.client.query(collection_name, filter=filter_expr, output_fields=["count(*)"])
                     if res and isinstance(res, list) and "count(*)" in res[0]:
                        return int(res[0]["count(*)"])
        except Exception:
            logger.exception("Get collection row count error for %s", collection_name)
            
        # Fallback: 使用 get_collection_stats (可能包含已删除未 Compaction 的数据)
        try:
            stats = milvus_service.client.get_collection_stats(collection_name)
            return int(stats.get("row_count", 0))
        except Exception:
            return 0

    async def _infer_and_save_metadata(self, db: AsyncSession, kb: KnowledgeBase) -> None:
        """
        [Internal] 从 Milvus 数据中推断元数据并保存到 DB
        仅当 DB 中没有定义元数据时调用
        """
        target_col = kb.collection_name_parent
        if not target_col:
            return

        try:
            # 检查是否已加载（避免不必要的错误）
            # if milvus_service.get_collection_state(target_col) != "Loaded":
            #     milvus_service.load_collection(target_col)
            
            # 采样查询 (获取前10条)
            res = []
            try:
                # 先检查 metadata 字段是否存在，避免报错 field metadata not exist
                desc = milvus_service.client.describe_collection(target_col)
                fields = [f['name'] for f in desc.get('fields', [])]
                if "metadata" not in fields:
                    return # 集合无 metadata 字段，无需推断

                # [Fix] 动态检测 is_deleted 类型，自适应 Int64 或 Bool
                is_del_field = next((f for f in fields if f['name'] == 'is_deleted'), None)
                filter_expr = "is_deleted == 0"
                if is_del_field and is_del_field.get('type') == 1: # Bool=1
                    filter_expr = "is_deleted == false"
                elif not is_del_field:
                    # 如果没有 is_deleted 字段，则不使用过滤
                    filter_expr = ""

                # 如果 filter_expr 为空，使用恒真表达式以避免空 filter 错误
                if not filter_expr:
                    pk_field = "id"
                    for f in fields:
                        if f.get('is_primary') or f.get('primary_key'):
                            pk_field = f.get('name')
                            break
                    filter_expr = f"{pk_field} >= 0"
                    pk_type = next((f.get('type') for f in fields if f.get('name') == pk_field), None)
                    if pk_type == 21 or pk_type == 'VarChar':
                         filter_expr = f'{pk_field} != ""'

                res = milvus_service.client.query(
                    collection_name=target_col,
                    filter=filter_expr,
                    output_fields=["metadata"],
                    limit=10
                )
            except Exception as e:
                try:
                    # 如果 filter 查询失败（可能不支持 is_deleted），尝试无 filter 查询
                    # 再次确认 metadata 是否存在 (防止上面的检查因某种原因失效或异常被捕获)
                    if "metadata" in fields:
                        res = milvus_service.client.query(
                            collection_name=target_col,
                            filter="",
                            output_fields=["metadata"],
                            limit=10
                        )
                except Exception:
                    pass
            
            if res:
                inferred_keys = set()
                for item in res:
                    meta = item.get("metadata") or {}
                    # Milvus 可能会返回 JSON 字符串
                    if isinstance(meta, str):
                        try:
                            import json
                            meta = json.loads(meta)
                        except:
                            meta = {}
                            
                    if isinstance(meta, dict):
                        inferred_keys.update(meta.keys())
                
                # 过滤掉默认字段
                ignore_keys = {"doc_name", "file_name", "title", "source", "chunk_id"}
                inferred_keys = inferred_keys - ignore_keys
                
                if inferred_keys:
                    # 自动生成并保存到 DB
                    for key in inferred_keys:
                        new_metadata = SampleMetadata(
                            id=str(uuid.uuid4()),
                            knowledge_base_id=kb.id,
                            field_zh_name=key, # 默认用英文名
                            field_en_name=key,
                            field_type="text", # 默认推断为 text
                            remark="Auto inferred from Milvus data"
                        )
                        db.add(new_metadata)
                    
                    # 注意：调用方负责 commit，这里不 commit 以支持批量事务
                    logger.info("Auto inferred metadata for %s: %s", target_col, inferred_keys)
        except Exception as e:
            logger.exception("Failed to infer metadata for %s", target_col)

    async def get_list(
        self, 
        db: AsyncSession,
        page: int = 1,
        page_size: int = 10,
        keyword: Optional[str] = None,
        status: Optional[str] = None
    ) -> Tuple[List[KnowledgeBase], PaginationSchema]:
        """获取知识库列表"""
        
        # --- 同步 Milvus 数据 (简化版：仅更新现有KB的计数和状态) ---
        try:
            # 1. 获取 Milvus 所有集合
            milvus_names = milvus_service.client.list_collections()
            
            # 2. 获取 DB 中已有的集合
            result = await db.execute(select(KnowledgeBase).where(KnowledgeBase.is_deleted == 0))
            existing_kbs = result.scalars().all()
            
            # 3. 更新现有KB的统计
            has_changes = False
            for kb in existing_kbs:
                total_count = 0
                
                # 统计 collection_name_parent
                if kb.collection_name_parent and kb.collection_name_parent in milvus_names:
                    total_count += await self._get_collection_row_count(kb.collection_name_parent)
                    
                # 统计 collection_name_children
                if kb.collection_name_children and kb.collection_name_children in milvus_names:
                    total_count += await self._get_collection_row_count(kb.collection_name_children)
                    
                if kb.document_count != total_count:
                    kb.document_count = total_count
                    has_changes = True

            if has_changes:
                await db.commit()

        except Exception as e:
            logger.exception("Sync Milvus collections failed")
        # ----------------------

        # 查询未删除的 KB
        query = select(KnowledgeBase).where(KnowledgeBase.is_deleted == 0)
        
        if keyword:
            query = query.where(or_(
                KnowledgeBase.name.like(f"%{keyword}%"),
                KnowledgeBase.collection_name_parent.like(f"%{keyword}%"),
                KnowledgeBase.collection_name_children.like(f"%{keyword}%")
            ))
        
        if status:
            query = query.where(KnowledgeBase.status == status)

        query = query.order_by(KnowledgeBase.created_time.desc())
        
        # 分页
        total = await db.scalar(select(func.count()).select_from(query.subquery()))
        
        # 计算总页数
        total_pages = ceil(total / page_size) if page_size else 0
        
        # 获取当前页数据
        result = await db.execute(query.offset((page - 1) * page_size).limit(page_size))
        items = result.scalars().all()

        # 设置 is_synced (辅助字段，不存库)
        milvus_names_set = set(milvus_service.client.list_collections())
        for item in items:
            # 父集合可能为空（未勾选创建父集合），此时只要子集合存在也算已同步
            c1_ok = True if not getattr(item, "has_parent_collection", 1) else item.collection_name_parent in milvus_names_set
            c2_ok = True
            if item.collection_name_children:
                c2_ok = item.collection_name_children in milvus_names_set
            
            item.is_synced = c1_ok and c2_ok
        
        meta = PaginationSchema(
            page=page,
            page_size=page_size,
            total=total,
            total_pages=total_pages,
        )
        
        return items, meta

    async def create(self, db: AsyncSession, payload: KnowledgeBaseCreate) -> KnowledgeBase:
        """创建新知识库"""
        name = (payload.name or "").strip()
        if not name:
            raise ValueError("请输入知识库名称")

        parent_name = (payload.collection_name_parent or "").strip() or None
        child_name = (payload.collection_name_children or "").strip()

        if not child_name:
            raise ValueError("请输入子集合名称")

        # 检查知识库名称是否重名
        exists_name = await db.execute(select(KnowledgeBase).where(
            KnowledgeBase.name == name,
            KnowledgeBase.is_deleted == 0
        ))
        if exists_name.scalars().first():
            raise ValueError(f"知识库名称 {name} 已存在")

        # 1. 检查 DB 是否已存在
        # 检查父子集合名称不能相同
        if parent_name and parent_name == child_name:
            raise ValueError("父集合名称和子集合名称不能相同")

        # 检查 collection_name_parent (可选)
        if parent_name:
            exists1 = await db.execute(select(KnowledgeBase).where(
                or_(
                    KnowledgeBase.collection_name_parent == parent_name,
                    KnowledgeBase.collection_name_children == parent_name
                ),
                KnowledgeBase.is_deleted == 0
            ))
            if exists1.scalars().first():
                raise ValueError(f"父集合名称 {parent_name} 已存在")
            
        # 检查 collection_name_children
        exists2 = await db.execute(select(KnowledgeBase).where(
            or_(
                KnowledgeBase.collection_name_parent == child_name,
                KnowledgeBase.collection_name_children == child_name
            ),
            KnowledgeBase.is_deleted == 0
        ))
        if exists2.scalars().first():
            raise ValueError(f"子集合名称 {child_name} 已存在")

        try:
            # 3. 创建 DB 记录
            now = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
            new_kb = KnowledgeBase(
                id=str(uuid.uuid4()),
                name=name,
                collection_name_parent=parent_name,
                collection_name_children=child_name,
                has_parent_collection=1 if parent_name else 0,
                description=payload.description,
                # 默认创建为禁用状态，待同步成功后再启用
                status="disabled",
                created_by="admin",
                updated_by="admin",
                created_time=now,
                updated_time=now,
                document_count=0
            )
            db.add(new_kb)
            
            # 4. 保存元数据定义 (如果有)
            if payload.metadata_fields:
                for field in payload.metadata_fields:
                    new_metadata = SampleMetadata(
                        id=str(uuid.uuid4()),
                        knowledge_base_id=new_kb.id,
                        field_zh_name=field.field_zh_name,
                        field_en_name=field.field_en_name,
                        field_type=field.field_type,
                        remark=field.remark
                    )
                    db.add(new_metadata)

            await db.commit()
            await db.refresh(new_kb)

            return new_kb
        except Exception as e:
            await db.rollback()
            raise e

    async def update(self, db: AsyncSession, id: str, payload: KnowledgeBaseUpdate) -> KnowledgeBase:
        """更新知识库信息"""
        result = await db.execute(select(KnowledgeBase).where(KnowledgeBase.id == id, KnowledgeBase.is_deleted == 0))
        kb = result.scalars().first()
        
        if not kb:
            raise ValueError("知识库不存在")

        try:
            if payload.name is not None:
                new_name = (payload.name or "").strip()
                if not new_name:
                    raise ValueError("请输入知识库名称")
                if new_name != kb.name:
                    exists_name = await db.execute(select(KnowledgeBase).where(
                        KnowledgeBase.name == new_name,
                        KnowledgeBase.is_deleted == 0,
                        KnowledgeBase.id != id
                    ))
                    if exists_name.scalars().first():
                        raise ValueError(f"知识库名称 {new_name} 已存在")
                    kb.name = new_name

            if payload.description is not None:
                kb.description = payload.description
                # 同步更新 Milvus 描述 (如果 Milvus 中存在该集合)
                if kb.collection_name_parent and milvus_service.has_collection(kb.collection_name_parent):
                    milvus_service.update_collection_description(kb.collection_name_parent, payload.description)
                if kb.collection_name_children and milvus_service.has_collection(kb.collection_name_children):
                    milvus_service.update_collection_description(kb.collection_name_children, payload.description)

            if payload.status is not None:
                kb.status = payload.status
            
            # 更新元数据字段 (Metadata Fields)
            if payload.metadata_fields is not None:
                # 1. 删除旧的元数据字段
                await db.execute(sql_delete(SampleMetadata).where(SampleMetadata.knowledge_base_id == id))
                
                # 2. 插入新的元数据字段
                for field in payload.metadata_fields:
                    new_metadata = SampleMetadata(
                        id=str(uuid.uuid4()),
                        knowledge_base_id=kb.id,
                        field_zh_name=field.field_zh_name,
                        field_en_name=field.field_en_name,
                        field_type=field.field_type,
                        remark=field.remark
                    )
                    db.add(new_metadata)
            
            kb.updated_by = "admin" # 暂时硬编码
            kb.updated_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")

            await db.commit()
            await db.refresh(kb)
            
            return kb
        except Exception as e:
            await db.rollback()
            raise e

    async def update_status(self, db: AsyncSession, id: str, status: str) -> KnowledgeBase:
        """更新知识库状态（启用/禁用）"""
        result = await db.execute(select(KnowledgeBase).where(KnowledgeBase.id == id, KnowledgeBase.is_deleted == 0))
        kb = result.scalars().first()
        
        if not kb:
            raise ValueError("知识库不存在")
        
        try:
            kb.status = status
            kb.updated_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
            
            # 同步操作 Milvus Load/Release
            targets = []
            if kb.collection_name_parent: targets.append(kb.collection_name_parent)
            if kb.collection_name_children: targets.append(kb.collection_name_children)
            
            for col in targets:
                if milvus_service.has_collection(col):
                    if status == "normal":
                        milvus_service.set_collection_state(col, "load")
                    elif status == "disabled":
                        milvus_service.set_collection_state(col, "release")
                
            await db.commit()
            await db.refresh(kb)
            return kb
        except Exception as e:
            await db.rollback()
            raise e

    async def delete(self, db: AsyncSession, id: str) -> None:
        """删除知识库"""
        result = await db.execute(select(KnowledgeBase).where(KnowledgeBase.id == id))
        kb = result.scalars().first()
        
        if not kb:
            raise ValueError("知识库不存在")

        # 检查文档数量
        # 只要任意一个集合有文档，就阻止删除
        current_count = kb.document_count
        
        # 尝试获取最新计数
        try:
            real_count = 0
            if kb.collection_name_parent and milvus_service.has_collection(kb.collection_name_parent):
                real_count += await self._get_collection_row_count(kb.collection_name_parent)
            if kb.collection_name_children and milvus_service.has_collection(kb.collection_name_children):
                real_count += await self._get_collection_row_count(kb.collection_name_children)
            current_count = real_count
        except:
            pass
        
        if current_count > 0:
            raise ValueError(f"知识库中仍有 {current_count} 条文档，请先清空文档后再删除")

        try:
            # 1. 删除 Milvus 集合
            targets = []
            if kb.collection_name_parent: targets.append(kb.collection_name_parent)
            if kb.collection_name_children: targets.append(kb.collection_name_children)
            
            for col in targets:
                try:
                    if milvus_service.has_collection(col):
                        milvus_service.drop_collection(col)
                except Exception as milvus_err:
                    logger.warning("Ignore Milvus error during delete %s: %s", col, milvus_err)
            
            # 2. 解除文档关联 (将 kb_id 置空，状态改为未入库)
            await db.execute(
                sql_update(DocumentMain)
                .where(DocumentMain.kb_id == id)
                .values(kb_id=None, whether_to_enter=0)
            )
            
            # 3. 删除关联的元数据 (硬删除)
            await db.execute(sql_delete(SampleMetadata).where(SampleMetadata.knowledge_base_id == id))
            
            # 4. 删除关联的自定义Schema (硬删除)
            await db.execute(sql_delete(CustomSchema).where(CustomSchema.knowledge_base_id == id))

            # 5. 硬删除 KnowledgeBase 记录
            await db.execute(sql_delete(KnowledgeBase).where(KnowledgeBase.id == id))

            await db.commit()
        except Exception as e:
            await db.rollback()
            raise e

    async def sync_to_milvus(self, db: AsyncSession, id: str) -> KnowledgeBase:
        """同步知识库到Milvus"""
        result = await db.execute(select(KnowledgeBase).where(KnowledgeBase.id == id, KnowledgeBase.is_deleted == 0))
        kb = result.scalars().first()
        
        if not kb:
            raise ValueError("知识库不存在")
            
        # 使用固定的 Schema 定义
        fields = [
            {"name": "pk", "type": "INT64", "is_primary": True, "description": "主键"},
            {"name": "text", "type": "VARCHAR", "max_length": 65535, "description": "内容"},
            {"name": "dense", "type": "FLOAT_VECTOR", "description": "向量列"},
            {"name": "sparse", "type": "BM25", "description": "内容的BM25关键字检索"},
            {"name": "document_id", "type": "VARCHAR", "max_length": 128, "description": "样本中心上传文档ID"},
            {"name": "parent_id", "type": "VARCHAR", "max_length": 128, "description": "父段ID"},
            {"name": "index", "type": "INT64", "description": "索引序号"},
            {"name": "tag_list", "type": "VARCHAR", "max_length": 2048, "description": "标签"},
            {"name": "permission", "type": "JSON", "description": "权限"},
            {"name": "metadata", "type": "JSON", "description": "元数据"},
            {"name": "is_deleted", "type": "BOOL", "description": "删除标志"},
            {"name": "created_by", "type": "VARCHAR", "max_length": 128, "description": "创建人"},
            {"name": "created_time", "type": "INT64", "description": "创建时间"},
            {"name": "updated_by", "type": "VARCHAR", "max_length": 128, "description": "修改人"},
            {"name": "updated_time", "type": "INT64", "description": "修改时间"}
        ]
        
        try:
            # 创建 collection1
            if kb.collection_name_parent and not milvus_service.has_collection(kb.collection_name_parent):
                milvus_service.create_collection(
                    name=kb.collection_name_parent,
                    dimension=milvus_service.DENSE_DIM, 
                    description=kb.description or "",
                    fields=fields
                )
            
            # 创建 collection2
            if kb.collection_name_children and not milvus_service.has_collection(kb.collection_name_children):
                milvus_service.create_collection(
                    name=kb.collection_name_children,
                    dimension=milvus_service.DENSE_DIM, 
                    description=kb.description or "",
                    fields=fields
                )

            # 同步完成后将知识库状态置为 normal，并触发 load
            kb.status = "normal"
            kb.updated_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")

            targets = []
            if kb.collection_name_parent: targets.append(kb.collection_name_parent)
            if kb.collection_name_children: targets.append(kb.collection_name_children)
            for col in targets:
                if milvus_service.has_collection(col):
                    milvus_service.set_collection_state(col, "load")

            await db.commit()
            await db.refresh(kb)
            return kb
        except Exception as e:
            raise e

    async def get_metadata_and_schema(self, db: AsyncSession, kb_id: str) -> Dict[str, List[dict]]:
        """获取知识库的元数据字段列表 (如果 DB 中没有定义，尝试从 Milvus 数据推断)"""
        # 检查知识库是否存在
        result = await db.execute(select(KnowledgeBase).where(KnowledgeBase.id == kb_id, KnowledgeBase.is_deleted == 0))
        kb = result.scalars().first()
        if not kb:
            raise ValueError("知识库不存在")

        # 查询元数据表
        meta_query = select(SampleMetadata).where(SampleMetadata.knowledge_base_id == kb_id)
        meta_result = await db.execute(meta_query)
        metadata_fields = [f.to_dict() for f in meta_result.scalars().all()]
        
        # 自动推断逻辑：如果 DB 中没有定义元数据，且 Milvus 中有数据，尝试推断 (优先推断 collection_name_parent)
        target_col = kb.collection_name_parent
        if not metadata_fields and target_col and milvus_service.has_collection(target_col):
            try:
                # 采样查询 (获取前10条)
                try:
                    res = milvus_service.client.query(
                        collection_name=target_col,
                        filter="is_deleted == false",
                        output_fields=["metadata"],
                        limit=10
                    )
                except Exception as e:
                    # 如果 filter 查询失败（可能不支持 is_deleted），尝试无 filter 查询
                    res = milvus_service.client.query(
                        collection_name=target_col,
                        filter="",
                        output_fields=["metadata"],
                        limit=10
                    )
                
                if res:
                    inferred_keys = set()
                    for item in res:
                        meta = item.get("metadata") or {}
                        # Milvus 可能会返回 JSON 字符串，尝试解析
                        if isinstance(meta, str):
                            try:
                                import json
                                meta = json.loads(meta)
                            except:
                                meta = {}
                                
                        if isinstance(meta, dict):
                            inferred_keys.update(meta.keys())
                    
                    # 过滤掉一些默认字段，避免干扰
                    ignore_keys = {"doc_name", "file_name", "title", "source", "chunk_id"}
                    inferred_keys = inferred_keys - ignore_keys
                    
                    if inferred_keys:
                        # 自动生成并保存到 DB
                        new_fields = []
                        for key in inferred_keys:
                            new_metadata = SampleMetadata(
                                id=str(uuid.uuid4()),
                                knowledge_base_id=kb.id,
                                field_zh_name=key, # 默认用英文名
                                field_en_name=key,
                                field_type="text", # 默认推断为 text
                                remark="Auto inferred from Milvus data"
                            )
                            db.add(new_metadata)
                            new_fields.append(new_metadata.to_dict())
                        
                        await db.commit()
                        metadata_fields = new_fields
                        logger.info("Auto inferred metadata for %s: %s", target_col, inferred_keys)
            except Exception as e:
                logger.exception("Failed to infer metadata for %s", target_col)
                # 推断失败不影响正常返回
        
        # 返回空的 custom_schemas，因为现在是固定 Schema
        return {
            "metadata_fields": metadata_fields,
            "custom_schemas": [] 
        }

    async def get_metadata_fields(self, db: AsyncSession, kb_id: str) -> List[dict]:
        """获取知识库的元数据字段列表"""
        # 检查知识库是否存在
        result = await db.execute(select(KnowledgeBase).where(KnowledgeBase.id == kb_id, KnowledgeBase.is_deleted == 0))
        kb = result.scalars().first()
        if not kb:
            raise ValueError("知识库不存在")

        # 查询元数据表
        query = select(SampleMetadata).where(SampleMetadata.knowledge_base_id == kb_id)
        result = await db.execute(query)
        fields = result.scalars().all()
        
        return [f.to_dict() for f in fields]

    async def update_doc_count(self, db: AsyncSession, collection_name: str) -> int:
        """根据 Milvus 实时数据更新知识库文档数量
        返回最新的行计数（int），调用方可以使用该返回值立即更新前端显示。
        """
        # 查找知识库 (匹配任意一个 collection_name)
        result = await db.execute(select(KnowledgeBase).where(
            or_(
                KnowledgeBase.collection_name_parent == collection_name,
                KnowledgeBase.collection_name_children == collection_name
            ),
            KnowledgeBase.is_deleted == 0
        ))
        kb = result.scalars().first()
        
        if kb and milvus_service.has_collection(collection_name):
            try:
                # 确保集合已加载
                state = milvus_service.get_collection_state(collection_name)
                if state != "Loaded":
                    milvus_service.set_collection_state(collection_name, "load")
                
                # 获取该集合的计数
                # 注意：这里我们应该更新总计数，所以需要获取两个集合的计数
                count1 = 0
                count2 = 0
                
                if kb.collection_name_parent:
                     count1 = await self._get_collection_row_count(kb.collection_name_parent)
                if kb.collection_name_children:
                     count2 = await self._get_collection_row_count(kb.collection_name_children)
                     
                total_count = count1 + count2
                
                # 更新数据库
                if kb.document_count != total_count:
                    kb.document_count = total_count
                    kb.updated_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
                    await db.commit()
                return total_count
            except Exception as e:
                logger.exception(f"Failed to update doc count for {collection_name}: {e}")
                return 0
        return 0

knowledge_base_service = KnowledgeBaseService()
