use crate::proto::control::v1::{
    field_type_list::ListValueType, field_type_matrix::MatrixValueType,
};

use super::*;

impl FieldSpec {
    pub fn with_index(mut self, index: FieldIndex) -> Self {
        assert!(self.index.is_none(), "Field index is already set");
        self.index = Some(index);
        self
    }

    pub fn text(required: bool, index_type: Option<KeywordIndexType>) -> FieldSpec {
        FieldSpec {
            data_type: Some(FieldType {
                data_type: Some(field_type::DataType::Text(FieldTypeText {})),
            }),
            required,
            index: index_type.map(|index_type| FieldIndex {
                index: Some(field_index::Index::KeywordIndex(KeywordIndex {
                    index_type: index_type as i32,
                })),
            }),
        }
    }

    pub fn integer(required: bool) -> FieldSpec {
        FieldSpec {
            data_type: Some(FieldType {
                data_type: Some(field_type::DataType::Integer(FieldTypeInteger {})),
            }),
            required,
            index: None,
        }
    }

    pub fn float(required: bool) -> FieldSpec {
        FieldSpec {
            data_type: Some(FieldType {
                data_type: Some(field_type::DataType::Float(FieldTypeFloat {})),
            }),
            required,
            index: None,
        }
    }

    pub fn boolean(required: bool) -> FieldSpec {
        FieldSpec {
            data_type: Some(FieldType {
                data_type: Some(field_type::DataType::Boolean(FieldTypeBoolean {})),
            }),
            required,
            index: None,
        }
    }

    pub fn bytes(required: bool) -> FieldSpec {
        FieldSpec {
            data_type: Some(FieldType {
                data_type: Some(field_type::DataType::Bytes(FieldTypeBytes {})),
            }),
            required,
            index: None,
        }
    }

    pub fn list(required: bool, value_type: ListValueType) -> FieldSpec {
        FieldSpec {
            data_type: Some(FieldType {
                data_type: Some(field_type::DataType::List(FieldTypeList {
                    value_type: value_type.into(),
                })),
            }),
            required,
            index: None,
        }
    }

    pub fn f32_vector(dimension: u32, required: bool, metric: VectorDistanceMetric) -> FieldSpec {
        Self::indexed_vector_field(
            field_type::DataType::f32_vector(dimension),
            required,
            metric,
        )
    }

    pub fn f16_vector(dimension: u32, required: bool, metric: VectorDistanceMetric) -> FieldSpec {
        Self::indexed_vector_field(
            field_type::DataType::f16_vector(dimension),
            required,
            metric,
        )
    }

    pub fn f8_vector(dimension: u32, required: bool, metric: VectorDistanceMetric) -> FieldSpec {
        Self::indexed_vector_field(field_type::DataType::f8_vector(dimension), required, metric)
    }

    pub fn u8_vector(dimension: u32, required: bool, metric: VectorDistanceMetric) -> FieldSpec {
        Self::indexed_vector_field(field_type::DataType::u8_vector(dimension), required, metric)
    }

    pub fn i8_vector(dimension: u32, required: bool, metric: VectorDistanceMetric) -> FieldSpec {
        Self::indexed_vector_field(field_type::DataType::i8_vector(dimension), required, metric)
    }

    pub fn binary_vector(
        dimension: u32,
        required: bool,
        metric: VectorDistanceMetric,
    ) -> FieldSpec {
        Self::indexed_vector_field(
            field_type::DataType::binary_vector(dimension),
            required,
            metric,
        )
    }

    pub fn f32_sparse_vector(required: bool, metric: VectorDistanceMetric) -> FieldSpec {
        Self::indexed_vector_field(
            field_type::DataType::F32SparseVector(FieldTypeF32SparseVector {}),
            required,
            metric,
        )
    }

    pub fn u8_sparse_vector(required: bool, metric: VectorDistanceMetric) -> FieldSpec {
        Self::indexed_vector_field(
            field_type::DataType::U8SparseVector(FieldTypeU8SparseVector {}),
            required,
            metric,
        )
    }

    fn indexed_vector_field(
        data_type: field_type::DataType,
        required: bool,
        metric: VectorDistanceMetric,
    ) -> FieldSpec {
        FieldSpec {
            data_type: Some(FieldType {
                data_type: Some(data_type),
            }),
            required,
            index: Some(FieldIndex {
                index: Some(field_index::Index::VectorIndex(VectorIndex {
                    metric: metric as i32,
                    exact: None,
                })),
            }),
        }
    }

    pub fn semantic(
        required: bool,
        model: Option<String>,
        embedding_type: Option<EmbeddingDataType>,
    ) -> FieldSpec {
        FieldSpec {
            data_type: Some(FieldType {
                data_type: Some(field_type::DataType::Text(FieldTypeText {})),
            }),
            required,
            index: Some(FieldIndex {
                index: Some(field_index::Index::SemanticIndex(SemanticIndex {
                    model,
                    embedding_type: embedding_type.map(|dt| dt.into()),
                })),
            }),
        }
    }

    pub fn matrix(required: bool, num_cols: u32, value_type: MatrixValueType) -> FieldSpec {
        FieldSpec {
            data_type: Some(FieldType {
                data_type: Some(field_type::DataType::Matrix(FieldTypeMatrix {
                    dimension: num_cols,
                    value_type: value_type.into(),
                })),
            }),
            required,
            index: None,
        }
    }
}
