use core::{
    mem::MaybeUninit,
    ops::{Deref, DerefMut},
    slice::memchr,
};

use alloc::{boxed::Box, vec::Vec};

use crate::error::EResult;

use super::{Read, Write};

// Read

enum ReadBufferBacking<'a> {
    Owned(Box<[MaybeUninit<u8>]>),
    Borrowed(&'a mut [MaybeUninit<u8>]),
}

pub struct ReadBuffer<'a> {
    data: ReadBufferBacking<'a>,
    position: usize,
    len: usize,
}

// Write

pub trait FileWriter: Write {
    fn is_line_buffered(&self) -> bool;
    fn capacity(&self) -> usize;
    fn pending_bytes(&self) -> usize;
}

pub struct UnbufferedWriter<W: Write> {
    inner: W,
}

pub struct FullBufferedWriter<W: Write> {
    inner: W,
    buffer: Vec<u8>,
}

pub struct LineBufferedWriter<W: Write> {
    inner: FullBufferedWriter<W>,
    need_flush: bool,
}

// Read impl

impl<'a> ReadBuffer<'a> {
    pub fn owned(capacity: usize) -> Self {
        Self {
            data: ReadBufferBacking::Owned(Box::new_uninit_slice(capacity)),
            position: 0,
            len: 0,
        }
    }

    pub fn borrowed(data: &'a mut [MaybeUninit<u8>]) -> Self {
        Self {
            data: ReadBufferBacking::Borrowed(data),
            position: 0,
            len: 0,
        }
    }

    pub fn fill_from<R: Read + ?Sized>(&mut self, source: &mut R) -> EResult<&[u8]> {
        let buffer = unsafe { self.data.as_slice_mut() };
        if self.position == self.len {
            let amount = match source.read(buffer) {
                EResult::Ok(0) => return EResult::Ok(&buffer[..0]),
                EResult::Ok(n) => n,
                EResult::Err(err) => return EResult::Err(err),
            };

            self.position = 0;
            self.len = amount;
        }

        EResult::Ok(&buffer[self.position..self.len])
    }

    pub fn consume(&mut self, amount: usize) {
        self.position = (self.position + amount).min(self.len);
    }
}

impl ReadBufferBacking<'_> {
    unsafe fn as_slice_mut(&mut self) -> &mut [u8] {
        self.deref_mut().assume_init_mut()
    }
}

impl Deref for ReadBufferBacking<'_> {
    type Target = [MaybeUninit<u8>];

    fn deref(&self) -> &Self::Target {
        match self {
            Self::Owned(owned) => owned.deref(),
            Self::Borrowed(borrowed) => borrowed.deref(),
        }
    }
}

impl DerefMut for ReadBufferBacking<'_> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        match self {
            Self::Owned(owned) => owned.deref_mut(),
            Self::Borrowed(borrowed) => borrowed.deref_mut(),
        }
    }
}

// Write impl

impl<W: Write> UnbufferedWriter<W> {
    pub fn new(inner: W) -> Self {
        Self { inner }
    }
}

impl<W: Write> Write for UnbufferedWriter<W> {
    fn write(&mut self, data: &[u8]) -> EResult<usize> {
        self.inner.write(data)
    }

    fn flush(&mut self) -> EResult<()> {
        EResult::Ok(())
    }
}

impl<W: Write> FileWriter for UnbufferedWriter<W> {
    fn pending_bytes(&self) -> usize {
        0
    }

    fn capacity(&self) -> usize {
        0
    }

    fn is_line_buffered(&self) -> bool {
        false
    }
}

impl<W: Write> FullBufferedWriter<W> {
    pub fn with_capacity(inner: W, capacity: usize) -> Self {
        Self {
            inner,
            buffer: Vec::with_capacity(capacity),
        }
    }
}

impl<W: Write> Write for FullBufferedWriter<W> {
    fn write(&mut self, mut data: &[u8]) -> EResult<usize> {
        if data.len() + self.buffer.len() >= self.buffer.capacity() {
            self.flush()?;
            assert_eq!(self.buffer.len(), 0);
        }

        let len = data.len();

        if len >= self.buffer.capacity() {
            // Write directly from `data` until it can be fit into the buffer
            let (amount, result) = self
                .inner
                .write_chunked(&data[..len - self.buffer.capacity()]);

            // Fail if write fails
            result?;

            data = &data[amount..];
        }

        // Store the data in the buffer
        assert!(data.len() <= self.buffer.capacity());
        self.buffer.extend_from_slice(data);
        EResult::Ok(len)
    }

    fn flush(&mut self) -> EResult<()> {
        let (amount, result) = self.inner.write_chunked(&self.buffer);
        if amount != 0 {
            self.buffer.drain(..amount);
        }
        result
    }
}

impl<W: Write> FileWriter for FullBufferedWriter<W> {
    fn pending_bytes(&self) -> usize {
        self.buffer.len()
    }

    fn capacity(&self) -> usize {
        self.buffer.capacity()
    }

    fn is_line_buffered(&self) -> bool {
        false
    }
}

impl<W: Write> LineBufferedWriter<W> {
    pub fn with_capacity(inner: W, capacity: usize) -> Self {
        Self {
            inner: FullBufferedWriter::with_capacity(inner, capacity),
            need_flush: false,
        }
    }
}

impl<W: Write> Write for LineBufferedWriter<W> {
    fn write(&mut self, data: &[u8]) -> EResult<usize> {
        if self.need_flush {
            self.flush()?;
        }

        let Some(nl_pos) = memchr::memchr(b'\n', data) else {
            return self.inner.write(data);
        };
        // Write up to the last newline and flush
        let amount = self.inner.write(&data[..nl_pos + 1])?;
        self.need_flush = true;
        if self.flush().is_err() || amount != nl_pos + 1 {
            return EResult::Ok(amount);
        }

        // Write the rest to the buffer
        match self.inner.write(&data[nl_pos + 1..]) {
            EResult::Ok(amount_to_buffer) => EResult::Ok(amount_to_buffer + amount),
            EResult::Err(_) => EResult::Ok(amount),
        }
    }

    fn flush(&mut self) -> EResult<()> {
        self.inner.flush()?;
        self.need_flush = false;
        EResult::Ok(())
    }
}

impl<W: Write> FileWriter for LineBufferedWriter<W> {
    fn pending_bytes(&self) -> usize {
        self.inner.pending_bytes()
    }

    fn capacity(&self) -> usize {
        self.inner.capacity()
    }

    fn is_line_buffered(&self) -> bool {
        true
    }
}
