"""Milvus collection schemas for the mutual_fund knowledge base."""
from __future__ import annotations

from pymilvus import AsyncMilvusClient, DataType

from config.database.milvus import client as configured_milvus_client


KNOWLEDGE_COLLECTIONS = ("fin_faq", "fin_fund_doc", "fin_policy")
EMBEDDING_DIMENSION = 768


def build_knowledge_schema():
    schema = AsyncMilvusClient.create_schema(auto_id=False, enable_dynamic_field=False)
    schema.add_field("chunk_id", DataType.VARCHAR, is_primary=True, max_length=128)
    schema.add_field("doc_id", DataType.VARCHAR, max_length=128)
    schema.add_field("title", DataType.VARCHAR, max_length=512)
    schema.add_field("section_title", DataType.VARCHAR, max_length=1024)
    schema.add_field("text", DataType.VARCHAR, max_length=65535)
    schema.add_field("strategy", DataType.VARCHAR, max_length=32)
    schema.add_field("vector", DataType.FLOAT_VECTOR, dim=EMBEDDING_DIMENSION)
    return schema


def build_knowledge_index_params():
    params = AsyncMilvusClient.prepare_index_params()
    params.add_index(
        field_name="vector",
        index_type="HNSW",
        metric_type="COSINE",
        params={"M": 16, "efConstruction": 200},
    )
    return params


async def ensure_collections(milvus_client: AsyncMilvusClient | None = None) -> None:
    client = milvus_client or configured_milvus_client()
    schema = build_knowledge_schema()
    index_params = build_knowledge_index_params()
    for collection_name in KNOWLEDGE_COLLECTIONS:
        if not await client.has_collection(collection_name):
            await client.create_collection(
                collection_name=collection_name,
                schema=schema,
                index_params=index_params,
            )