use core::{ffi::c_char, fmt, ptr::NonNull};

use alloc::vec::Vec;

use crate::{
    allocator::c_alloc,
    error::{self, EResult},
    headers::errno::Errno,
    io::Write,
};

pub(super) struct FileWriter<'w, W: Write>(pub(super) &'w mut W);
pub(super) struct StringWriter {
    buffer: NonNull<c_char>,
    position: usize,
    capacity: usize,
}
pub(super) struct AllocatedStringWriter {
    buffer: Vec<u8>,
    limit: usize,
}

pub(super) trait FmtWriter: Write + fmt::Write {}

// File

impl<W: Write> fmt::Write for FileWriter<'_, W> {
    fn write_str(&mut self, s: &str) -> fmt::Result {
        self.0
            .write_all(s.as_bytes())
            .into_result(|_| fmt::Error, true)
    }
}

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

    fn flush(&mut self) -> EResult<()> {
        self.0.flush()
    }
}

impl<W: Write> FmtWriter for FileWriter<'_, W> {}

// String

impl StringWriter {
    pub fn new(buffer: NonNull<c_char>, capacity: usize) -> Self {
        Self {
            buffer,
            capacity,
            position: 0,
        }
    }

    pub fn finish(&mut self) {
        if self.position < self.capacity {
            unsafe { self.buffer.add(self.position).write(0) };
        }
    }
}

impl fmt::Write for StringWriter {
    fn write_str(&mut self, s: &str) -> fmt::Result {
        self.write(s.as_bytes()).into_result(|_| fmt::Error, true)?;
        Ok(())
    }
}

impl Write for StringWriter {
    fn write(&mut self, data: &[u8]) -> EResult<usize> {
        let count = core::cmp::min(self.capacity - self.position, data.len());
        if count > 0 {
            let dst = unsafe {
                core::slice::from_raw_parts_mut(
                    self.buffer.add(self.position).as_ptr().cast(),
                    count,
                )
            };
            dst.copy_from_slice(&data[..count]);
            self.position += count;
        }
        EResult::Ok(count)
    }

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

impl FmtWriter for StringWriter {}

// Allocated string

impl AllocatedStringWriter {
    pub fn new(limit: usize) -> Self {
        Self {
            buffer: Vec::new(),
            limit,
        }
    }

    pub fn finish(self) -> EResult<NonNull<c_char>> {
        assert!(self.buffer.len() < self.limit);
        let data = c_alloc(self.buffer.len() + 1, 1, false)?.cast::<c_char>();
        unsafe {
            let dst_slice =
                NonNull::slice_from_raw_parts(data.cast::<u8>(), self.buffer.len() + 1).as_mut();
            dst_slice[..self.buffer.len()].copy_from_slice(&self.buffer[..]);
            dst_slice[self.buffer.len()] = 0;
        }
        EResult::Ok(data)
    }
}

impl Write for AllocatedStringWriter {
    fn write(&mut self, data: &[u8]) -> EResult<usize> {
        let amount = (self.limit - self.buffer.len()).min(data.len());
        if self.buffer.try_reserve(amount).is_err() {
            return EResult::Err(Errno::ENOMEM);
        }
        self.buffer.extend_from_slice(&data[..amount]);
        EResult::Ok(data.len())
    }

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

impl fmt::Write for AllocatedStringWriter {
    fn write_str(&mut self, s: &str) -> fmt::Result {
        match self.write_all(s.as_bytes()) {
            EResult::Ok(()) => Ok(()),
            EResult::Err(e) => {
                error::set_errno(e);
                Err(fmt::Error)
            }
        }
    }
}

impl FmtWriter for AllocatedStringWriter {}
