"""
Milvus Service：业务层（直接用 manager.client 调 Milvus 原生方法）
"""
from __future__ import annotations

import sys
import os

# 添加src目录到Python路径
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '../..'))
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '../../..'))

import logging
from typing import List, Dict, Any
from datetime import datetime

from app.base import get_milvus_manager, get_milvus_vectorstore, get_embedding_model

logger = logging.getLogger(__name__)


class MilvusService:
    def __init__(self):
        self.client = get_milvus_manager().client
        # 获取embedding model
        self.emdmodel = get_embedding_model()

    def create_collection(self, name: str, dimension: int = 768, description: str = "") -> None:
        """创建 Milvus 集合"""
        if self.client.has_collection(name):
            logger.info(f"Collection {name} already exists.")
            return
        
        # 使用简化的 create_collection API
        self.client.create_collection(
            collection_name=name,
            dimension=dimension,
            description=description,
            auto_id=True,  # 自动生成 ID
            id_type="int", # ID 类型
            metric_type="COSINE" # 默认使用余弦相似度
        )
        logger.info(f"Created collection {name} with dimension {dimension}")

    def drop_collection(self, name: str) -> None:
        """删除 Milvus 集合"""
        if self.client.has_collection(name):
            self.client.drop_collection(name)
            logger.info(f"Dropped collection {name}")

    def has_collection(self, name: str) -> bool:
        """检查集合是否存在"""
        return self.client.has_collection(name)

    def get_collection_details(self) -> List[Dict[str, Any]]:
        """
        获取所有 Collections 详细信息
        """
        details: List[Dict[str, Any]] = []

        names = self.client.list_collections()

        for name in names:
            desc = self.client.describe_collection(collection_name=name)
            stats = self.client.get_collection_stats(collection_name=name)
            load_state = self.client.get_load_state(collection_name=name)

            # ===== 时间戳转换（按你指定写法，无封装）=====
            created_time = None
            updated_time = None

            if desc.get("created_timestamp") is not None:
                ts_int = int(desc["created_timestamp"])
                physical_ms = ts_int >> 18
                created_time = datetime.fromtimestamp(physical_ms / 1000).strftime("%Y-%m-%d %H:%M:%S")

            if desc.get("update_timestamp") is not None:
                ts_int = int(desc["update_timestamp"])
                physical_ms = ts_int >> 18
                updated_time = datetime.fromtimestamp(physical_ms / 1000).strftime("%Y-%m-%d %H:%M:%S")

            # ===== 数量：不保底（要求返回结构必须有 row_count）=====
            entity_count = stats["row_count"]

            # ===== 状态：不保底（要求返回结构必须有 state）=====
            status = load_state["state"]

            details.append(
                {
                    "name": name,
                    "status": status,
                    "entity_count": entity_count,
                    "description": desc.get("description", ""),
                    "created_time": created_time,
                    "updated_time": updated_time,
                }
            )

        logger.info(f"成功获取Collections详细信息，共{len(details)}个")
        return details

    def set_collection_state(self, name: str, action: str) -> Dict[str, Any]:
        """
        改变指定 Collection 的加载状态。

        参数:
        - name: 集合名称
        - action: 操作，取值 'load' 或 'release'

        返回:
        - 包含集合名称和当前状态的字典，例如: {"name": name, "state": "Loaded"}
        """
        action_norm = (action or "").strip().lower()
        if action_norm not in {"load", "release"}:
            raise ValueError("action 必须为 'load' 或 'release'")

        # 执行加载/释放
        if action_norm == "load":
            self.client.load_collection(collection_name=name)
        else:
            self.client.release_collection(collection_name=name)

        # 返回最新状态
        load_state = self.client.get_load_state(collection_name=name)
        state = load_state.get("state") if isinstance(load_state, dict) else load_state
        result = {"name": name, "state": state, "action": action_norm}
        logger.info(f"集合 {name} 状态更新为 {state} (action={action_norm})")
        return result

    def delete_collection_if_empty(self, name: str) -> Dict[str, Any]:
        """仅当集合内容为空时删除集合，否则抛出异常"""
        stats = self.client.get_collection_stats(collection_name=name)
        row_count = stats.get("row_count") if isinstance(stats, dict) else None
        if row_count is None:
            raise ValueError("无法获取集合行数，禁止删除")
        if int(row_count) > 0:
            raise ValueError("集合内容不为空，不能删除")

        self.client.drop_collection(collection_name=name)
        logger.info(f"集合 {name} 已删除")
        return {"name": name, "deleted": True}

    def get_collection_detail(self, name: str) -> Dict[str, Any]:
        """获取单个集合的详细信息，包含schema、索引等所有desc字段"""
        desc = self.client.describe_collection(collection_name=name)
        stats = self.client.get_collection_stats(collection_name=name)
        load_state = self.client.get_load_state(collection_name=name)

        # 时间戳转换
        created_time = None
        updated_time = None

        if desc.get("created_timestamp") is not None:
            ts_int = int(desc["created_timestamp"])
            physical_ms = ts_int >> 18
            created_time = datetime.fromtimestamp(physical_ms / 1000).strftime("%Y-%m-%d %H:%M:%S")

        if desc.get("update_timestamp") is not None:
            ts_int = int(desc["update_timestamp"])
            physical_ms = ts_int >> 18
            updated_time = datetime.fromtimestamp(physical_ms / 1000).strftime("%Y-%m-%d %H:%M:%S")

        entity_count = stats.get("row_count", 0)
        status = load_state.get("state") if isinstance(load_state, dict) else load_state

        # 提取字段schema
        fields = []
        if "fields" in desc:
            for field in desc["fields"]:
                field_info = {
                    "name": field.get("name"),
                    "type": str(field.get("type")),
                    "description": field.get("description", ""),
                    "is_primary": field.get("is_primary", False),
                    "auto_id": field.get("auto_id"),
                }
                # 向量维度
                if "params" in field and "dim" in field["params"]:
                    field_info["dim"] = field["params"]["dim"]
                # 字符串长度
                if "params" in field and "max_length" in field["params"]:
                    field_info["max_length"] = field["params"]["max_length"]
                # 其他params
                if "params" in field:
                    field_info["params"] = field["params"]
                fields.append(field_info)

        # 提取索引信息
        indices = []
        if "indexes" in desc:
            for idx in desc["indexes"]:
                index_info = {
                    "field_name": idx.get("field_name"),
                    "index_name": idx.get("index_name"),
                    "index_type": idx.get("index_type"),
                    "metric_type": idx.get("metric_type"),
                    "params": idx.get("params"),
                }
                indices.append(index_info)

        detail = {
            "name": name,
            "description": desc.get("description", ""),
            "status": status,
            "entity_count": entity_count,
            "created_time": created_time,
            "updated_time": updated_time,
            "fields": fields,
            "enable_dynamic_field": desc.get("enable_dynamic_field", False),
            "consistency_level": desc.get("consistency_level"),
            "num_shards": desc.get("num_shards"),
            "num_partitions": desc.get("num_partitions"),
            "indices": indices,
            "properties": desc.get("properties"),
            "aliases": desc.get("aliases", []),
        }

        logger.info(f"成功获取集合 {name} 的详细信息")
        return detail

    
    def update_collection_description(self, name: str, description: str) -> Dict[str, Any]:
        """使用 alter_collection_properties 更新集合描述"""
        description = description or ""

        # 1. 更新集合 description（唯一修改点）
        self.client.alter_collection_properties(
            collection_name=name,
            properties={"collection.description": description},
        )

        # 2. 重新获取集合信息
        desc = self.client.describe_collection(collection_name=name)
        print(desc)
        stats = self.client.get_collection_stats(collection_name=name)
        load_state = self.client.get_load_state(collection_name=name)

        # 3. 时间戳转换（Milvus TSO -> 物理时间）
        def ts_to_str(ts):
            if ts is None:
                return None
            ts_int = int(ts)
            physical_ms = ts_int >> 18
            return datetime.fromtimestamp(physical_ms / 1000).strftime("%Y-%m-%d %H:%M:%S")

        created_time = ts_to_str(desc.get("created_timestamp"))
        updated_time = ts_to_str(desc.get("update_timestamp"))

        entity_count = stats.get("row_count") if isinstance(stats, dict) else None
        status = load_state.get("state") if isinstance(load_state, dict) else load_state

        return {
            "name": name,
            "status": status,
            "entity_count": entity_count,
            "description": desc.get("description", ""),
            "created_time": created_time,
            "updated_time": updated_time,
        }

    def hybrid_search(self, collection_name: str, query_text: str,
                     top_k: int = 3, ranker_type: str = "weighted",
                     dense_weight: float = 0.7, sparse_weight: float = 0.3):
        """
        混合搜索（参考 test_hybrid_v2.6.py 的实现）

        Args:
            param: 包含collection_name的参数字典
            query_text: 查询文本
            top_k: 返回结果数量
            ranker_type: 重排序类型 "weighted" 或 "rrf"
            dense_weight: 密集向量权重（当ranker_type="weighted"时使用）
            sparse_weight: 稀疏向量权重（当ranker_type="weighted"时使用）

        Returns:
            List[Dict]: 搜索结果列表
        """
        try:
            collection_name = collection_name

            # 获取 vectorstore 实例（包含 Milvus 和 BM25BuiltInFunction）
            vectorstore = get_milvus_vectorstore(
                collection_name=collection_name,
                consistency_level="Strong"
            )

            # 执行混合搜索 (完全按照 test_hybrid_v2.6.py 的逻辑)
            if ranker_type == "weighted":
                results = vectorstore.similarity_search(
                    query=query_text,
                    k=top_k,
                    ranker_type="weighted",
                    ranker_params={"weights": [dense_weight, sparse_weight]}
                )
            else:  # rrf
                results = vectorstore.similarity_search(
                    query=query_text,
                    k=top_k,
                    ranker_type="rrf",
                    ranker_params={"k": 60}
                )

            # 格式化结果，保持与其他搜索方法一致
            formatted_results = []
            for doc in results:
                formatted_results.append({
                    'id': doc.metadata.get('pk', 0),
                    'text_content': doc.page_content,
                    'metadata': doc.metadata,
                    'distance': 0.0,
                    'similarity': 1.0
                })

            logger.info(f"Hybrid search returned {len(formatted_results)} results")
            return formatted_results

        except Exception as e:
            logger.error(f"Error in hybrid search: {e}")
            # 回退到传统的向量搜索
            logger.info("Falling back to traditional vector search")


# 可选：单例
milvus_service = MilvusService()


if __name__ == "__main__":
    # 推荐这样跑：
    # uv run python -m src.app.services.milvus_service
    import json

    service = MilvusService()
    
    # 测试混合搜索 hybrid_search
    print("=" * 50)
    print("测试混合检索 (Hybrid Search)")
    print("=" * 50)
    
    try:
        # 示例参数，需要根据实际情况修改
        collection_name = "first_bfp_collection_status" 
        query_text = "《公路水运工程临时用电技术规程》（JTT1499-2024）状态为现行"  # 修改为实际查询内容
        
        # 测试 weighted 模式
        print("\n1. 测试 Weighted 重排序模式:")
        print(f"   集合: {collection_name}")
        print(f"   查询: {query_text}")
        print(f"   密集权重: 0.7, 稀疏权重: 0.3")
        
        results_weighted = service.hybrid_search(
            collection_name=collection_name,
            query_text=query_text,
            top_k=5,
            ranker_type="weighted",
            dense_weight=0.7,
            sparse_weight=0.3
        )
        
        print(f"\n   结果数量: {len(results_weighted)}")
        for i, result in enumerate(results_weighted, 1):
            print(f"   [{i}] ID: {result.get('id')}, Text: {result.get('text_content')[:50]}...")
        
        # 测试 RRF 模式
        print("\n2. 测试 RRF (Reciprocal Rank Fusion) 重排序模式:")
        print(f"   集合: {collection_name}")
        print(f"   查询: {query_text}")
        
        results_rrf = service.hybrid_search(
            collection_name=collection_name,
            query_text=query_text,
            top_k=5,
            ranker_type="rrf"
        )
        
        print(f"\n   结果数量: {len(results_rrf)}")
        for i, result in enumerate(results_rrf, 1):
            print(f"   [{i}] ID: {result.get('id')}, Text: {result.get('text_content')[:50]}...")
        
        print("\n✓ 混合检索测试完成")
        
    except Exception as e:
        print(f"\n✗ 混合检索测试失败: {e}")
        import traceback
        traceback.print_exc()
    
    # 也可以查看集合详情
    print("\n" + "=" * 50)
    print("获取所有集合信息:")
    print("=" * 50)
    data = service.get_collection_details()
    for item in data:
        print(json.dumps(item, ensure_ascii=False, indent=2))
