# generated by datamodel-codegen:
#   filename:  openapi.json
#   timestamp: 2022-12-08T02:46:08+00:00

from __future__ import annotations

from typing import List, Optional

from pydantic import BaseModel, Field

from models import common, schema


class DescribeCollectionRequest(BaseModel):
    collection_name: Optional[str] = None
    collectionID: Optional[int] = Field(
        None, description='The collection ID you want to describe'
    )
    time_stamp: Optional[int] = Field(
        None,
        description='If time_stamp is not zero, will describe collection success when time_stamp >= created collection timestamp, otherwise will throw error.',
    )


class DropCollectionRequest(BaseModel):
    collection_name: Optional[str] = Field(
        None, description='The unique collection name in milvus.(Required)'
    )


class FieldData(BaseModel):
    field: Optional[List] = None
    field_id: Optional[int] = None
    field_name: Optional[str] = None
    type: Optional[int] = Field(
        None,
        description='0:   "None",\n1:   "Bool",\n2:   "Int8",\n3:   "Int16",\n4:   "Int32",\n5:   "Int64",\n10:  "Float",\n11:  "Double",\n20:  "String",\n21:  "VarChar",\n100: "BinaryVector",\n101: "FloatVector",',
    )


class GetCollectionStatisticsRequest(BaseModel):
    collection_name: Optional[str] = Field(
        None, description='The collection name you want get statistics'
    )


class HasCollectionRequest(BaseModel):
    collection_name: Optional[str] = Field(
        None, description='The unique collection name in milvus.(Required)'
    )
    time_stamp: Optional[int] = Field(
        None,
        description='If time_stamp is not zero, will return true when time_stamp >= created collection timestamp, otherwise will return false.',
    )


class InsertRequest(BaseModel):
    base: Optional[common.MsgBase] = None
    collection_name: Optional[str] = None
    db_name: Optional[str] = None
    fields_data: Optional[List[FieldData]] = None
    hash_keys: Optional[List[int]] = None
    num_rows: Optional[int] = None
    partition_name: Optional[str] = None


class LoadCollectionRequest(BaseModel):
    collection_name: Optional[str] = Field(
        None, description='The collection name you want to load'
    )
    replica_number: Optional[int] = Field(
        None, description='The replica number to load, default by 1'
    )


class ReleaseCollectionRequest(BaseModel):
    collection_name: Optional[str] = Field(
        None, description='The collection name you want to release'
    )


class ShowCollectionsRequest(BaseModel):
    collection_names: Optional[List[str]] = Field(
        None,
        description="When type is InMemory, will return these collection's inMemory_percentages.(Optional)",
    )
    type: Optional[int] = Field(
        None,
        description='Decide return Loaded collections or All collections(Optional)',
    )


class VectorIDs(BaseModel):
    collection_name: Optional[str] = None
    field_name: Optional[str] = None
    id_array: Optional[List[int]] = None
    partition_names: Optional[List[str]] = None


class VectorsArray(BaseModel):
    binary_vectors: Optional[List[int]] = Field(
        None,
        description='Vectors is an array of binary vector divided by given dim. Disabled when IDs is set',
    )
    dim: Optional[int] = Field(
        None, description='Dim of vectors or binary_vectors, not needed when use ids'
    )
    ids: Optional[VectorIDs] = None
    vectors: Optional[List[float]] = Field(
        None,
        description='Vectors is an array of vector divided by given dim. Disabled when ids or binary_vectors is set',
    )


class CalcDistanceRequest(BaseModel):
    base: Optional[common.MsgBase] = None
    op_left: Optional[VectorsArray] = None
    op_right: Optional[VectorsArray] = None
    params: Optional[List[common.KeyValuePair]] = None


class CreateCollectionRequest(BaseModel):
    collection_name: str = Field(
        ...,
        description='The unique collection name in milvus.(Required)',
        example='book',
    )
    consistency_level: int = Field(
        ...,
        description='The consistency level that the collection used, modification is not supported now.\n"Strong":     0,\n"Session":    1,\n"Bounded":    2,\n"Eventually": 3,\n"Customized": 4,',
        example=1,
    )
    schema_: schema.CollectionSchema = Field(..., alias='schema')
    shards_num: Optional[int] = Field(
        None,
        description='Once set, no modification is allowed (Optional)\nhttps://github.com/milvus-io/milvus/issues/6690',
        example=1,
    )
