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

from pymilvus import MilvusClient, DataType

from app.client.milvus_client_manager import milvus_client
from app.conf.app_config import app_config
from app.core.log import logger


class QARepository:
    collection_name: str = 'message_qa'

    def __init__(self):
        self.dim = app_config.milvus.embedding_size
        self._ensure_connection()

    def _ensure_connection(self):
        if milvus_client.client is None:
            milvus_client.init()

    def create_collection(self, dim: int = None) -> bool:
        """
        创建消息摘要集合
        
        Args:
            dim: 向量维度，默认从配置读取
        
        Returns:
            是否成功创建（已存在返回False）
        """
        if dim is not None:
            self.dim = dim

        if milvus_client.has_collection(self.collection_name):
            return False

        schema = milvus_client.client.create_schema(
            auto_id=True,
            enable_dynamic_field=True,
            description="QA向量集合"
        )

        schema.add_field(field_name="id", datatype=DataType.INT64, is_primary=True, auto_id=True)
        schema.add_field(field_name="question", datatype=DataType.VARCHAR, max_length=1024)
        schema.add_field(field_name="answer", datatype=DataType.VARCHAR, max_length=65535)
        schema.add_field(field_name="parent_id", datatype=DataType.INT64)
        schema.add_field(field_name="question_dense_vector", datatype=DataType.FLOAT_VECTOR, dim=1024)
        schema.add_field(field_name="created_at", datatype=DataType.INT64)

        index_params = milvus_client.client.prepare_index_params()
        index_params.add_index(
            field_name="question_dense_vector",
            index_type="HNSW",
            metric_type="COSINE",
            params={"M": 32, "efConstruction": 300}
        )
        milvus_client.client.create_collection(
            collection_name=self.collection_name,
            schema=schema,
            index_params=index_params
        )

        return True

    async def save_qa_to_milvus(self, qa_list: List[Dict[str, Any]]) -> List[int]:
        """
        将QA数据入库到Milvus

        Args:
            qa_pairs: QA对列表，格式 [{"question": "...", "answer": "..."}, ...]
            parent_id: 父记录ID，默认为0

        Returns:
            插入的记录ID列表
        """
        try:
            # 构建入库数据
            created_at = int(datetime.now().timestamp())
            data = []
            for qa in qa_list:
                data.append({
                    "question": qa.get("question", ""),
                    "answer": qa.get("answer", ""),
                    "parent_id": qa.get("parent_id", 0),
                    "question_dense_vector": qa.get("question_dense_vector", []),
                    "created_at": created_at
                })

            # 插入数据
            result = milvus_client.client.insert(
                collection_name=self.collection_name,
                data=data
            )

            logger.info(f"QA数据入库成功，插入 {len(data)} 条记录")
            return result.get("ids", [])

        except Exception as e:
            logger.error(f"QA数据入库失败: {str(e)}", exc_info=True)
            return []