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 IntentClassificationRepository:
    collection_name: str = 'intent_classification'

    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:
        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="意图分类向量集合"
        )

        schema.add_field(field_name="id", datatype=DataType.INT64, is_primary=True, auto_id=True)
        schema.add_field(field_name="intent_category", datatype=DataType.VARCHAR, max_length=128)
        schema.add_field(field_name="content", datatype=DataType.VARCHAR, max_length=65535)
        schema.add_field(field_name="content_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="content_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_intent_data(self, data: List[Dict[str, Any]]) -> List[int]:
        try:
            created_at = int(datetime.now().timestamp())
            insert_data = []
            for item in data:
                insert_data.append({
                    "intent_category": item.get("intent_category", ""),
                    "content": item.get("content", ""),
                    "content_vector": item.get("content_vector", []),
                    "created_at": created_at
                })

            result = milvus_client.client.insert(
                collection_name=self.collection_name,
                data=insert_data
            )

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

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

    async def search_by_intent_category(self, intent_category: List[str], top_k: int = 100) -> List[Dict[str, Any]]:
        if isinstance(intent_category, str):
            intent_category = [intent_category]
        categories_str = ", ".join(f'"{c}"' for c in intent_category)
        expr = f'intent_category in [{categories_str}]'
        result = milvus_client.client.query(
            collection_name=self.collection_name,
            filter=expr,
            output_fields=["intent_category", "content", "created_at"],
            limit=top_k
        )
        return result

    async def search_by_vector(self, query_vector: List[float], top_k: int = 5) -> List[Dict[str, Any]]:
        res = milvus_client.client.search(
            collection_name=self.collection_name,
            anns_field="content_vector",
            data=[query_vector],
            limit=top_k,
            search_params={
                "metric_type": "COSINE",
                "efSearch": 300
            },
            output_fields=["intent_category", "content", "created_at"]
        )

        results = []
        for hits in res:
            for hit in hits:
                item = {
                    "intent_category": hit.entity.get("intent_category"),
                    "content": hit.entity.get("content"),
                    "created_at": hit.entity.get("created_at"),
                    "score": hit.get("distance", 0)
                }
                results.append(item)
        return results

    async def get_all_intent_categories(self) -> List[str]:
        result = milvus_client.client.query(
            collection_name=self.collection_name,
            output_fields=["intent_category"],
            limit=1000
        )
        categories = list(set([item.get("intent_category", "") for item in result]))
        return categories