use std::collections::BTreeMap;

use pyo3::prelude::*;
use qroissant_core::ValueType;

#[pyclass(
    module = "qroissant",
    eq,
    eq_int,
    frozen,
    rename_all = "SCREAMING_SNAKE_CASE"
)]
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub enum Attribute {
    None,
    Sorted,
    Unique,
    Parted,
    Grouped,
}

#[pymethods]
impl Attribute {}

impl From<qroissant_core::Attribute> for Attribute {
    fn from(value: qroissant_core::Attribute) -> Self {
        match value {
            qroissant_core::Attribute::None => Self::None,
            qroissant_core::Attribute::Sorted => Self::Sorted,
            qroissant_core::Attribute::Unique => Self::Unique,
            qroissant_core::Attribute::Parted => Self::Parted,
            qroissant_core::Attribute::Grouped => Self::Grouped,
        }
    }
}

impl From<Attribute> for qroissant_core::Attribute {
    fn from(value: Attribute) -> Self {
        match value {
            Attribute::None => Self::None,
            Attribute::Sorted => Self::Sorted,
            Attribute::Unique => Self::Unique,
            Attribute::Parted => Self::Parted,
            Attribute::Grouped => Self::Grouped,
        }
    }
}

#[pyclass(
    module = "qroissant",
    eq,
    eq_int,
    frozen,
    rename_all = "SCREAMING_SNAKE_CASE"
)]
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub enum Shape {
    Atom,
    Vector,
    List,
    Dictionary,
    Table,
    UnaryPrimitive,
    Error,
}

#[pymethods]
impl Shape {}

impl From<qroissant_core::Shape> for Shape {
    fn from(value: qroissant_core::Shape) -> Self {
        match value {
            qroissant_core::Shape::Atom => Self::Atom,
            qroissant_core::Shape::Vector => Self::Vector,
            qroissant_core::Shape::List => Self::List,
            qroissant_core::Shape::Dictionary => Self::Dictionary,
            qroissant_core::Shape::Table => Self::Table,
            qroissant_core::Shape::UnaryPrimitive => Self::UnaryPrimitive,
            qroissant_core::Shape::Error => Self::Error,
        }
    }
}

impl From<Shape> for qroissant_core::Shape {
    fn from(value: Shape) -> Self {
        match value {
            Shape::Atom => Self::Atom,
            Shape::Vector => Self::Vector,
            Shape::List => Self::List,
            Shape::Dictionary => Self::Dictionary,
            Shape::Table => Self::Table,
            Shape::UnaryPrimitive => Self::UnaryPrimitive,
            Shape::Error => Self::Error,
        }
    }
}

#[pyclass(
    module = "qroissant",
    eq,
    eq_int,
    frozen,
    rename_all = "SCREAMING_SNAKE_CASE"
)]
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub enum Primitive {
    Boolean,
    Guid,
    Byte,
    Short,
    Int,
    Long,
    Real,
    Float,
    Char,
    Symbol,
    Timestamp,
    Month,
    Date,
    Datetime,
    Timespan,
    Minute,
    Second,
    Time,
    Mixed,
}

#[pymethods]
impl Primitive {}

impl From<qroissant_core::Primitive> for Primitive {
    fn from(value: qroissant_core::Primitive) -> Self {
        match value {
            qroissant_core::Primitive::Boolean => Self::Boolean,
            qroissant_core::Primitive::Guid => Self::Guid,
            qroissant_core::Primitive::Byte => Self::Byte,
            qroissant_core::Primitive::Short => Self::Short,
            qroissant_core::Primitive::Int => Self::Int,
            qroissant_core::Primitive::Long => Self::Long,
            qroissant_core::Primitive::Real => Self::Real,
            qroissant_core::Primitive::Float => Self::Float,
            qroissant_core::Primitive::Char => Self::Char,
            qroissant_core::Primitive::Symbol => Self::Symbol,
            qroissant_core::Primitive::Timestamp => Self::Timestamp,
            qroissant_core::Primitive::Month => Self::Month,
            qroissant_core::Primitive::Date => Self::Date,
            qroissant_core::Primitive::Datetime => Self::Datetime,
            qroissant_core::Primitive::Timespan => Self::Timespan,
            qroissant_core::Primitive::Minute => Self::Minute,
            qroissant_core::Primitive::Second => Self::Second,
            qroissant_core::Primitive::Time => Self::Time,
            qroissant_core::Primitive::Mixed => Self::Mixed,
        }
    }
}

impl From<Primitive> for qroissant_core::Primitive {
    fn from(value: Primitive) -> Self {
        match value {
            Primitive::Boolean => Self::Boolean,
            Primitive::Guid => Self::Guid,
            Primitive::Byte => Self::Byte,
            Primitive::Short => Self::Short,
            Primitive::Int => Self::Int,
            Primitive::Long => Self::Long,
            Primitive::Real => Self::Real,
            Primitive::Float => Self::Float,
            Primitive::Char => Self::Char,
            Primitive::Symbol => Self::Symbol,
            Primitive::Timestamp => Self::Timestamp,
            Primitive::Month => Self::Month,
            Primitive::Date => Self::Date,
            Primitive::Datetime => Self::Datetime,
            Primitive::Timespan => Self::Timespan,
            Primitive::Minute => Self::Minute,
            Primitive::Second => Self::Second,
            Primitive::Time => Self::Time,
            Primitive::Mixed => Self::Mixed,
        }
    }
}

#[pyclass(
    module = "qroissant",
    eq,
    eq_int,
    frozen,
    rename_all = "SCREAMING_SNAKE_CASE"
)]
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub enum Encoding {
    LittleEndian,
    BigEndian,
}

#[pymethods]
impl Encoding {}

impl From<qroissant_core::Encoding> for Encoding {
    fn from(value: qroissant_core::Encoding) -> Self {
        match value {
            qroissant_core::Encoding::LittleEndian => Self::LittleEndian,
            qroissant_core::Encoding::BigEndian => Self::BigEndian,
        }
    }
}

impl From<Encoding> for qroissant_core::Encoding {
    fn from(value: Encoding) -> Self {
        match value {
            Encoding::LittleEndian => Self::LittleEndian,
            Encoding::BigEndian => Self::BigEndian,
        }
    }
}

#[pyclass(
    module = "qroissant",
    eq,
    eq_int,
    frozen,
    rename_all = "SCREAMING_SNAKE_CASE"
)]
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub enum Compression {
    Uncompressed,
    Compressed,
    CompressedLarge,
}

#[pymethods]
impl Compression {}

impl From<qroissant_core::Compression> for Compression {
    fn from(value: qroissant_core::Compression) -> Self {
        match value {
            qroissant_core::Compression::Uncompressed => Self::Uncompressed,
            qroissant_core::Compression::Compressed => Self::Compressed,
            qroissant_core::Compression::CompressedLarge => Self::CompressedLarge,
        }
    }
}

impl From<Compression> for qroissant_core::Compression {
    fn from(value: Compression) -> Self {
        match value {
            Compression::Uncompressed => Self::Uncompressed,
            Compression::Compressed => Self::Compressed,
            Compression::CompressedLarge => Self::CompressedLarge,
        }
    }
}

#[pyclass(
    module = "qroissant",
    eq,
    eq_int,
    frozen,
    rename_all = "SCREAMING_SNAKE_CASE"
)]
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub enum MessageType {
    Asynchronous,
    Synchronous,
    Response,
}

#[pymethods]
impl MessageType {}

impl From<qroissant_core::MessageType> for MessageType {
    fn from(value: qroissant_core::MessageType) -> Self {
        match value {
            qroissant_core::MessageType::Asynchronous => Self::Asynchronous,
            qroissant_core::MessageType::Synchronous => Self::Synchronous,
            qroissant_core::MessageType::Response => Self::Response,
        }
    }
}

impl From<MessageType> for qroissant_core::MessageType {
    fn from(value: MessageType) -> Self {
        match value {
            MessageType::Asynchronous => Self::Asynchronous,
            MessageType::Synchronous => Self::Synchronous,
            MessageType::Response => Self::Response,
        }
    }
}

#[pyclass(
    module = "qroissant",
    eq,
    eq_int,
    frozen,
    rename_all = "SCREAMING_SNAKE_CASE"
)]
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Default)]
pub enum SymbolInterpretation {
    #[default]
    Utf8,
    LargeUtf8,
    Utf8View,
    Dictionary,
    RawBytes,
}

#[pymethods]
impl SymbolInterpretation {}

#[pyclass(
    module = "qroissant",
    eq,
    eq_int,
    frozen,
    rename_all = "SCREAMING_SNAKE_CASE"
)]
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Default)]
pub enum ListInterpretation {
    #[default]
    List,
    LargeList,
    ListView,
}

#[pymethods]
impl ListInterpretation {}

#[pyclass(
    module = "qroissant",
    eq,
    eq_int,
    frozen,
    rename_all = "SCREAMING_SNAKE_CASE"
)]
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Default)]
pub enum StringInterpretation {
    #[default]
    Utf8,
    Binary,
}

#[pymethods]
impl StringInterpretation {}

#[pyclass(
    module = "qroissant",
    eq,
    eq_int,
    frozen,
    rename_all = "SCREAMING_SNAKE_CASE"
)]
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Default)]
pub enum UnionMode {
    #[default]
    Dense,
    Sparse,
}

#[pymethods]
impl UnionMode {}

#[pyclass(module = "qroissant", frozen, eq)]
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Type {
    primitive: Option<Primitive>,
    shape: Shape,
    attribute: Option<Attribute>,
    sorted: Option<bool>,
}

#[pymethods]
impl Type {
    #[new]
    #[pyo3(signature = (primitive, shape, attribute=None, sorted=None))]
    fn new(
        primitive: Option<Primitive>,
        shape: Shape,
        attribute: Option<Attribute>,
        sorted: Option<bool>,
    ) -> Self {
        Self {
            primitive,
            shape,
            attribute,
            sorted,
        }
    }

    #[getter]
    fn primitive(&self) -> Option<Primitive> {
        self.primitive
    }

    #[getter]
    fn shape(&self) -> Shape {
        self.shape
    }

    #[getter]
    fn attribute(&self) -> Option<Attribute> {
        self.attribute
    }

    #[getter]
    fn sorted(&self) -> Option<bool> {
        self.sorted
    }
}

impl From<ValueType> for Type {
    fn from(value: ValueType) -> Self {
        Self {
            primitive: value.primitive.map(Primitive::from),
            shape: Shape::from(value.shape),
            attribute: value.attribute.map(Attribute::from),
            sorted: value.sorted,
        }
    }
}

impl Type {
    pub fn to_core(&self) -> ValueType {
        ValueType {
            primitive: self.primitive.map(qroissant_core::Primitive::from),
            shape: qroissant_core::Shape::from(self.shape),
            attribute: self.attribute.map(qroissant_core::Attribute::from),
            sorted: self.sorted,
        }
    }

    pub fn primitive_value(&self) -> Option<Primitive> {
        self.primitive
    }

    pub fn shape_value(&self) -> Shape {
        self.shape
    }

    pub fn attribute_value(&self) -> Option<Attribute> {
        self.attribute
    }

    pub fn sorted_value(&self) -> Option<bool> {
        self.sorted
    }
}

#[pyclass(module = "qroissant", frozen, eq)]
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct MessageHeader {
    encoding: Encoding,
    message_type: MessageType,
    compression: Compression,
    size: usize,
}

#[pymethods]
impl MessageHeader {
    #[new]
    fn new(
        encoding: Encoding,
        message_type: MessageType,
        compression: Compression,
        size: usize,
    ) -> Self {
        Self {
            encoding,
            message_type,
            compression,
            size,
        }
    }

    #[getter]
    fn encoding(&self) -> Encoding {
        self.encoding
    }

    #[getter]
    fn message_type(&self) -> MessageType {
        self.message_type
    }

    #[getter]
    fn compression(&self) -> Compression {
        self.compression
    }

    #[getter]
    fn size(&self) -> usize {
        self.size
    }
}

impl From<qroissant_core::MessageHeader> for MessageHeader {
    fn from(value: qroissant_core::MessageHeader) -> Self {
        Self {
            encoding: Encoding::from(value.encoding()),
            message_type: MessageType::from(value.message_type()),
            compression: Compression::from(value.compression()),
            size: value.size(),
        }
    }
}

impl MessageHeader {
    pub fn new_native(
        encoding: Encoding,
        message_type: MessageType,
        compression: Compression,
        size: usize,
    ) -> Self {
        Self {
            encoding,
            message_type,
            compression,
            size,
        }
    }
}

#[pyclass(module = "qroissant", frozen, eq)]
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Endpoint {
    scheme: String,
    host: Option<String>,
    port: Option<u16>,
    path: Option<String>,
    username: Option<String>,
    password: Option<String>,
    timeout_ms: Option<u64>,
}

#[pymethods]
impl Endpoint {
    #[staticmethod]
    #[pyo3(signature = (host, port, *, username=None, password=None, timeout_ms=None))]
    fn tcp(
        host: String,
        port: u16,
        username: Option<String>,
        password: Option<String>,
        timeout_ms: Option<u64>,
    ) -> Self {
        Self {
            scheme: "tcp".to_string(),
            host: Some(host),
            port: Some(port),
            path: None,
            username,
            password,
            timeout_ms,
        }
    }

    #[staticmethod]
    #[pyo3(signature = (path, *, username=None, password=None, timeout_ms=None))]
    fn unix(
        path: String,
        username: Option<String>,
        password: Option<String>,
        timeout_ms: Option<u64>,
    ) -> Self {
        Self {
            scheme: "unix".to_string(),
            host: None,
            port: None,
            path: Some(path),
            username,
            password,
            timeout_ms,
        }
    }

    #[getter]
    fn scheme(&self) -> String {
        self.scheme.clone()
    }

    #[getter]
    fn host(&self) -> Option<String> {
        self.host.clone()
    }

    #[getter]
    fn port(&self) -> Option<u16> {
        self.port
    }

    #[getter]
    fn path(&self) -> Option<String> {
        self.path.clone()
    }

    #[getter]
    fn username(&self) -> Option<String> {
        self.username.clone()
    }

    #[getter]
    fn password(&self) -> Option<String> {
        self.password.clone()
    }

    #[getter]
    fn timeout_ms(&self) -> Option<u64> {
        self.timeout_ms
    }
}

impl Endpoint {
    pub fn validate(&self) -> Result<(), String> {
        match self.scheme.as_str() {
            "tcp" => {
                if self.host.is_none() {
                    return Err("tcp endpoints require a host".to_string());
                }
                if self.port.is_none() {
                    return Err("tcp endpoints require a port".to_string());
                }
            }
            "unix" => {
                if self.path.is_none() {
                    return Err("unix endpoints require a path".to_string());
                }
            }
            _ => {
                return Err(format!("unsupported endpoint scheme {:?}", self.scheme));
            }
        }
        Ok(())
    }

    pub fn username_deref(&self) -> Option<&str> {
        self.username.as_deref()
    }

    pub fn password_deref(&self) -> Option<&str> {
        self.password.as_deref()
    }

    pub fn scheme_value(&self) -> &str {
        &self.scheme
    }

    pub fn host_value(&self) -> Option<&str> {
        self.host.as_deref()
    }

    pub fn port_value(&self) -> Option<u16> {
        self.port
    }

    pub fn path_value(&self) -> Option<&str> {
        self.path.as_deref()
    }

    pub fn timeout_ms_value(&self) -> Option<u64> {
        self.timeout_ms
    }
}

#[pyclass(module = "qroissant", frozen, eq)]
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct PoolOptions {
    max_size: u32,
    min_idle: Option<u32>,
    checkout_timeout_ms: u64,
    idle_timeout_ms: Option<u64>,
    max_lifetime_ms: Option<u64>,
    test_on_checkout: bool,
    healthcheck_query: Option<String>,
    retry_attempts: u32,
    retry_backoff_ms: u64,
}

#[pymethods]
impl PoolOptions {
    #[new]
    #[pyo3(signature = (*, max_size=10, min_idle=None, checkout_timeout_ms=30_000, idle_timeout_ms=None, max_lifetime_ms=None, test_on_checkout=true, healthcheck_query=Some("::".to_string()), retry_attempts=0, retry_backoff_ms=0))]
    #[allow(clippy::too_many_arguments)]
    fn new(
        max_size: u32,
        min_idle: Option<u32>,
        checkout_timeout_ms: u64,
        idle_timeout_ms: Option<u64>,
        max_lifetime_ms: Option<u64>,
        test_on_checkout: bool,
        healthcheck_query: Option<String>,
        retry_attempts: u32,
        retry_backoff_ms: u64,
    ) -> Self {
        Self {
            max_size,
            min_idle,
            checkout_timeout_ms,
            idle_timeout_ms,
            max_lifetime_ms,
            test_on_checkout,
            healthcheck_query,
            retry_attempts,
            retry_backoff_ms,
        }
    }

    #[getter]
    fn max_size(&self) -> u32 {
        self.max_size
    }

    #[getter]
    fn min_idle(&self) -> Option<u32> {
        self.min_idle
    }

    #[getter]
    fn checkout_timeout_ms(&self) -> u64 {
        self.checkout_timeout_ms
    }

    #[getter]
    fn idle_timeout_ms(&self) -> Option<u64> {
        self.idle_timeout_ms
    }

    #[getter]
    fn max_lifetime_ms(&self) -> Option<u64> {
        self.max_lifetime_ms
    }

    #[getter]
    fn test_on_checkout(&self) -> bool {
        self.test_on_checkout
    }

    #[getter]
    fn healthcheck_query(&self) -> Option<String> {
        self.healthcheck_query.clone()
    }

    #[getter]
    fn retry_attempts(&self) -> u32 {
        self.retry_attempts
    }

    #[getter]
    fn retry_backoff_ms(&self) -> u64 {
        self.retry_backoff_ms
    }
}

impl Default for PoolOptions {
    fn default() -> Self {
        Self {
            max_size: 10,
            min_idle: None,
            checkout_timeout_ms: 30_000,
            idle_timeout_ms: None,
            max_lifetime_ms: None,
            test_on_checkout: true,
            healthcheck_query: Some("::".to_string()),
            retry_attempts: 0,
            retry_backoff_ms: 0,
        }
    }
}

impl PoolOptions {
    pub fn max_size_value(&self) -> u32 {
        self.max_size
    }

    pub fn min_idle_value(&self) -> Option<u32> {
        self.min_idle
    }

    pub fn checkout_timeout_ms_value(&self) -> u64 {
        self.checkout_timeout_ms
    }

    pub fn idle_timeout_ms_value(&self) -> Option<u64> {
        self.idle_timeout_ms
    }

    pub fn max_lifetime_ms_value(&self) -> Option<u64> {
        self.max_lifetime_ms
    }

    pub fn test_on_checkout_value(&self) -> bool {
        self.test_on_checkout
    }

    pub fn healthcheck_query_value(&self) -> Option<&str> {
        self.healthcheck_query.as_deref()
    }

    pub fn retry_attempts_value(&self) -> u32 {
        self.retry_attempts
    }

    pub fn retry_backoff_ms_value(&self) -> u64 {
        self.retry_backoff_ms
    }
}

#[pyclass(module = "qroissant", frozen, eq)]
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct DecodeOptions {
    list_interpretation: ListInterpretation,
    union_mode: UnionMode,
    string_interpretation: StringInterpretation,
    symbol_interpretation: SymbolInterpretation,
    assume_symbol_utf8: bool,
    parallel: bool,
    preserve_original_body: bool,
    validate_compressed_trailing_bytes: bool,
    temporal_nulls: bool,
    treat_infinity_as_null: bool,
    coalesce_temporals: bool,
}

#[pymethods]
impl DecodeOptions {
    #[staticmethod]
    fn builder() -> DecodeOptionsBuilder {
        DecodeOptionsBuilder::default()
    }

    #[getter]
    fn list_interpretation(&self) -> ListInterpretation {
        self.list_interpretation
    }

    #[getter]
    fn union_mode(&self) -> UnionMode {
        self.union_mode
    }

    #[getter]
    fn string_interpretation(&self) -> StringInterpretation {
        self.string_interpretation
    }

    #[getter]
    fn symbol_interpretation(&self) -> SymbolInterpretation {
        self.symbol_interpretation
    }

    #[getter]
    fn assume_symbol_utf8(&self) -> bool {
        self.assume_symbol_utf8
    }

    #[getter]
    fn parallel(&self) -> bool {
        self.parallel
    }

    #[getter]
    fn preserve_original_body(&self) -> bool {
        self.preserve_original_body
    }

    #[getter]
    fn validate_compressed_trailing_bytes(&self) -> bool {
        self.validate_compressed_trailing_bytes
    }

    #[getter]
    fn temporal_nulls(&self) -> bool {
        self.temporal_nulls
    }

    #[getter]
    fn get_treat_infinity_as_null(&self) -> bool {
        self.treat_infinity_as_null
    }

    #[getter]
    fn coalesce_temporals(&self) -> bool {
        self.coalesce_temporals
    }
}

impl Default for DecodeOptions {
    fn default() -> Self {
        Self {
            list_interpretation: ListInterpretation::List,
            union_mode: UnionMode::Dense,
            string_interpretation: StringInterpretation::Utf8,
            symbol_interpretation: SymbolInterpretation::Utf8,
            assume_symbol_utf8: true,
            parallel: true,
            preserve_original_body: true,
            validate_compressed_trailing_bytes: true,
            temporal_nulls: true,
            treat_infinity_as_null: false,
            coalesce_temporals: false,
        }
    }
}

impl DecodeOptions {
    pub fn is_default(&self) -> bool {
        self == &Self::default()
    }

    pub(crate) fn list_interpretation_value(&self) -> ListInterpretation {
        self.list_interpretation
    }

    pub(crate) fn string_interpretation_value(&self) -> StringInterpretation {
        self.string_interpretation
    }

    pub(crate) fn symbol_interpretation_value(&self) -> SymbolInterpretation {
        self.symbol_interpretation
    }

    pub(crate) fn union_mode_value(&self) -> UnionMode {
        self.union_mode
    }

    pub(crate) fn treat_infinity_as_null(&self) -> bool {
        self.treat_infinity_as_null
    }

    pub(crate) fn parallel_value(&self) -> bool {
        self.parallel
    }

    pub(crate) fn assume_symbol_utf8_value(&self) -> bool {
        self.assume_symbol_utf8
    }

    pub(crate) fn coalesce_temporals_value(&self) -> bool {
        self.coalesce_temporals
    }
}

#[pyclass(module = "qroissant", frozen, eq)]
#[derive(Clone, Debug, PartialEq, Eq, Default)]
pub struct DecodeOptionsBuilder {
    options: DecodeOptions,
}

#[pymethods]
impl DecodeOptionsBuilder {
    fn with_list_interpretation(&self, value: ListInterpretation) -> Self {
        let mut next = self.clone();
        next.options.list_interpretation = value;
        next
    }

    fn with_union_mode(&self, value: UnionMode) -> Self {
        let mut next = self.clone();
        next.options.union_mode = value;
        next
    }

    fn with_string_interpretation(&self, value: StringInterpretation) -> Self {
        let mut next = self.clone();
        next.options.string_interpretation = value;
        next
    }

    fn with_symbol_interpretation(&self, value: SymbolInterpretation) -> Self {
        let mut next = self.clone();
        next.options.symbol_interpretation = value;
        next
    }

    fn with_assume_symbol_utf8(&self, value: bool) -> Self {
        let mut next = self.clone();
        next.options.assume_symbol_utf8 = value;
        next
    }

    fn with_parallel(&self, value: bool) -> Self {
        let mut next = self.clone();
        next.options.parallel = value;
        next
    }

    fn with_preserve_original_body(&self, value: bool) -> Self {
        let mut next = self.clone();
        next.options.preserve_original_body = value;
        next
    }

    fn with_validate_compressed_trailing_bytes(&self, value: bool) -> Self {
        let mut next = self.clone();
        next.options.validate_compressed_trailing_bytes = value;
        next
    }

    fn with_temporal_nulls(&self, value: bool) -> Self {
        let mut next = self.clone();
        next.options.temporal_nulls = value;
        next
    }

    fn with_treat_infinity_as_null(&self, value: bool) -> Self {
        let mut next = self.clone();
        next.options.treat_infinity_as_null = value;
        next
    }

    fn with_coalesce_temporals(&self, value: bool) -> Self {
        let mut next = self.clone();
        next.options.coalesce_temporals = value;
        next
    }

    fn build(&self) -> DecodeOptions {
        self.options.clone()
    }

    #[getter]
    fn list_interpretation(&self) -> ListInterpretation {
        self.options.list_interpretation
    }

    #[getter]
    fn union_mode(&self) -> UnionMode {
        self.options.union_mode
    }

    #[getter]
    fn string_interpretation(&self) -> StringInterpretation {
        self.options.string_interpretation
    }

    #[getter]
    fn symbol_interpretation(&self) -> SymbolInterpretation {
        self.options.symbol_interpretation
    }

    #[getter]
    fn assume_symbol_utf8(&self) -> bool {
        self.options.assume_symbol_utf8
    }

    #[getter]
    fn parallel(&self) -> bool {
        self.options.parallel
    }

    #[getter]
    fn preserve_original_body(&self) -> bool {
        self.options.preserve_original_body
    }

    #[getter]
    fn validate_compressed_trailing_bytes(&self) -> bool {
        self.options.validate_compressed_trailing_bytes
    }

    #[getter]
    fn temporal_nulls(&self) -> bool {
        self.options.temporal_nulls
    }

    #[getter]
    fn treat_infinity_as_null(&self) -> bool {
        self.options.treat_infinity_as_null
    }

    #[getter]
    fn coalesce_temporals(&self) -> bool {
        self.options.coalesce_temporals
    }
}

#[pyclass(module = "qroissant", frozen, eq)]
#[derive(Clone, Debug, PartialEq, Eq, Default)]
pub struct EncodeField {
    primitive: Option<Primitive>,
    shape: Option<Shape>,
    attribute: Option<Attribute>,
    sorted: Option<bool>,
}

#[pymethods]
impl EncodeField {
    #[staticmethod]
    fn builder() -> EncodeFieldBuilder {
        EncodeFieldBuilder::default()
    }

    #[getter]
    fn primitive(&self) -> Option<Primitive> {
        self.primitive
    }

    #[getter]
    fn shape(&self) -> Option<Shape> {
        self.shape
    }

    #[getter]
    fn attribute(&self) -> Option<Attribute> {
        self.attribute
    }

    #[getter]
    fn sorted(&self) -> Option<bool> {
        self.sorted
    }
}

#[pyclass(module = "qroissant", frozen, eq)]
#[derive(Clone, Debug, PartialEq, Eq, Default)]
pub struct EncodeFieldBuilder {
    field: EncodeField,
}

#[pymethods]
impl EncodeFieldBuilder {
    fn with_primitive(&self, value: Primitive) -> Self {
        let mut next = self.clone();
        next.field.primitive = Some(value);
        next
    }

    fn with_shape(&self, value: Shape) -> Self {
        let mut next = self.clone();
        next.field.shape = Some(value);
        next
    }

    fn with_attribute(&self, value: Attribute) -> Self {
        let mut next = self.clone();
        next.field.attribute = Some(value);
        next
    }

    fn with_sorted(&self, value: bool) -> Self {
        let mut next = self.clone();
        next.field.sorted = Some(value);
        next
    }

    fn build(&self) -> EncodeField {
        self.field.clone()
    }
}

#[pyclass(module = "qroissant", frozen, eq)]
#[derive(Clone, Debug, PartialEq, Eq, Default)]
pub struct EncodeOptions {
    primitive: Option<Primitive>,
    shape: Option<Shape>,
    attribute: Option<Attribute>,
    strict: bool,
    fields: BTreeMap<String, EncodeField>,
}

#[pymethods]
impl EncodeOptions {
    #[staticmethod]
    fn builder() -> EncodeOptionsBuilder {
        EncodeOptionsBuilder::default()
    }

    #[getter]
    fn primitive(&self) -> Option<Primitive> {
        self.primitive
    }

    #[getter]
    fn shape(&self) -> Option<Shape> {
        self.shape
    }

    #[getter]
    fn attribute(&self) -> Option<Attribute> {
        self.attribute
    }

    #[getter]
    fn strict(&self) -> bool {
        self.strict
    }

    fn field(&self, name: &str) -> Option<EncodeField> {
        self.fields.get(name).cloned()
    }
}

#[pyclass(module = "qroissant", frozen, eq)]
#[derive(Clone, Debug, PartialEq, Eq, Default)]
pub struct EncodeOptionsBuilder {
    options: EncodeOptions,
}

#[pymethods]
impl EncodeOptionsBuilder {
    fn with_primitive(&self, value: Primitive) -> Self {
        let mut next = self.clone();
        next.options.primitive = Some(value);
        next
    }

    fn with_shape(&self, value: Shape) -> Self {
        let mut next = self.clone();
        next.options.shape = Some(value);
        next
    }

    fn with_attribute(&self, value: Attribute) -> Self {
        let mut next = self.clone();
        next.options.attribute = Some(value);
        next
    }

    fn with_strict(&self, value: bool) -> Self {
        let mut next = self.clone();
        next.options.strict = value;
        next
    }

    fn with_field(&self, name: String, field: EncodeField) -> Self {
        let mut next = self.clone();
        next.options.fields.insert(name, field);
        next
    }

    fn build(&self) -> EncodeOptions {
        self.options.clone()
    }
}

impl EncodeOptions {
    pub fn is_default(&self) -> bool {
        self == &Self::default()
    }
}

#[pyclass(module = "qroissant", frozen, eq)]
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct PoolMetrics {
    connections: u32,
    idle_connections: u32,
    max_size: u32,
    min_idle: Option<u32>,
    initialized: bool,
    closed: bool,
}

impl PoolMetrics {
    pub fn new_native(
        connections: u32,
        idle_connections: u32,
        max_size: u32,
        min_idle: Option<u32>,
        initialized: bool,
        closed: bool,
    ) -> Self {
        Self {
            connections,
            idle_connections,
            max_size,
            min_idle,
            initialized,
            closed,
        }
    }
}

#[pymethods]
impl PoolMetrics {
    #[new]
    fn new(
        connections: u32,
        idle_connections: u32,
        max_size: u32,
        min_idle: Option<u32>,
        initialized: bool,
        closed: bool,
    ) -> Self {
        Self {
            connections,
            idle_connections,
            max_size,
            min_idle,
            initialized,
            closed,
        }
    }

    #[getter]
    fn connections(&self) -> u32 {
        self.connections
    }

    #[getter]
    fn idle_connections(&self) -> u32 {
        self.idle_connections
    }

    #[getter]
    fn max_size(&self) -> u32 {
        self.max_size
    }

    #[getter]
    fn min_idle(&self) -> Option<u32> {
        self.min_idle
    }

    #[getter]
    fn initialized(&self) -> bool {
        self.initialized
    }

    #[getter]
    fn closed(&self) -> bool {
        self.closed
    }
}

pub fn register(module: &Bound<'_, PyModule>) -> PyResult<()> {
    module.add_class::<Attribute>()?;
    module.add_class::<Shape>()?;
    module.add_class::<Primitive>()?;
    module.add_class::<Encoding>()?;
    module.add_class::<Compression>()?;
    module.add_class::<MessageType>()?;
    module.add_class::<SymbolInterpretation>()?;
    module.add_class::<ListInterpretation>()?;
    module.add_class::<StringInterpretation>()?;
    module.add_class::<UnionMode>()?;
    module.add_class::<Type>()?;
    module.add_class::<MessageHeader>()?;
    module.add_class::<Endpoint>()?;
    module.add_class::<PoolOptions>()?;
    module.add_class::<DecodeOptions>()?;
    module.add_class::<DecodeOptionsBuilder>()?;
    module.add_class::<EncodeField>()?;
    module.add_class::<EncodeFieldBuilder>()?;
    module.add_class::<EncodeOptions>()?;
    module.add_class::<EncodeOptionsBuilder>()?;
    module.add_class::<PoolMetrics>()?;
    Ok(())
}
