use std::fmt;
use std::io::Read;
use std::sync::Arc;
use std::sync::Mutex;
use std::sync::MutexGuard;

use pyo3::buffer::PyBuffer;
use pyo3::prelude::*;
use pyo3::types::PyAny;
use pyo3::types::PyBytes;
use pyo3::types::PyModule;
use pyo3_async_runtimes::tokio::future_into_py;
use qroissant_core::HEADER_LEN;
use qroissant_core::MessageHeader as CoreMessageHeader;
use tokio::io::AsyncRead;
use tokio::io::AsyncReadExt;
use tokio::io::AsyncWrite;
use tokio::task::spawn_blocking;

use crate::serde::decode_core_value;
use crate::types::Compression;
use crate::types::DecodeOptions;
use crate::types::Encoding;
use crate::types::MessageHeader;
use crate::types::MessageType;
use crate::values::core_value_to_python_with_opts;

pub(crate) trait SyncRawLease: Read + Send {
    fn mark_reusable(&mut self);
    fn abandon(&mut self);
}

pub(crate) trait AsyncStreamingLease: AsyncRead + AsyncWrite + Send + Unpin {
    fn mark_reusable(&mut self);
    fn abandon(&mut self);
}

pub(crate) struct BlockingAsyncBridge<T> {
    inner: T,
    handle: tokio::runtime::Handle,
}

impl<T> BlockingAsyncBridge<T> {
    pub(crate) fn new(inner: T) -> Self {
        Self {
            inner,
            handle: tokio::runtime::Handle::current(),
        }
    }
}

impl<T> Read for BlockingAsyncBridge<T>
where
    T: AsyncStreamingLease,
{
    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
        let handle = self.handle.clone();
        let inner = &mut self.inner;
        let fut = async move { inner.read(buf).await };
        if tokio::runtime::Handle::try_current().is_ok() {
            tokio::task::block_in_place(|| handle.block_on(fut))
        } else {
            handle.block_on(fut)
        }
    }
}

impl<T> SyncRawLease for BlockingAsyncBridge<T>
where
    T: AsyncStreamingLease,
{
    fn mark_reusable(&mut self) {
        self.inner.mark_reusable();
    }

    fn abandon(&mut self) {
        self.inner.abandon();
    }
}

fn closed_raw_response_error() -> PyErr {
    pyo3::exceptions::PyValueError::new_err("I/O operation on closed qroissant raw response")
}

fn backend_lock_error() -> PyErr {
    pyo3::exceptions::PyRuntimeError::new_err("qroissant raw response state is poisoned")
}

fn unsupported_seek_error() -> PyErr {
    pyo3::exceptions::PyOSError::new_err(
        "qroissant raw streaming responses are forward-only and do not support seek()",
    )
}

fn readonly_buffer_error() -> PyErr {
    pyo3::exceptions::PyTypeError::new_err("readinto() requires a writable buffer")
}

fn non_contiguous_buffer_error() -> PyErr {
    pyo3::exceptions::PyTypeError::new_err("readinto() requires a C-contiguous buffer")
}

#[derive(Debug)]
enum RawReadError {
    Closed,
    BackendPoisoned,
    PartiallyConsumed,
    Io(std::io::Error),
}

impl From<std::io::Error> for RawReadError {
    fn from(error: std::io::Error) -> Self {
        Self::Io(error)
    }
}

impl fmt::Display for RawReadError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Closed => write!(f, "raw response is closed"),
            Self::BackendPoisoned => write!(f, "raw response backend is poisoned"),
            Self::PartiallyConsumed => {
                write!(f, "raw response has already been partially consumed")
            }
            Self::Io(error) => error.fmt(f),
        }
    }
}

fn raw_read_error_to_py(error: RawReadError) -> PyErr {
    match error {
        RawReadError::Closed => closed_raw_response_error(),
        RawReadError::BackendPoisoned => backend_lock_error(),
        RawReadError::PartiallyConsumed => pyo3::exceptions::PyValueError::new_err(
            "cannot decode a partially consumed raw response",
        ),
        RawReadError::Io(error) => PyErr::from(error),
    }
}

fn extract_writable_contiguous_u8_buffer(payload: &Bound<'_, PyAny>) -> PyResult<PyBuffer<u8>> {
    let buffer = PyBuffer::<u8>::get(payload)?;
    if buffer.readonly() {
        return Err(readonly_buffer_error());
    }
    if !buffer.is_c_contiguous() {
        return Err(non_contiguous_buffer_error());
    }
    Ok(buffer)
}

enum RawResponseBackend {
    Buffered {
        payload: Vec<u8>,
        position: usize,
    },
    Streaming {
        header_bytes: [u8; HEADER_LEN],
        header_position: usize,
        remaining_body: usize,
        position: usize,
        lease: Option<Box<dyn SyncRawLease>>,
    },
    Closed,
}

struct RawResponseState {
    header: MessageHeader,
    backend: RawResponseBackend,
}

impl RawResponseState {
    fn streaming_remaining_total(header_position: usize, remaining_body: usize) -> usize {
        (HEADER_LEN - header_position) + remaining_body
    }

    fn finalize_stream(lease: &mut Option<Box<dyn SyncRawLease>>, reusable: bool) {
        if let Some(mut lease) = lease.take() {
            if reusable {
                lease.mark_reusable();
            } else {
                lease.abandon();
            }
        }
    }

    fn close_backend(backend: &mut RawResponseBackend) {
        let backend = std::mem::replace(backend, RawResponseBackend::Closed);
        match backend {
            RawResponseBackend::Buffered { .. } | RawResponseBackend::Closed => {}
            RawResponseBackend::Streaming {
                remaining_body,
                header_position,
                mut lease,
                ..
            } => {
                let reusable =
                    Self::streaming_remaining_total(header_position, remaining_body) == 0;
                Self::finalize_stream(&mut lease, reusable);
            }
        }
    }

    fn read_streaming_into(
        header_bytes: &[u8; HEADER_LEN],
        header_position: &mut usize,
        remaining_body: &mut usize,
        position: &mut usize,
        lease: &mut Option<Box<dyn SyncRawLease>>,
        out: &mut [u8],
    ) -> Result<usize, RawReadError> {
        let total_remaining = Self::streaming_remaining_total(*header_position, *remaining_body);
        if total_remaining == 0 {
            Self::finalize_stream(lease, true);
            return Ok(0);
        }

        let target = out.len().min(total_remaining);
        let mut filled = 0_usize;
        let header_copied = if *header_position < HEADER_LEN && filled < target {
            let available = HEADER_LEN - *header_position;
            let to_copy = (target - filled).min(available);
            out[..to_copy]
                .copy_from_slice(&header_bytes[*header_position..*header_position + to_copy]);
            *header_position += to_copy;
            filled += to_copy;
            to_copy
        } else {
            0
        };

        if filled < target {
            while filled < target {
                let lease_ref = lease
                    .as_mut()
                    .expect("streaming raw responses must hold an active lease");
                let read = lease_ref.read(&mut out[filled..target])?;
                if read == 0 {
                    Self::finalize_stream(lease, false);
                    return Err(std::io::Error::from(std::io::ErrorKind::UnexpectedEof).into());
                }
                filled += read;
            }
        }

        let body_bytes = filled.saturating_sub(header_copied);
        if body_bytes != 0 {
            *remaining_body = remaining_body.saturating_sub(body_bytes);
        }
        *position = position.saturating_add(filled);

        if Self::streaming_remaining_total(*header_position, *remaining_body) == 0 {
            Self::finalize_stream(lease, true);
        }
        Ok(filled)
    }
}

fn header_from_payload(payload: &[u8]) -> PyResult<MessageHeader> {
    if payload.len() < HEADER_LEN {
        return Ok(MessageHeader::new_native(
            Encoding::LittleEndian,
            MessageType::Response,
            Compression::Uncompressed,
            payload.len(),
        ));
    }
    let header = CoreMessageHeader::parse(payload)
        .map_err(|error| pyo3::exceptions::PyValueError::new_err(error.to_string()))?;
    Ok(MessageHeader::from(header))
}

#[pyclass(module = "qroissant")]
pub struct RawResponse {
    state: Arc<Mutex<RawResponseState>>,
}

impl RawResponse {
    fn lock_state_result(&self) -> Result<MutexGuard<'_, RawResponseState>, RawReadError> {
        self.state.lock().map_err(|_| RawReadError::BackendPoisoned)
    }

    fn lock_state(&self) -> PyResult<MutexGuard<'_, RawResponseState>> {
        self.lock_state_result().map_err(raw_read_error_to_py)
    }

    fn ensure_open(backend: &RawResponseBackend) -> PyResult<()> {
        if matches!(backend, RawResponseBackend::Closed) {
            return Err(closed_raw_response_error());
        }
        Ok(())
    }

    fn ensure_open_result(backend: &RawResponseBackend) -> Result<(), RawReadError> {
        if matches!(backend, RawResponseBackend::Closed) {
            return Err(RawReadError::Closed);
        }
        Ok(())
    }

    pub(crate) fn buffered(payload: Vec<u8>) -> PyResult<Self> {
        let header = header_from_payload(&payload)?;
        Ok(Self {
            state: Arc::new(Mutex::new(RawResponseState {
                header,
                backend: RawResponseBackend::Buffered {
                    payload,
                    position: 0,
                },
            })),
        })
    }

    pub(crate) fn streaming(
        header: MessageHeader,
        header_bytes: [u8; HEADER_LEN],
        remaining_body: usize,
        lease: Box<dyn SyncRawLease>,
    ) -> Self {
        Self {
            state: Arc::new(Mutex::new(RawResponseState {
                header,
                backend: RawResponseBackend::Streaming {
                    header_bytes,
                    header_position: 0,
                    remaining_body,
                    position: 0,
                    lease: Some(lease),
                },
            })),
        }
    }

    pub(crate) fn into_async(self) -> AsyncRawResponse {
        let this = std::mem::ManuallyDrop::new(self);
        // SAFETY: `ManuallyDrop` suppresses `RawResponse::drop`, so it is safe
        // to move the owned `Arc` into the async wrapper without closing the
        // underlying raw-response state.
        let state = unsafe { std::ptr::read(&this.state) };
        AsyncRawResponse { state }
    }

    fn materialize_result(&self) -> Result<Vec<u8>, RawReadError> {
        let position = {
            let state = self.lock_state_result()?;
            Self::ensure_open_result(&state.backend)?;
            match &state.backend {
                RawResponseBackend::Buffered { position, .. }
                | RawResponseBackend::Streaming { position, .. } => *position,
                RawResponseBackend::Closed => {
                    unreachable!("closed raw responses are handled above")
                }
            }
        };
        if position != 0 {
            return Err(RawReadError::PartiallyConsumed);
        }
        self.read_owned_result(None)
    }

    fn read_owned_result(&self, size: Option<isize>) -> Result<Vec<u8>, RawReadError> {
        let mut state = self.lock_state_result()?;
        Self::ensure_open_result(&state.backend)?;
        match &mut state.backend {
            RawResponseBackend::Buffered { payload, position } => {
                if *position >= payload.len() {
                    return Ok(Vec::new());
                }
                let remaining = payload.len() - *position;
                let to_read = match size {
                    Some(size) if size >= 0 => remaining.min(size as usize),
                    _ => remaining,
                };
                let start = *position;
                let end = start + to_read;
                *position = end;
                Ok(payload[start..end].to_vec())
            }
            RawResponseBackend::Streaming {
                header_bytes,
                header_position,
                remaining_body,
                position,
                lease,
            } => {
                let total_remaining =
                    RawResponseState::streaming_remaining_total(*header_position, *remaining_body);
                let target = match size {
                    Some(size) if size >= 0 => total_remaining.min(size as usize),
                    _ => total_remaining,
                };
                let mut out = vec![0_u8; target];
                match RawResponseState::read_streaming_into(
                    header_bytes,
                    header_position,
                    remaining_body,
                    position,
                    lease,
                    &mut out,
                ) {
                    Ok(filled) => {
                        out.truncate(filled);
                        Ok(out)
                    }
                    Err(error) => {
                        state.backend = RawResponseBackend::Closed;
                        Err(error)
                    }
                }
            }
            RawResponseBackend::Closed => Err(RawReadError::Closed),
        }
    }

    fn read_into_result(&self, out: &mut [u8]) -> Result<usize, RawReadError> {
        let mut state = self.lock_state_result()?;
        Self::ensure_open_result(&state.backend)?;
        match &mut state.backend {
            RawResponseBackend::Buffered { payload, position } => {
                if *position >= payload.len() {
                    return Ok(0);
                }
                let remaining = payload.len() - *position;
                let to_read = remaining.min(out.len());
                let start = *position;
                let end = start + to_read;
                out[..to_read].copy_from_slice(&payload[start..end]);
                *position = end;
                Ok(to_read)
            }
            RawResponseBackend::Streaming {
                header_bytes,
                header_position,
                remaining_body,
                position,
                lease,
            } => match RawResponseState::read_streaming_into(
                header_bytes,
                header_position,
                remaining_body,
                position,
                lease,
                out,
            ) {
                Ok(filled) => Ok(filled),
                Err(error) => {
                    state.backend = RawResponseBackend::Closed;
                    Err(error)
                }
            },
            RawResponseBackend::Closed => Err(RawReadError::Closed),
        }
    }
}

impl Drop for RawResponse {
    fn drop(&mut self) {
        // Clean up even if the mutex is poisoned (panic in another thread).
        let mut state = match self.state.lock() {
            Ok(guard) => guard,
            Err(poisoned) => poisoned.into_inner(),
        };
        RawResponseState::close_backend(&mut state.backend);
    }
}

#[pymethods]
impl RawResponse {
    #[new]
    fn new(payload: Vec<u8>) -> PyResult<Self> {
        Self::buffered(payload)
    }

    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.close()
    }

    #[getter]
    fn closed(&self) -> bool {
        self.state
            .lock()
            .map(|state| matches!(state.backend, RawResponseBackend::Closed))
            .unwrap_or(true)
    }

    #[getter]
    fn header(&self) -> PyResult<MessageHeader> {
        let state = self.lock_state()?;
        Ok(state.header.clone())
    }

    fn close(&self) -> PyResult<()> {
        let mut state = self.lock_state()?;
        RawResponseState::close_backend(&mut state.backend);
        Ok(())
    }

    fn readable(&self) -> bool {
        !self.closed()
    }

    fn seekable(&self) -> bool {
        self.state
            .lock()
            .map(|state| matches!(state.backend, RawResponseBackend::Buffered { .. }))
            .unwrap_or(false)
    }

    #[pyo3(signature = (size=None))]
    fn read<'py>(&self, py: Python<'py>, size: Option<isize>) -> PyResult<Bound<'py, PyBytes>> {
        let bytes = py
            .detach(|| self.read_owned_result(size))
            .map_err(raw_read_error_to_py)?;
        Ok(PyBytes::new(py, &bytes))
    }

    #[pyo3(signature = (size=None))]
    fn read1<'py>(&self, py: Python<'py>, size: Option<isize>) -> PyResult<Bound<'py, PyBytes>> {
        self.read(py, size)
    }

    fn readinto(&self, py: Python<'_>, buffer: &Bound<'_, PyAny>) -> PyResult<usize> {
        let writable = extract_writable_contiguous_u8_buffer(buffer)?;
        let len = writable.len_bytes();
        if len == 0 {
            let mut empty = [];
            return py
                .detach(|| self.read_into_result(&mut empty))
                .map_err(raw_read_error_to_py);
        }
        let ptr = writable.buf_ptr() as usize;
        py.detach(move || {
            let ptr = ptr as *mut u8;
            // SAFETY: the writable Python buffer outlives this detached call and
            // the slice length is bounded by the exported buffer length.
            let slice = unsafe { std::slice::from_raw_parts_mut(ptr, len) };
            self.read_into_result(slice)
        })
        .map_err(raw_read_error_to_py)
    }

    fn readinto1(&self, py: Python<'_>, buffer: &Bound<'_, PyAny>) -> PyResult<usize> {
        self.readinto(py, buffer)
    }

    fn tell(&self) -> PyResult<usize> {
        let state = self.lock_state()?;
        Self::ensure_open(&state.backend)?;
        match &state.backend {
            RawResponseBackend::Buffered { position, .. }
            | RawResponseBackend::Streaming { position, .. } => Ok(*position),
            RawResponseBackend::Closed => Err(closed_raw_response_error()),
        }
    }

    #[pyo3(signature = (offset, whence=0))]
    fn seek(&self, offset: i64, whence: i32) -> PyResult<usize> {
        let mut state = self.lock_state()?;
        Self::ensure_open(&state.backend)?;
        match &mut state.backend {
            RawResponseBackend::Buffered { payload, position } => {
                let base = match whence {
                    0 => 0_i64,
                    1 => i64::try_from(*position).map_err(|_| {
                        pyo3::exceptions::PyOverflowError::new_err(
                            "raw response position exceeds supported seek range",
                        )
                    })?,
                    2 => i64::try_from(payload.len()).map_err(|_| {
                        pyo3::exceptions::PyOverflowError::new_err(
                            "raw response length exceeds supported seek range",
                        )
                    })?,
                    _ => {
                        return Err(pyo3::exceptions::PyValueError::new_err(format!(
                            "invalid seek whence value {whence}; expected 0, 1, or 2"
                        )));
                    }
                };
                let position_i64 = base.checked_add(offset).ok_or_else(|| {
                    pyo3::exceptions::PyOverflowError::new_err(
                        "raw response seek position overflowed",
                    )
                })?;
                if position_i64 < 0 {
                    return Err(pyo3::exceptions::PyValueError::new_err(
                        "negative seek position is not allowed",
                    ));
                }
                *position = usize::try_from(position_i64).map_err(|_| {
                    pyo3::exceptions::PyOverflowError::new_err(
                        "raw response seek position overflowed",
                    )
                })?;
                Ok(*position)
            }
            RawResponseBackend::Streaming { .. } => Err(unsupported_seek_error()),
            RawResponseBackend::Closed => Err(closed_raw_response_error()),
        }
    }

    #[pyo3(signature = (*, options=None))]
    fn decode(&self, py: Python<'_>, options: Option<&DecodeOptions>) -> PyResult<Py<PyAny>> {
        let payload = py
            .detach(|| self.materialize_result())
            .map_err(raw_read_error_to_py)?;
        let (value, opts) =
            decode_core_value(bytes::Bytes::from(payload), options)
                .map_err(crate::errors::to_py_err)?;
        core_value_to_python_with_opts(py, value, opts)
    }

    fn __repr__(&self) -> String {
        match self.state.lock() {
            Ok(state) => match &state.backend {
                RawResponseBackend::Buffered { payload, position } => format!(
                    "RawResponse(mode='buffered', len={}, position={}, closed=false)",
                    payload.len(),
                    position
                ),
                RawResponseBackend::Streaming {
                    header_position,
                    remaining_body,
                    position,
                    ..
                } => format!(
                    "RawResponse(mode='streaming', remaining={}, position={}, closed=false)",
                    RawResponseState::streaming_remaining_total(*header_position, *remaining_body),
                    position
                ),
                RawResponseBackend::Closed => "RawResponse(mode='closed', closed=true)".to_string(),
            },
            Err(_) => "RawResponse(mode='poisoned', closed=true)".to_string(),
        }
    }
}

#[pyclass(module = "qroissant")]
pub struct AsyncRawResponse {
    state: Arc<Mutex<RawResponseState>>,
}

#[pymethods]
impl AsyncRawResponse {
    fn __aenter__<'py>(slf: PyRef<'py, Self>, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
        let state = slf.state.clone();
        future_into_py(py, async move {
            Python::attach(|py| Py::new(py, Self { state }).map(|value| value.into_any()))
        })
    }

    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 state = self.state.clone();
        future_into_py(py, async move {
            let mut state = state.lock().map_err(|_| backend_lock_error())?;
            RawResponseState::close_backend(&mut state.backend);
            Ok(false)
        })
    }

    #[getter]
    fn closed(&self) -> bool {
        self.state
            .lock()
            .map(|state| matches!(state.backend, RawResponseBackend::Closed))
            .unwrap_or(true)
    }

    #[getter]
    fn header(&self) -> PyResult<MessageHeader> {
        let state = self.state.lock().map_err(|_| backend_lock_error())?;
        Ok(state.header.clone())
    }

    fn close<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
        let state = self.state.clone();
        future_into_py(py, async move {
            let mut state = state.lock().map_err(|_| backend_lock_error())?;
            RawResponseState::close_backend(&mut state.backend);
            Ok(())
        })
    }

    #[pyo3(signature = (size=None))]
    fn read<'py>(&self, py: Python<'py>, size: Option<isize>) -> PyResult<Bound<'py, PyAny>> {
        let raw = RawResponse {
            state: self.state.clone(),
        };
        future_into_py(py, async move {
            let bytes = spawn_blocking(move || raw.read_owned_result(size))
                .await
                .map_err(|error| pyo3::exceptions::PyRuntimeError::new_err(error.to_string()))?
                .map_err(raw_read_error_to_py)?;
            Python::attach(|py| Ok(PyBytes::new(py, &bytes).unbind().into_any()))
        })
    }

    #[pyo3(signature = (size=None))]
    fn read1<'py>(&self, py: Python<'py>, size: Option<isize>) -> PyResult<Bound<'py, PyAny>> {
        self.read(py, size)
    }

    fn readinto<'py>(&self, py: Python<'py>, buffer: Py<PyAny>) -> PyResult<Bound<'py, PyAny>> {
        let state = self.state.clone();
        future_into_py(py, async move {
            Python::attach(|py| {
                let buffer = buffer.bind(py);
                let writable = extract_writable_contiguous_u8_buffer(buffer)?;
                let len = writable.len_bytes();
                if len == 0 {
                    return Ok(0);
                }
                let ptr = writable.buf_ptr() as usize;
                let raw = RawResponse {
                    state: state.clone(),
                };
                drop(writable);
                let read = py
                    .detach(move || {
                        let ptr = ptr as *mut u8;
                        // SAFETY: the writable Python buffer outlives this detached call and
                        // the slice length is bounded by the exported buffer length.
                        let slice = unsafe { std::slice::from_raw_parts_mut(ptr, len) };
                        raw.read_into_result(slice)
                    })
                    .map_err(raw_read_error_to_py)?;
                Ok(read)
            })
        })
    }

    fn readinto1<'py>(&self, py: Python<'py>, buffer: Py<PyAny>) -> PyResult<Bound<'py, PyAny>> {
        self.readinto(py, buffer)
    }

    #[pyo3(signature = (*, options=None))]
    fn decode<'py>(
        &self,
        py: Python<'py>,
        options: Option<DecodeOptions>,
    ) -> PyResult<Bound<'py, PyAny>> {
        let raw = RawResponse {
            state: self.state.clone(),
        };
        future_into_py(py, async move {
            let payload = spawn_blocking(move || raw.materialize_result())
                .await
                .map_err(|error| pyo3::exceptions::PyRuntimeError::new_err(error.to_string()))?
                .map_err(raw_read_error_to_py)?;
            let (value, opts) =
                decode_core_value(bytes::Bytes::from(payload), options.as_ref())
                    .map_err(crate::errors::to_py_err)?;
            Python::attach(|py| core_value_to_python_with_opts(py, value, opts))
        })
    }
}

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