use std::io::Read;
use std::io::Write;
use std::pin::Pin;
use std::sync::Arc;
use std::sync::Mutex;
use std::task::Context;
use std::task::Poll;
use std::thread;
use std::time::Duration;

use bb8::ManageConnection as AsyncManageConnection;
use pyo3::PyClass;
use pyo3::prelude::*;
use pyo3::types::PyAny;
use pyo3::types::PyModule;
use pyo3_async_runtimes::tokio::future_into_py;
use qroissant_core::Value;
use qroissant_transport::AsyncPooledTransport;
use qroissant_transport::AsyncTransport;
use qroissant_transport::QIPC_HEADER_LEN;
use qroissant_transport::SyncPooledTransport;
use qroissant_transport::SyncTransport;
use qroissant_transport::begin_streaming_frame_over;
use qroissant_transport::begin_streaming_frame_over_async;
use qroissant_transport::connect_tcp_transport;
use qroissant_transport::connect_tcp_transport_async;
#[cfg(unix)]
use qroissant_transport::connect_unix_transport;
#[cfg(unix)]
use qroissant_transport::connect_unix_transport_async;
use qroissant_transport::encode_sync_query;
use qroissant_transport::extract_q_error;
use qroissant_transport::request_frame_streaming_over;
use qroissant_transport::request_value_pipelined_over_async;
use qroissant_transport::validate_response_frame;
use qroissant_transport::validate_response_header_bytes;
use r2d2::ManageConnection;
use r2d2::PooledConnection;
use tokio::io::AsyncRead;
use tokio::io::AsyncWrite;
use tokio::io::ReadBuf;
use tokio::sync::Mutex as AsyncMutex;
use tokio::time::sleep;

use crate::errors::PythonError;
use crate::errors::to_py_err;
use crate::raw_response::AsyncStreamingLease;
use crate::raw_response::BlockingAsyncBridge;
use crate::raw_response::RawResponse;
use crate::raw_response::SyncRawLease;
use crate::serde::decode_core_value;
use crate::types::DecodeOptions;
use crate::types::Endpoint;
use crate::types::MessageHeader;
use crate::types::PoolMetrics;
use crate::types::PoolOptions;
use crate::values::core_value_to_python_with_opts;

fn async_self<T, F>(py: Python<'_>, constructor: F) -> PyResult<Bound<'_, PyAny>>
where
    T: PyClass + Send + Sync + 'static,
    F: FnOnce(Python<'_>) -> PyResult<Py<T>> + Send + 'static,
{
    future_into_py(py, async move {
        Python::attach(|py| constructor(py).map(|value| value.into_any()))
    })
}

fn connect_sync_transport(endpoint: &Endpoint) -> Result<SyncTransport, PythonError> {
    endpoint.validate().map_err(PythonError::Transport)?;
    match endpoint.scheme_value() {
        "tcp" => connect_tcp_transport(
            endpoint.host_value().unwrap_or_default(),
            endpoint.port_value().unwrap_or_default(),
            endpoint.username_deref(),
            endpoint.password_deref(),
            endpoint.timeout_ms_value(),
        )
        .map_err(crate::errors::map_transport_error),
        #[cfg(unix)]
        "unix" => connect_unix_transport(
            endpoint.path_value().unwrap_or_default(),
            endpoint.username_deref(),
            endpoint.password_deref(),
            endpoint.timeout_ms_value(),
        )
        .map_err(crate::errors::map_transport_error),
        other => Err(PythonError::Transport(format!(
            "unsupported endpoint scheme {other:?}"
        ))),
    }
}

async fn connect_async_transport(endpoint: &Endpoint) -> Result<AsyncTransport, PythonError> {
    endpoint.validate().map_err(PythonError::Transport)?;
    match endpoint.scheme_value() {
        "tcp" => connect_tcp_transport_async(
            endpoint.host_value().unwrap_or_default(),
            endpoint.port_value().unwrap_or_default(),
            endpoint.username_deref(),
            endpoint.password_deref(),
            endpoint.timeout_ms_value(),
        )
        .await
        .map_err(crate::errors::map_transport_error),
        #[cfg(unix)]
        "unix" => connect_unix_transport_async(
            endpoint.path_value().unwrap_or_default(),
            endpoint.username_deref(),
            endpoint.password_deref(),
            endpoint.timeout_ms_value(),
        )
        .await
        .map_err(crate::errors::map_transport_error),
        other => Err(PythonError::Transport(format!(
            "unsupported endpoint scheme {other:?}"
        ))),
    }
}

fn validate_success_response(frame: &[u8]) -> Result<(), PythonError> {
    validate_response_frame(frame).map_err(crate::errors::map_transport_error)?;
    if let Some(message) = extract_q_error(frame).map_err(crate::errors::map_transport_error)? {
        return Err(PythonError::QRuntime(message));
    }
    Ok(())
}

fn response_header(header_bytes: [u8; QIPC_HEADER_LEN]) -> Result<MessageHeader, PythonError> {
    let header =
        validate_response_header_bytes(header_bytes).map_err(crate::errors::map_transport_error)?;
    Ok(MessageHeader::from(header))
}

fn retryable_pool_error(error: &PythonError) -> bool {
    matches!(error, PythonError::Transport(_) | PythonError::Pool(_))
}

type SyncRawQuery = (
    MessageHeader,
    [u8; QIPC_HEADER_LEN],
    usize,
    Box<dyn SyncRawLease>,
);

type AsyncConnectionRawQuery = (
    MessageHeader,
    [u8; QIPC_HEADER_LEN],
    usize,
    AsyncConnectionLease,
);

type AsyncPoolRawQuery = (
    MessageHeader,
    [u8; QIPC_HEADER_LEN],
    usize,
    AsyncPooledConnectionLease,
);

#[derive(Debug, thiserror::Error)]
enum PoolBackendError {
    #[error("{0}")]
    Message(String),
}

impl From<PythonError> for PoolBackendError {
    fn from(value: PythonError) -> Self {
        Self::Message(value.to_string())
    }
}

fn sync_backoff(delay_ms: u64) {
    if delay_ms > 0 {
        thread::sleep(Duration::from_millis(delay_ms));
    }
}

async fn async_backoff(delay_ms: u64) {
    if delay_ms > 0 {
        sleep(Duration::from_millis(delay_ms)).await;
    }
}

struct ConnectionSlotState {
    transport: Option<SyncTransport>,
    closed: bool,
}

struct ConnectionLease {
    state: Arc<Mutex<ConnectionSlotState>>,
    transport: Option<SyncTransport>,
    reusable: bool,
}

impl ConnectionLease {
    fn new(state: Arc<Mutex<ConnectionSlotState>>, transport: SyncTransport) -> Self {
        Self {
            state,
            transport: Some(transport),
            reusable: false,
        }
    }
}

impl Read for ConnectionLease {
    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
        self.transport
            .as_mut()
            .expect("active connection lease must hold a transport")
            .read(buf)
    }
}

impl Write for ConnectionLease {
    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
        self.transport
            .as_mut()
            .expect("active connection lease must hold a transport")
            .write(buf)
    }

    fn flush(&mut self) -> std::io::Result<()> {
        self.transport
            .as_mut()
            .expect("active connection lease must hold a transport")
            .flush()
    }
}

impl SyncRawLease for ConnectionLease {
    fn mark_reusable(&mut self) {
        self.reusable = true;
    }

    fn abandon(&mut self) {
        self.reusable = false;
    }
}

impl Drop for ConnectionLease {
    fn drop(&mut self) {
        let Some(mut transport) = self.transport.take() else {
            return;
        };

        let mut state = match self.state.lock() {
            Ok(state) => state,
            Err(_) => {
                let _ = transport.shutdown();
                return;
            }
        };

        if self.reusable && !state.closed && state.transport.is_none() {
            state.transport = Some(transport);
            return;
        }

        let _ = transport.shutdown();
    }
}

struct ConnectionCore {
    state: Arc<Mutex<ConnectionSlotState>>,
}

impl ConnectionCore {
    fn connect(endpoint: &Endpoint) -> Result<Self, PythonError> {
        Ok(Self {
            state: Arc::new(Mutex::new(ConnectionSlotState {
                transport: Some(connect_sync_transport(endpoint)?),
                closed: false,
            })),
        })
    }

    fn checkout(&self) -> Result<ConnectionLease, PythonError> {
        let mut state = self
            .state
            .lock()
            .map_err(|_| PythonError::Operation("connection lock is poisoned".to_string()))?;
        if state.closed {
            return Err(PythonError::Operation("connection is closed".to_string()));
        }
        let transport = state.transport.take().ok_or_else(|| {
            PythonError::Operation("connection is busy with an active raw response".to_string())
        })?;
        Ok(ConnectionLease::new(self.state.clone(), transport))
    }

    /// Sends a synchronous query and reads the full response frame.
    fn query_frame(&self, expr: &str) -> Result<Vec<u8>, PythonError> {
        let payload = encode_sync_query(expr).map_err(crate::errors::map_transport_error)?;
        let mut lease = self.checkout()?;
        match request_frame_streaming_over(&mut lease, &payload) {
            Ok(frame) => {
                validate_response_frame(&frame).map_err(crate::errors::map_transport_error)?;
                lease.mark_reusable();
                Ok(frame)
            }
            Err(error) => {
                lease.abandon();
                Err(crate::errors::map_transport_error(error))
            }
        }
    }

    /// Sends a synchronous query and returns a streaming raw-response handle.
    fn begin_raw_query(&self, expr: &str) -> Result<SyncRawQuery, PythonError> {
        let payload = encode_sync_query(expr).map_err(crate::errors::map_transport_error)?;
        let mut lease = self.checkout()?;
        match begin_streaming_frame_over(&mut lease, &payload) {
            Ok((header_bytes, remaining_body)) => {
                let header = match response_header(header_bytes) {
                    Ok(header) => header,
                    Err(error) => {
                        lease.abandon();
                        return Err(error);
                    }
                };
                Ok((header, header_bytes, remaining_body, Box::new(lease)))
            }
            Err(error) => {
                lease.abandon();
                Err(crate::errors::map_transport_error(error))
            }
        }
    }

    fn close(&self) -> Result<(), PythonError> {
        let transport = {
            let mut state = self
                .state
                .lock()
                .map_err(|_| PythonError::Operation("connection lock is poisoned".to_string()))?;
            state.closed = true;
            state.transport.take()
        };
        if let Some(mut transport) = transport {
            transport
                .shutdown()
                .map_err(|error| PythonError::Transport(error.to_string()))?;
        }
        Ok(())
    }
}

/// Connection slot state machine for async connections.
///
/// ```text
/// Disconnected ─── connect ──→ Busy ─── success ──→ Ready
///       ↑                        │                    │
///       └── connect fail ────────┘                    │
///       └── mark_broken ─────────────────────────────┘
///                                                     │
///                                 Closed ←── close ───┘
/// ```
enum AsyncConnectionSlotState {
    Disconnected,
    Ready(AsyncTransport),
    Busy,
    Closed,
}

struct AsyncConnectionLease {
    state: Arc<AsyncMutex<AsyncConnectionSlotState>>,
    transport: Option<AsyncTransport>,
    reusable: bool,
}

impl AsyncConnectionLease {
    fn new(state: Arc<AsyncMutex<AsyncConnectionSlotState>>, transport: AsyncTransport) -> Self {
        Self {
            state,
            transport: Some(transport),
            reusable: false,
        }
    }

    fn restore_transport(
        state: &mut AsyncConnectionSlotState,
        transport: AsyncTransport,
        reusable: bool,
    ) {
        if reusable {
            if matches!(state, AsyncConnectionSlotState::Busy) {
                *state = AsyncConnectionSlotState::Ready(transport);
                return;
            }
        } else if matches!(state, AsyncConnectionSlotState::Busy) {
            *state = AsyncConnectionSlotState::Disconnected;
        }
        drop(transport);
    }
}

impl AsyncRead for AsyncConnectionLease {
    fn poll_read(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &mut ReadBuf<'_>,
    ) -> Poll<std::io::Result<()>> {
        Pin::new(
            self.transport
                .as_mut()
                .expect("active async connection lease must hold a transport"),
        )
        .poll_read(cx, buf)
    }
}

impl AsyncWrite for AsyncConnectionLease {
    fn poll_write(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &[u8],
    ) -> Poll<std::io::Result<usize>> {
        Pin::new(
            self.transport
                .as_mut()
                .expect("active async connection lease must hold a transport"),
        )
        .poll_write(cx, buf)
    }

    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
        Pin::new(
            self.transport
                .as_mut()
                .expect("active async connection lease must hold a transport"),
        )
        .poll_flush(cx)
    }

    fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
        Pin::new(
            self.transport
                .as_mut()
                .expect("active async connection lease must hold a transport"),
        )
        .poll_shutdown(cx)
    }
}

impl AsyncStreamingLease for AsyncConnectionLease {
    fn mark_reusable(&mut self) {
        self.reusable = true;
    }

    fn abandon(&mut self) {
        self.reusable = false;
    }
}

impl Drop for AsyncConnectionLease {
    fn drop(&mut self) {
        let Some(transport) = self.transport.take() else {
            return;
        };

        if let Ok(mut state) = self.state.try_lock() {
            Self::restore_transport(&mut state, transport, self.reusable);
            return;
        }

        if let Ok(handle) = tokio::runtime::Handle::try_current() {
            let state = self.state.clone();
            let reusable = self.reusable;
            handle.spawn(async move {
                let mut state = state.lock().await;
                Self::restore_transport(&mut state, transport, reusable);
            });
            return;
        }

        let mut state = self.state.blocking_lock();
        Self::restore_transport(&mut state, transport, self.reusable);
    }
}

struct AsyncConnectionCore {
    endpoint: Endpoint,
    state: Arc<AsyncMutex<AsyncConnectionSlotState>>,
}

impl AsyncConnectionCore {
    fn new(endpoint: Endpoint) -> Result<Self, PythonError> {
        endpoint.validate().map_err(PythonError::Transport)?;
        Ok(Self {
            endpoint,
            state: Arc::new(AsyncMutex::new(AsyncConnectionSlotState::Disconnected)),
        })
    }

    async fn checkout(&self) -> Result<AsyncConnectionLease, PythonError> {
        let should_connect = {
            let mut state = self.state.lock().await;
            match std::mem::replace(&mut *state, AsyncConnectionSlotState::Busy) {
                AsyncConnectionSlotState::Ready(transport) => {
                    return Ok(AsyncConnectionLease::new(self.state.clone(), transport));
                }
                AsyncConnectionSlotState::Disconnected => true,
                AsyncConnectionSlotState::Busy => {
                    *state = AsyncConnectionSlotState::Busy;
                    return Err(PythonError::Operation(
                        "connection is busy with an active raw response".to_string(),
                    ));
                }
                AsyncConnectionSlotState::Closed => {
                    *state = AsyncConnectionSlotState::Closed;
                    return Err(PythonError::Operation("connection is closed".to_string()));
                }
            }
        };

        debug_assert!(should_connect);
        match connect_async_transport(&self.endpoint).await {
            Ok(transport) => Ok(AsyncConnectionLease::new(self.state.clone(), transport)),
            Err(error) => {
                let mut state = self.state.lock().await;
                if !matches!(&*state, AsyncConnectionSlotState::Closed) {
                    *state = AsyncConnectionSlotState::Disconnected;
                }
                Err(error)
            }
        }
    }

    async fn query_value(&self, expr: &str) -> Result<Value, PythonError> {
        let payload = encode_sync_query(expr).map_err(crate::errors::map_transport_error)?;
        let mut lease = self.checkout().await?;
        match qroissant_transport::request_value_pipelined_over_async(&mut lease, &payload).await {
            Ok(value) => {
                lease.mark_reusable();
                Ok(value)
            }
            Err(error) => {
                lease.abandon();
                Err(crate::errors::map_transport_error(error))
            }
        }
    }

    /// Sends an async query and returns a streaming raw-response handle.
    async fn begin_raw_query(&self, expr: &str) -> Result<AsyncConnectionRawQuery, PythonError> {
        let payload = encode_sync_query(expr).map_err(crate::errors::map_transport_error)?;
        let mut lease = self.checkout().await?;
        match begin_streaming_frame_over_async(&mut lease, &payload).await {
            Ok((header_bytes, remaining_body)) => {
                let header = match response_header(header_bytes) {
                    Ok(header) => header,
                    Err(error) => {
                        lease.abandon();
                        return Err(error);
                    }
                };
                Ok((header, header_bytes, remaining_body, lease))
            }
            Err(error) => {
                lease.abandon();
                Err(crate::errors::map_transport_error(error))
            }
        }
    }

    async fn close(&self) -> Result<(), PythonError> {
        let transport = {
            let mut state = self.state.lock().await;
            match std::mem::replace(&mut *state, AsyncConnectionSlotState::Closed) {
                AsyncConnectionSlotState::Ready(transport) => Some(transport),
                AsyncConnectionSlotState::Disconnected
                | AsyncConnectionSlotState::Busy
                | AsyncConnectionSlotState::Closed => None,
            }
        };
        if let Some(mut transport) = transport {
            transport
                .shutdown()
                .await
                .map_err(|error| PythonError::Transport(error.to_string()))?;
        }
        Ok(())
    }
}

fn sync_request(conn: &mut SyncPooledTransport, payload: &[u8]) -> Result<Vec<u8>, PythonError> {
    match request_frame_streaming_over(conn, payload) {
        Ok(frame) => Ok(frame),
        Err(error) => {
            conn.mark_broken();
            Err(crate::errors::map_transport_error(error))
        }
    }
}

fn sync_send_query(conn: &mut SyncPooledTransport, query: &str) -> Result<Vec<u8>, PythonError> {
    let payload = encode_sync_query(query).map_err(crate::errors::map_transport_error)?;
    sync_request(conn, &payload)
}

fn sync_validate_connection(
    conn: &mut SyncPooledTransport,
    healthcheck_query: Option<&str>,
) -> Result<(), PoolBackendError> {
    if conn.is_broken() {
        return Err(PoolBackendError::Message(
            "pooled connection is broken".to_string(),
        ));
    }
    if let Some(query) = healthcheck_query {
        let response = sync_send_query(conn, query)?;
        validate_success_response(&response)?;
    }
    Ok(())
}

#[derive(Clone)]
struct SyncTransportManager {
    endpoint: Endpoint,
    healthcheck_query: Option<String>,
}

impl ManageConnection for SyncTransportManager {
    type Connection = SyncPooledTransport;
    type Error = PoolBackendError;

    fn connect(&self) -> Result<Self::Connection, Self::Error> {
        connect_sync_transport(&self.endpoint)
            .map(SyncPooledTransport::new)
            .map_err(Into::into)
    }

    fn is_valid(&self, conn: &mut Self::Connection) -> Result<(), Self::Error> {
        sync_validate_connection(conn, self.healthcheck_query.as_deref())
    }

    fn has_broken(&self, conn: &mut Self::Connection) -> bool {
        conn.is_broken()
    }
}

struct SyncPooledConnectionLease {
    connection: Option<PooledConnection<SyncTransportManager>>,
    reusable: bool,
}

impl SyncPooledConnectionLease {
    fn new(connection: PooledConnection<SyncTransportManager>) -> Self {
        Self {
            connection: Some(connection),
            reusable: false,
        }
    }

    fn mark_inner_broken(&mut self) {
        if let Some(connection) = self.connection.as_mut() {
            connection.mark_broken();
        }
    }
}

impl Read for SyncPooledConnectionLease {
    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
        self.connection
            .as_mut()
            .expect("active pooled lease must hold a checked-out connection")
            .read(buf)
    }
}

impl SyncRawLease for SyncPooledConnectionLease {
    fn mark_reusable(&mut self) {
        self.reusable = true;
    }

    fn abandon(&mut self) {
        self.reusable = false;
        self.mark_inner_broken();
    }
}

impl Drop for SyncPooledConnectionLease {
    fn drop(&mut self) {
        if !self.reusable {
            self.mark_inner_broken();
        }
    }
}

struct SyncPoolState {
    pool: Option<r2d2::Pool<SyncTransportManager>>,
    closed: bool,
}

struct SyncPoolCore {
    manager: SyncTransportManager,
    options: PoolOptions,
    state: Mutex<SyncPoolState>,
}

impl SyncPoolCore {
    fn new(endpoint: Endpoint, options: PoolOptions) -> Result<Self, PythonError> {
        endpoint.validate().map_err(PythonError::Pool)?;
        Ok(Self {
            manager: SyncTransportManager {
                endpoint,
                healthcheck_query: options.healthcheck_query_value().map(str::to_string),
            },
            options,
            state: Mutex::new(SyncPoolState {
                pool: None,
                closed: false,
            }),
        })
    }

    fn build_pool(&self) -> Result<r2d2::Pool<SyncTransportManager>, PoolBackendError> {
        let mut builder = r2d2::Pool::builder()
            .max_size(self.options.max_size_value())
            .test_on_check_out(self.options.test_on_checkout_value())
            .connection_timeout(Duration::from_millis(
                self.options.checkout_timeout_ms_value(),
            ));
        builder = builder.min_idle(self.options.min_idle_value());
        builder = builder.idle_timeout(
            self.options
                .idle_timeout_ms_value()
                .map(Duration::from_millis),
        );
        builder = builder.max_lifetime(
            self.options
                .max_lifetime_ms_value()
                .map(Duration::from_millis),
        );
        builder
            .build(self.manager.clone())
            .map_err(|error| PoolBackendError::Message(error.to_string()))
    }

    fn pool(&self) -> Result<r2d2::Pool<SyncTransportManager>, PythonError> {
        let mut state = self
            .state
            .lock()
            .map_err(|_| PythonError::Pool("pool lock is poisoned".to_string()))?;
        if state.closed {
            return Err(PythonError::PoolClosed);
        }
        if let Some(pool) = &state.pool {
            return Ok(pool.clone());
        }
        let pool = self
            .build_pool()
            .map_err(|error| PythonError::Pool(error.to_string()))?;
        state.pool = Some(pool.clone());
        Ok(pool)
    }

    fn close(&self) -> Result<(), PythonError> {
        let mut state = self
            .state
            .lock()
            .map_err(|_| PythonError::Pool("pool lock is poisoned".to_string()))?;
        state.closed = true;
        state.pool.take();
        Ok(())
    }

    fn metrics(&self) -> Result<PoolMetrics, PythonError> {
        let state = self
            .state
            .lock()
            .map_err(|_| PythonError::Pool("pool lock is poisoned".to_string()))?;
        let (connections, idle_connections, initialized) = state
            .pool
            .as_ref()
            .map(|pool| {
                let pool_state = pool.state();
                (pool_state.connections, pool_state.idle_connections, true)
            })
            .unwrap_or((0, 0, false));
        Ok(PoolMetrics::new_native(
            connections,
            idle_connections,
            self.options.max_size_value(),
            self.options.min_idle_value(),
            initialized,
            state.closed,
        ))
    }

    fn prewarm(&self) -> Result<PoolMetrics, PythonError> {
        let target = self.options.min_idle_value().unwrap_or(1);
        if target == 0 {
            return self.metrics();
        }
        let pool = self.pool()?;
        let mut connections = Vec::with_capacity(target as usize);
        for _ in 0..target {
            connections.push(
                pool.get()
                    .map_err(|error| PythonError::Pool(error.to_string()))?,
            );
        }
        drop(connections);
        self.metrics()
    }

    fn try_query_frame_once(&self, expr: &str) -> Result<Vec<u8>, PythonError> {
        let payload = encode_sync_query(expr).map_err(crate::errors::map_transport_error)?;
        let pool = self.pool()?;
        let mut connection = pool
            .get()
            .map_err(|error| PythonError::Pool(error.to_string()))?;
        let frame = request_frame_streaming_over(&mut *connection, &payload).map_err(|error| {
            connection.mark_broken();
            crate::errors::map_transport_error(error)
        })?;
        validate_response_frame(&frame).map_err(crate::errors::map_transport_error)?;
        Ok(frame)
    }

    fn query_frame(&self, expr: &str) -> Result<Vec<u8>, PythonError> {
        for attempt in 0..=self.options.retry_attempts_value() {
            match self.try_query_frame_once(expr) {
                Ok(frame) => return Ok(frame),
                Err(error) => {
                    if attempt == self.options.retry_attempts_value()
                        || !retryable_pool_error(&error)
                    {
                        return Err(error);
                    }
                    sync_backoff(self.options.retry_backoff_ms_value());
                }
            }
        }
        Err(PythonError::Pool(
            "connection pool retry loop exited unexpectedly".to_string(),
        ))
    }

    fn try_begin_raw_query_once(&self, expr: &str) -> Result<SyncRawQuery, PythonError> {
        let payload = encode_sync_query(expr).map_err(crate::errors::map_transport_error)?;
        let pool = self.pool()?;
        let mut connection = pool
            .get()
            .map_err(|error| PythonError::Pool(error.to_string()))?;
        match begin_streaming_frame_over(&mut *connection, &payload) {
            Ok((header_bytes, remaining_body)) => {
                let header = match response_header(header_bytes) {
                    Ok(header) => header,
                    Err(error) => {
                        connection.mark_broken();
                        return Err(error);
                    }
                };
                Ok((
                    header,
                    header_bytes,
                    remaining_body,
                    Box::new(SyncPooledConnectionLease::new(connection)),
                ))
            }
            Err(error) => {
                connection.mark_broken();
                Err(crate::errors::map_transport_error(error))
            }
        }
    }

    fn begin_raw_query(&self, expr: &str) -> Result<SyncRawQuery, PythonError> {
        for attempt in 0..=self.options.retry_attempts_value() {
            match self.try_begin_raw_query_once(expr) {
                Ok(response) => return Ok(response),
                Err(error) => {
                    if attempt == self.options.retry_attempts_value()
                        || !retryable_pool_error(&error)
                    {
                        return Err(error);
                    }
                    sync_backoff(self.options.retry_backoff_ms_value());
                }
            }
        }
        Err(PythonError::Pool(
            "connection pool retry loop exited unexpectedly".to_string(),
        ))
    }
}

async fn async_request(
    conn: &mut AsyncPooledTransport,
    payload: &[u8],
) -> Result<Vec<u8>, PythonError> {
    match qroissant_transport::request_frame_streaming_over_async(conn, payload).await {
        Ok(frame) => Ok(frame),
        Err(error) => {
            conn.mark_broken();
            Err(crate::errors::map_transport_error(error))
        }
    }
}

async fn async_request_value(
    conn: &mut AsyncPooledTransport,
    payload: &[u8],
) -> Result<Value, PythonError> {
    match request_value_pipelined_over_async(conn, payload).await {
        Ok(value) => Ok(value),
        Err(error) => {
            conn.mark_broken();
            Err(crate::errors::map_transport_error(error))
        }
    }
}

async fn async_send_query(
    conn: &mut AsyncPooledTransport,
    query: &str,
) -> Result<Vec<u8>, PythonError> {
    let payload = encode_sync_query(query).map_err(crate::errors::map_transport_error)?;
    async_request(conn, &payload).await
}

async fn async_validate_connection(
    conn: &mut AsyncPooledTransport,
    healthcheck_query: Option<&str>,
) -> Result<(), PoolBackendError> {
    if conn.is_broken() {
        return Err(PoolBackendError::Message(
            "pooled connection is broken".to_string(),
        ));
    }
    if let Some(query) = healthcheck_query {
        let response = async_send_query(conn, query).await?;
        validate_success_response(&response)?;
    }
    Ok(())
}

#[derive(Clone)]
struct AsyncTransportManager {
    endpoint: Endpoint,
    healthcheck_query: Option<String>,
}

impl AsyncManageConnection for AsyncTransportManager {
    type Connection = AsyncPooledTransport;
    type Error = PoolBackendError;

    async fn connect(&self) -> Result<Self::Connection, Self::Error> {
        connect_async_transport(&self.endpoint)
            .await
            .map(AsyncPooledTransport::new)
            .map_err(Into::into)
    }

    async fn is_valid(&self, conn: &mut Self::Connection) -> Result<(), Self::Error> {
        async_validate_connection(conn, self.healthcheck_query.as_deref()).await
    }

    fn has_broken(&self, conn: &mut Self::Connection) -> bool {
        conn.is_broken()
    }
}

struct AsyncPooledConnectionLease {
    connection: Option<bb8::PooledConnection<'static, AsyncTransportManager>>,
    reusable: bool,
}

impl AsyncPooledConnectionLease {
    fn new(connection: bb8::PooledConnection<'static, AsyncTransportManager>) -> Self {
        Self {
            connection: Some(connection),
            reusable: false,
        }
    }

    fn mark_inner_broken(&mut self) {
        if let Some(connection) = self.connection.as_mut() {
            connection.mark_broken();
        }
    }
}

impl AsyncRead for AsyncPooledConnectionLease {
    fn poll_read(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &mut ReadBuf<'_>,
    ) -> Poll<std::io::Result<()>> {
        let connection = self
            .connection
            .as_mut()
            .expect("active async pooled lease must hold a checked-out connection");
        Pin::new(&mut **connection).poll_read(cx, buf)
    }
}

impl AsyncWrite for AsyncPooledConnectionLease {
    fn poll_write(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &[u8],
    ) -> Poll<std::io::Result<usize>> {
        let connection = self
            .connection
            .as_mut()
            .expect("active async pooled lease must hold a checked-out connection");
        Pin::new(&mut **connection).poll_write(cx, buf)
    }

    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
        let connection = self
            .connection
            .as_mut()
            .expect("active async pooled lease must hold a checked-out connection");
        Pin::new(&mut **connection).poll_flush(cx)
    }

    fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
        let connection = self
            .connection
            .as_mut()
            .expect("active async pooled lease must hold a checked-out connection");
        Pin::new(&mut **connection).poll_shutdown(cx)
    }
}

impl AsyncStreamingLease for AsyncPooledConnectionLease {
    fn mark_reusable(&mut self) {
        self.reusable = true;
    }

    fn abandon(&mut self) {
        self.reusable = false;
        self.mark_inner_broken();
    }
}

impl Drop for AsyncPooledConnectionLease {
    fn drop(&mut self) {
        if !self.reusable {
            self.mark_inner_broken();
        }
    }
}

struct AsyncPoolState {
    pool: Option<bb8::Pool<AsyncTransportManager>>,
    closed: bool,
}

struct AsyncPoolCore {
    manager: AsyncTransportManager,
    options: PoolOptions,
    state: Mutex<AsyncPoolState>,
}

impl AsyncPoolCore {
    fn new(endpoint: Endpoint, options: PoolOptions) -> Result<Self, PythonError> {
        endpoint.validate().map_err(PythonError::Pool)?;
        Ok(Self {
            manager: AsyncTransportManager {
                endpoint,
                healthcheck_query: options.healthcheck_query_value().map(str::to_string),
            },
            options,
            state: Mutex::new(AsyncPoolState {
                pool: None,
                closed: false,
            }),
        })
    }

    async fn build_pool(&self) -> Result<bb8::Pool<AsyncTransportManager>, PoolBackendError> {
        let mut builder = bb8::Pool::builder()
            .max_size(self.options.max_size_value())
            .connection_timeout(Duration::from_millis(
                self.options.checkout_timeout_ms_value(),
            ))
            .test_on_check_out(self.options.test_on_checkout_value());
        builder = builder.min_idle(self.options.min_idle_value());
        builder = builder.idle_timeout(
            self.options
                .idle_timeout_ms_value()
                .map(Duration::from_millis),
        );
        builder = builder.max_lifetime(
            self.options
                .max_lifetime_ms_value()
                .map(Duration::from_millis),
        );
        builder
            .build(self.manager.clone())
            .await
            .map_err(|error| PoolBackendError::Message(error.to_string()))
    }

    async fn pool(&self) -> Result<bb8::Pool<AsyncTransportManager>, PythonError> {
        {
            let state = self
                .state
                .lock()
                .map_err(|_| PythonError::Pool("pool lock is poisoned".to_string()))?;
            if state.closed {
                return Err(PythonError::PoolClosed);
            }
            if let Some(pool) = &state.pool {
                return Ok(pool.clone());
            }
        }

        let built = self
            .build_pool()
            .await
            .map_err(|error| PythonError::Pool(error.to_string()))?;
        let mut state = self
            .state
            .lock()
            .map_err(|_| PythonError::Pool("pool lock is poisoned".to_string()))?;
        if state.closed {
            return Err(PythonError::PoolClosed);
        }
        if let Some(pool) = &state.pool {
            return Ok(pool.clone());
        }
        state.pool = Some(built.clone());
        Ok(built)
    }

    fn close(&self) -> Result<(), PythonError> {
        let mut state = self
            .state
            .lock()
            .map_err(|_| PythonError::Pool("pool lock is poisoned".to_string()))?;
        state.closed = true;
        state.pool.take();
        Ok(())
    }

    fn metrics(&self) -> Result<PoolMetrics, PythonError> {
        let state = self
            .state
            .lock()
            .map_err(|_| PythonError::Pool("pool lock is poisoned".to_string()))?;
        let (connections, idle_connections, initialized) = state
            .pool
            .as_ref()
            .map(|pool| {
                let pool_state = pool.state();
                (pool_state.connections, pool_state.idle_connections, true)
            })
            .unwrap_or((0, 0, false));
        Ok(PoolMetrics::new_native(
            connections,
            idle_connections,
            self.options.max_size_value(),
            self.options.min_idle_value(),
            initialized,
            state.closed,
        ))
    }

    async fn prewarm(&self) -> Result<PoolMetrics, PythonError> {
        let target = self.options.min_idle_value().unwrap_or(1);
        if target == 0 {
            return self.metrics();
        }
        let pool = self.pool().await?;
        let mut connections = Vec::with_capacity(target as usize);
        for _ in 0..target {
            connections.push(
                pool.get()
                    .await
                    .map_err(|error| PythonError::Pool(error.to_string()))?,
            );
        }
        drop(connections);
        self.metrics()
    }

    async fn try_query_value_once(&self, expr: &str) -> Result<Value, PythonError> {
        let payload = encode_sync_query(expr).map_err(crate::errors::map_transport_error)?;
        let pool = self.pool().await?;
        let mut connection = pool
            .get()
            .await
            .map_err(|error| PythonError::Pool(error.to_string()))?;
        async_request_value(&mut connection, &payload).await
    }

    async fn query_value(&self, expr: &str) -> Result<Value, PythonError> {
        for attempt in 0..=self.options.retry_attempts_value() {
            match self.try_query_value_once(expr).await {
                Ok(value) => return Ok(value),
                Err(error) => {
                    if attempt == self.options.retry_attempts_value()
                        || !retryable_pool_error(&error)
                    {
                        return Err(error);
                    }
                    async_backoff(self.options.retry_backoff_ms_value()).await;
                }
            }
        }
        Err(PythonError::Pool(
            "connection pool retry loop exited unexpectedly".to_string(),
        ))
    }

    async fn try_begin_raw_query_once(&self, expr: &str) -> Result<AsyncPoolRawQuery, PythonError> {
        let payload = encode_sync_query(expr).map_err(crate::errors::map_transport_error)?;
        let pool = self.pool().await?;
        let mut connection = pool
            .get_owned()
            .await
            .map_err(|error| PythonError::Pool(error.to_string()))?;
        match begin_streaming_frame_over_async(&mut *connection, &payload).await {
            Ok((header_bytes, remaining_body)) => {
                let header = match response_header(header_bytes) {
                    Ok(header) => header,
                    Err(error) => {
                        connection.mark_broken();
                        return Err(error);
                    }
                };
                Ok((
                    header,
                    header_bytes,
                    remaining_body,
                    AsyncPooledConnectionLease::new(connection),
                ))
            }
            Err(error) => {
                connection.mark_broken();
                Err(crate::errors::map_transport_error(error))
            }
        }
    }

    async fn begin_raw_query(&self, expr: &str) -> Result<AsyncPoolRawQuery, PythonError> {
        for attempt in 0..=self.options.retry_attempts_value() {
            match self.try_begin_raw_query_once(expr).await {
                Ok(response) => return Ok(response),
                Err(error) => {
                    if attempt == self.options.retry_attempts_value()
                        || !retryable_pool_error(&error)
                    {
                        return Err(error);
                    }
                    async_backoff(self.options.retry_backoff_ms_value()).await;
                }
            }
        }
        Err(PythonError::Pool(
            "connection pool retry loop exited unexpectedly".to_string(),
        ))
    }
}

#[pyclass(module = "qroissant")]
pub struct Connection {
    core: Arc<ConnectionCore>,
    options: Option<DecodeOptions>,
}

#[pymethods]
impl Connection {
    #[new]
    #[pyo3(signature = (endpoint, *, options=None))]
    fn new(
        py: Python<'_>,
        endpoint: PyRef<'_, Endpoint>,
        options: Option<PyRef<'_, DecodeOptions>>,
    ) -> PyResult<Self> {
        let endpoint = endpoint.clone();
        let core = py
            .detach(move || ConnectionCore::connect(&endpoint))
            .map_err(to_py_err)?;
        Ok(Self {
            core: Arc::new(core),
            options: options.map(|value| value.clone()),
        })
    }

    fn __enter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> {
        slf
    }

    fn __exit__(
        &self,
        _exc_type: Option<&Bound<'_, PyAny>>,
        _exc_val: Option<&Bound<'_, PyAny>>,
        _exc_tb: Option<&Bound<'_, PyAny>>,
    ) -> PyResult<()> {
        self.core.close().map_err(to_py_err)
    }

    fn close(&self, py: Python<'_>) -> PyResult<()> {
        py.detach(|| self.core.close()).map_err(to_py_err)
    }

    #[pyo3(signature = (expr, /, *, raw=false, decode=None))]
    fn query(
        &self,
        py: Python<'_>,
        expr: String,
        raw: bool,
        decode: Option<PyRef<'_, DecodeOptions>>,
    ) -> PyResult<Py<PyAny>> {
        if raw {
            let (header, header_bytes, remaining_body, lease) = py
                .detach(|| self.core.begin_raw_query(&expr))
                .map_err(to_py_err)?;
            let response = RawResponse::streaming(header, header_bytes, remaining_body, lease);
            return Ok(Py::new(py, response)?.into_any());
        }

        let payload = py
            .detach(|| self.core.query_frame(&expr))
            .map_err(to_py_err)?;
        let options = decode.as_deref().or(self.options.as_ref());
        let (value, opts) =
            decode_core_value(bytes::Bytes::from(payload), options).map_err(to_py_err)?;
        core_value_to_python_with_opts(py, value, opts)
    }
}

#[pyclass(module = "qroissant")]
pub struct AsyncConnection {
    core: Arc<AsyncConnectionCore>,
    options: Option<DecodeOptions>,
}

#[pymethods]
impl AsyncConnection {
    #[new]
    #[pyo3(signature = (endpoint, *, options=None))]
    fn new(
        endpoint: PyRef<'_, Endpoint>,
        options: Option<PyRef<'_, DecodeOptions>>,
    ) -> PyResult<Self> {
        Ok(Self {
            core: Arc::new(AsyncConnectionCore::new(endpoint.clone()).map_err(to_py_err)?),
            options: options.map(|value| value.clone()),
        })
    }

    fn __aenter__<'py>(slf: PyRef<'py, Self>, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
        let core = slf.core.clone();
        let options = slf.options.clone();
        async_self(py, move |py| Py::new(py, Self { core, options }))
    }

    fn __aexit__<'py>(
        &self,
        py: Python<'py>,
        _exc_type: Option<&Bound<'_, PyAny>>,
        _exc_val: Option<&Bound<'_, PyAny>>,
        _exc_tb: Option<&Bound<'_, PyAny>>,
    ) -> PyResult<Bound<'py, PyAny>> {
        let core = self.core.clone();
        future_into_py(py, async move {
            core.close().await.map_err(to_py_err)?;
            Ok(false)
        })
    }

    fn close<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
        let core = self.core.clone();
        future_into_py(py, async move {
            core.close().await.map_err(to_py_err)?;
            Ok(())
        })
    }

    #[pyo3(signature = (expr, /, *, raw=false, decode=None))]
    fn query<'py>(
        &self,
        py: Python<'py>,
        expr: String,
        raw: bool,
        decode: Option<DecodeOptions>,
    ) -> PyResult<Bound<'py, PyAny>> {
        let core = self.core.clone();
        let default_options = self.options.clone();
        future_into_py(py, async move {
            if raw {
                let (header, header_bytes, remaining_body, lease) =
                    core.begin_raw_query(&expr).await.map_err(to_py_err)?;
                return Python::attach(|py| {
                    let response = RawResponse::streaming(
                        header,
                        header_bytes,
                        remaining_body,
                        Box::new(BlockingAsyncBridge::new(lease)),
                    );
                    Py::new(py, response.into_async()).map(|value| value.into_any())
                });
            }

            let value = core.query_value(&expr).await.map_err(to_py_err)?;
            let options = decode.as_ref().or(default_options.as_ref());
            let proj_opts = crate::serde::decode_options_to_proj_opts(options);
            Python::attach(|py| core_value_to_python_with_opts(py, value, proj_opts))
        })
    }
}

#[pyclass(module = "qroissant")]
pub struct Pool {
    core: Arc<SyncPoolCore>,
    options: Option<DecodeOptions>,
}

#[pymethods]
impl Pool {
    #[new]
    #[pyo3(signature = (endpoint, *, options=None, pool=None))]
    fn new(
        endpoint: PyRef<'_, Endpoint>,
        options: Option<PyRef<'_, DecodeOptions>>,
        pool: Option<PyRef<'_, PoolOptions>>,
    ) -> PyResult<Self> {
        let core = SyncPoolCore::new(
            endpoint.clone(),
            pool.map(|value| value.clone()).unwrap_or_default(),
        )
        .map_err(to_py_err)?;
        Ok(Self {
            core: Arc::new(core),
            options: options.map(|value| value.clone()),
        })
    }

    fn __enter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> {
        slf
    }

    fn __exit__(
        &self,
        _exc_type: Option<&Bound<'_, PyAny>>,
        _exc_val: Option<&Bound<'_, PyAny>>,
        _exc_tb: Option<&Bound<'_, PyAny>>,
    ) -> PyResult<()> {
        self.core.close().map_err(to_py_err)
    }

    fn close(&self, py: Python<'_>) -> PyResult<()> {
        py.detach(|| self.core.close()).map_err(to_py_err)
    }

    fn prewarm(&self, py: Python<'_>) -> PyResult<PoolMetrics> {
        py.detach(|| self.core.prewarm()).map_err(to_py_err)
    }

    fn metrics(&self) -> PyResult<PoolMetrics> {
        self.core.metrics().map_err(to_py_err)
    }

    #[pyo3(signature = (expr, /, *, raw=false, decode=None))]
    fn query(
        &self,
        py: Python<'_>,
        expr: String,
        raw: bool,
        decode: Option<PyRef<'_, DecodeOptions>>,
    ) -> PyResult<Py<PyAny>> {
        if raw {
            let (header, header_bytes, remaining_body, lease) = py
                .detach(|| self.core.begin_raw_query(&expr))
                .map_err(to_py_err)?;
            let response = RawResponse::streaming(header, header_bytes, remaining_body, lease);
            return Ok(Py::new(py, response)?.into_any());
        }

        let payload = py
            .detach(|| self.core.query_frame(&expr))
            .map_err(to_py_err)?;
        let options = decode.as_deref().or(self.options.as_ref());
        let (value, opts) =
            decode_core_value(bytes::Bytes::from(payload), options).map_err(to_py_err)?;
        core_value_to_python_with_opts(py, value, opts)
    }
}

#[pyclass(module = "qroissant")]
pub struct AsyncPool {
    core: Arc<AsyncPoolCore>,
    options: Option<DecodeOptions>,
}

#[pymethods]
impl AsyncPool {
    #[new]
    #[pyo3(signature = (endpoint, *, options=None, pool=None))]
    fn new(
        endpoint: PyRef<'_, Endpoint>,
        options: Option<PyRef<'_, DecodeOptions>>,
        pool: Option<PyRef<'_, PoolOptions>>,
    ) -> PyResult<Self> {
        let core = AsyncPoolCore::new(
            endpoint.clone(),
            pool.map(|value| value.clone()).unwrap_or_default(),
        )
        .map_err(to_py_err)?;
        Ok(Self {
            core: Arc::new(core),
            options: options.map(|value| value.clone()),
        })
    }

    fn __aenter__<'py>(slf: PyRef<'py, Self>, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
        let core = slf.core.clone();
        let options = slf.options.clone();
        async_self(py, move |py| Py::new(py, Self { core, options }))
    }

    fn __aexit__<'py>(
        &self,
        py: Python<'py>,
        _exc_type: Option<&Bound<'_, PyAny>>,
        _exc_val: Option<&Bound<'_, PyAny>>,
        _exc_tb: Option<&Bound<'_, PyAny>>,
    ) -> PyResult<Bound<'py, PyAny>> {
        let core = self.core.clone();
        future_into_py(py, async move {
            core.close().map_err(to_py_err)?;
            Ok(false)
        })
    }

    fn close<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
        let core = self.core.clone();
        future_into_py(py, async move {
            core.close().map_err(to_py_err)?;
            Ok(())
        })
    }

    fn prewarm<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
        let core = self.core.clone();
        future_into_py(py, async move {
            let metrics = core.prewarm().await.map_err(to_py_err)?;
            Python::attach(|py| Py::new(py, metrics).map(|value| value.into_any()))
        })
    }

    fn metrics<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
        let core = self.core.clone();
        future_into_py(py, async move {
            let metrics = core.metrics().map_err(to_py_err)?;
            Python::attach(|py| Py::new(py, metrics).map(|value| value.into_any()))
        })
    }

    #[pyo3(signature = (expr, /, *, raw=false, decode=None))]
    fn query<'py>(
        &self,
        py: Python<'py>,
        expr: String,
        raw: bool,
        decode: Option<DecodeOptions>,
    ) -> PyResult<Bound<'py, PyAny>> {
        let core = self.core.clone();
        let default_options = self.options.clone();
        future_into_py(py, async move {
            if raw {
                let (header, header_bytes, remaining_body, lease) =
                    core.begin_raw_query(&expr).await.map_err(to_py_err)?;
                return Python::attach(|py| {
                    let response = RawResponse::streaming(
                        header,
                        header_bytes,
                        remaining_body,
                        Box::new(BlockingAsyncBridge::new(lease)),
                    );
                    Py::new(py, response.into_async()).map(|value| value.into_any())
                });
            }

            let value = core.query_value(&expr).await.map_err(to_py_err)?;
            let options = decode.as_ref().or(default_options.as_ref());
            let proj_opts = crate::serde::decode_options_to_proj_opts(options);
            Python::attach(|py| core_value_to_python_with_opts(py, value, proj_opts))
        })
    }
}

pub fn register(module: &Bound<'_, PyModule>) -> PyResult<()> {
    module.add_class::<Connection>()?;
    module.add_class::<AsyncConnection>()?;
    module.add_class::<Pool>()?;
    module.add_class::<AsyncPool>()?;
    Ok(())
}
