//! Zero-copy-ish WSGI input stream (avoid BytesIO overhead)

use bytes::Bytes;
use pyo3::prelude::*;
use pyo3::types::PyBytes;
use std::sync::Arc;

#[pyclass]
pub struct WsgiInput {
    data: Arc<Bytes>,
    pos: usize,
}

impl WsgiInput {
    pub fn new(data: Arc<Bytes>) -> Self {
        Self { data, pos: 0 }
    }

    pub fn empty_input() -> Self {
        Self::new(Arc::new(Bytes::new()))
    }

    fn remaining(&self) -> usize {
        self.data.len().saturating_sub(self.pos)
    }

    fn take_bytes<'py>(&mut self, py: Python<'py>, size: usize) -> PyResult<Bound<'py, PyBytes>> {
        let end = (self.pos + size).min(self.data.len());
        let slice = &self.data[self.pos..end];
        self.pos = end;
        Ok(PyBytes::new(py, slice))
    }
}

#[pymethods]
impl WsgiInput {
    fn read<'py>(&mut self, py: Python<'py>, size: Option<isize>) -> PyResult<Bound<'py, PyBytes>> {
        let size = size.unwrap_or(-1);
        if size < 0 {
            let remaining = self.remaining();
            return self.take_bytes(py, remaining);
        }
        self.take_bytes(py, size as usize)
    }

    fn readline<'py>(
        &mut self,
        py: Python<'py>,
        size: Option<isize>,
    ) -> PyResult<Bound<'py, PyBytes>> {
        if self.pos >= self.data.len() {
            return Ok(PyBytes::new(py, &[]));
        }

        let mut end = self.data.len();
        if let Some(limit) = size {
            if limit >= 0 {
                end = (self.pos + limit as usize).min(end);
            }
        }

        // Find newline
        let slice = &self.data[self.pos..end];
        if let Some(idx) = slice.iter().position(|b| *b == b'\n') {
            let line_end = self.pos + idx + 1;
            let out = &self.data[self.pos..line_end];
            self.pos = line_end;
            return Ok(PyBytes::new(py, out));
        }

        // No newline, return up to end
        self.take_bytes(py, end.saturating_sub(self.pos))
    }

    fn readlines<'py>(
        &mut self,
        py: Python<'py>,
        hint: Option<isize>,
    ) -> PyResult<Vec<Bound<'py, PyBytes>>> {
        let mut lines = Vec::new();
        let mut total = 0usize;
        let limit = hint.unwrap_or(-1);

        while self.pos < self.data.len() {
            let line = self.readline(py, None)?;
            total += line.as_bytes().len();
            lines.push(line);
            if limit >= 0 && total >= limit as usize {
                break;
            }
        }
        Ok(lines)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_read_all_and_partial() {
        Python::attach(|py| {
            let mut input = WsgiInput::new(Arc::new(Bytes::from_static(b"abcdef")));

            let part = input.read(py, Some(2)).expect("read part");
            assert_eq!(part.as_bytes(), b"ab");

            let rest = input.read(py, None).expect("read rest");
            assert_eq!(rest.as_bytes(), b"cdef");
        });
    }

    #[test]
    fn test_readline_and_readlines() {
        Python::attach(|py| {
            let mut input = WsgiInput::new(Arc::new(Bytes::from_static(b"l1\nl2\nl3")));

            let line1 = input.readline(py, None).expect("line1");
            assert_eq!(line1.as_bytes(), b"l1\n");

            let lines = input.readlines(py, None).expect("readlines");
            assert_eq!(lines.len(), 2);
            assert_eq!(lines[0].as_bytes(), b"l2\n");
            assert_eq!(lines[1].as_bytes(), b"l3");
        });
    }
}
