"""
Milvus collection schema for the fashion_products multimodal collection.

Dataset: ashraq/fashion-product-images-small
Embedding model: CLIP ViT-B/32 (clip-ViT-B-32 via sentence-transformers)
"""

from pymilvus import DataType, MilvusClient
from config import EMBEDDING_DIM, NLIST

SCALAR_OUTPUT_FIELDS = [
    "id", "product_id", "product_name", "gender",
    "master_category", "sub_category", "article_type",
    "base_colour", "season", "year", "usage", "has_image",
]


def fashion_schema(client: MilvusClient):
    """
    Full schema for the fashion_products collection.

    Scalar fields    — Metadata used for filtering, display, and faceting.
    text_embedding   — CLIP encoding of the productDisplayName field.
    image_embedding  — CLIP encoding of the product photograph.
    has_image        — False when an image was unavailable during ingestion.
    """
    schema = client.create_schema(
        auto_id=True,
        enable_dynamic_fields=False,
        description=(
            "Multimodal fashion products (ashraq/fashion-product-images-small). "
            "Both vector fields use CLIP ViT-B/32 (512-dim, COSINE) in the same space. "
            "Supports: text-semantic, image-similarity, text-to-image, and image-to-text searches."
        ),
    )

    # ── Scalar Metadata ──────────────────────────────────────────────────────
    schema.add_field(
        "id", DataType.INT64, is_primary=True,
        description="Primary key automatically generated by Milvus",
    )
    schema.add_field(
        "product_id", DataType.INT64,
        description="Original numerical product ID from the dataset",
    )
    schema.add_field(
        "product_name", DataType.VARCHAR, max_length=512,
        description="productDisplayName — text used to build the text_embedding",
    )
    schema.add_field(
        "gender", DataType.VARCHAR, max_length=32,
        description="Product gender category"
    )
    schema.add_field(
        "master_category", DataType.VARCHAR, max_length=64,
        description="Main product category"
    )
    schema.add_field(
        "sub_category", DataType.VARCHAR, max_length=64,
        description="Product subcategory"
    )
    schema.add_field(
        "article_type", DataType.VARCHAR, max_length=64,
        description="Specific article type"
    )
    schema.add_field(
        "base_colour", DataType.VARCHAR, max_length=64,
        description="Primary color of the product"
    )
    schema.add_field(
        "season", DataType.VARCHAR, max_length=32,
        description="Season category"
    )
    schema.add_field(
        "year", DataType.INT32,
        description="Year of product release"
    )
    schema.add_field(
        "usage", DataType.VARCHAR, max_length=64,
        description="Usage category"
    )
    schema.add_field(
        "has_image", DataType.BOOL,
        description="True when image_embedding encodes an actual product photograph",
    )

    # ── Vector Fields — both reside in the same 512-dim CLIP space ───────────
    schema.add_field(
        "text_embedding",
        DataType.FLOAT_VECTOR,
        dim=EMBEDDING_DIM,
        description="CLIP text encoding from productDisplayName (512-dim)",
    )
    schema.add_field(
        "image_embedding",
        DataType.FLOAT_VECTOR,
        dim=EMBEDDING_DIM,
        description="CLIP image encoding of the product photo (512-dim)",
    )

    return schema


def fashion_index_params(client: MilvusClient):
    """
    Index parameters for the fashion_products collection.

    Scalar indexes — INVERTED for efficient WHERE filtering.
    Vector indexes — IVF_FLAT with COSINE; one index per vector field.
    """
    index_params = client.prepare_index_params()

    # Scalar indexes - primary key is indexed automatically
    index_params.add_index("product_id", index_type="INVERTED")
    index_params.add_index("year", index_type="INVERTED")
    index_params.add_index("gender", index_type="INVERTED")
    index_params.add_index("master_category", index_type="INVERTED")
    index_params.add_index("article_type", index_type="INVERTED")
    index_params.add_index("base_colour", index_type="INVERTED")
    index_params.add_index("season", index_type="INVERTED")
    index_params.add_index("usage", index_type="INVERTED")

    # text_embedding — used for: text-to-text, image-to-text
    index_params.add_index(
        field_name="text_embedding",
        index_type="IVF_FLAT",
        index_name="text_embedding_index",
        metric_type="COSINE",
        params={"nlist": NLIST}
    )
    
    # image_embedding — used for: image-to-image, text-to-image
    index_params.add_index(
        field_name="image_embedding",
        index_type="IVF_FLAT",
        index_name="image_embedding_index",
        metric_type="COSINE",
        params={"nlist": NLIST}
    )

    return index_params