/*
==--==--==--==--==--==--==--==--==--==--==--==--==--==--==--==--

Dia-Files

Copyright (C) 2019-2024  Anonymous

There are several releases over multiple years,
they are listed as ranges, such as: "2019-2024".

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU Lesser General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU Lesser General Public License for more details.

You should have received a copy of the GNU Lesser General Public License
along with this program.  If not, see <https://www.gnu.org/licenses/>.

::--::--::--::--::--::--::--::--::--::--::--::--::--::--::--::--
*/

use {
    std::io::{Error, ErrorKind, StdinLock},
    crate::Result,
};

#[cfg(not(feature="async-std"))]
use std::{
    fs::File,
    io::{self, BufReader, Read, Stdin},
    path::Path,
};

#[cfg(feature="async-std")]
use {
    core::{
        pin::Pin,
        task::{Context, Poll},
    },
    async_std::{
        fs::File,
        io::{self, BufReader, Read, ReadExt, Stdin},
        path::Path,
    },
};

#[cfg(test)]
mod tests;

macro_rules! read_to_end { ($self: ident, $capacity: ident) => {{
    let mut result = Vec::with_capacity(match $capacity {
        None => usize::MIN,
        Some(capacity) => match $self.limit.checked_sub($self.read) {
            None => return Err(err!()),
            Some(max) => match usize::try_from(max) {
                Ok(max) => max.min(capacity),
                Err(_) => capacity,
            },
        },
    });
    async_call!($self.inner.read_to_end(&mut result))?;
    Ok(result)
}}}

macro_rules! read { ($stream: ident, $limit: ident, $capacity: ident) => {{
    async_call!(Self::new($stream, $limit).read_to_end($capacity))
}}}

macro_rules! open_file { ($path: ident, $limit: ident) => {{
    let path = $path.as_ref();

    match async_call!(path.metadata())?.len() {
        size => if size > $limit {
            return Err(err!("File too large: {size} (limit: {limit})", limit=$limit));
        },
    };

    Ok(Self::new(BufReader::new(async_call!(File::open(path))?), $limit))
}}}

macro_rules! read_file { ($path: ident, $limit: ident) => {{
    let path = $path.as_ref();
    let mut limit = async_call!(Self::open_file(path, $limit))?;
    async_call!(limit.read_to_end(Some(async_call!(path.metadata())?.len().try_into().map_err(|_| err!())?)))
}}}

macro_rules! read_file_to_string { ($path: ident, $limit: ident) => {{
    let data = read_file!($path, $limit)?;
    match String::from_utf8(data) {
        Ok(s) => Ok(s),
        Err(err) => Err(Error::new(ErrorKind::InvalidInput, __!("{}", err.utf8_error()))),
    }
}}}

/// # Limit
///
/// This struct can be used for limiting data read from some source.
#[derive(Debug)]
pub struct Limit<R> {
    inner: R,
    read: u64,
    limit: u64,
}

impl<R> Limit<R> {

    /// # Makes new object
    pub fn new(inner: R, limit: u64) -> Self {
        Self {
            inner,
            read: u64::MIN,
            limit,
        }
    }

}

#[cfg(not(feature="async-std"))]
#[doc(cfg(not(feature="async-std")))]
impl<R> Limit<R> where R: Read {

    /// # Reads all bytes
    ///
    /// Be careful using this function (in case your source produces a lot of data).
    pub fn read_to_end(&mut self, capacity: Option<usize>) -> Result<Vec<u8>> {
        read_to_end!(self, capacity)
    }

    /// # Reads a whole stream
    pub fn read(stream: R, limit: u64, capacity: Option<usize>) -> Result<Vec<u8>> {
        read!(stream, limit, capacity)
    }

}

#[cfg(feature="async-std")]
#[doc(cfg(feature="async-std"))]
impl<R> Limit<R> where R: Unpin + Read {

    /// # Reads all bytes
    ///
    /// Be careful using this function (in case your source produces a lot of data).
    pub async fn read_to_end(&mut self, capacity: Option<usize>) -> Result<Vec<u8>> {
        read_to_end!(self, capacity)
    }

    /// # Reads a whole stream
    pub async fn read(stream: R, limit: u64, capacity: Option<usize>) -> Result<Vec<u8>> {
        read!(stream, limit, capacity)
    }

}

impl Limit<BufReader<File>> {

    /// # Opens a file for reading with a limit
    ///
    /// The file will be wrapped inside a [`BufReader`][struct:std/io/BufReader].
    ///
    /// [struct:std/io/BufReader]: https://doc.rust-lang.org/std/io/struct.BufReader.html
    #[cfg(not(feature="async-std"))]
    #[doc(cfg(not(feature="async-std")))]
    pub fn open_file<P>(path: P, limit: u64) -> Result<Self> where P: AsRef<Path> {
        open_file!(path, limit)
    }

    /// # Opens a file for reading with a limit
    ///
    /// The file will be wrapped inside a [`BufReader`][struct:std/io/BufReader].
    ///
    /// [struct:std/io/BufReader]: https://doc.rust-lang.org/std/io/struct.BufReader.html
    #[cfg(feature="async-std")]
    #[doc(cfg(feature="async-std"))]
    pub async fn open_file<P>(path: P, limit: u64) -> Result<Self> where P: AsRef<Path> {
        open_file!(path, limit)
    }

    /// # Reads a whole file
    #[cfg(not(feature="async-std"))]
    #[doc(cfg(not(feature="async-std")))]
    pub fn read_file<P>(path: P, limit: u64) -> Result<Vec<u8>> where P: AsRef<Path> {
        read_file!(path, limit)
    }

    /// # Reads a whole file
    #[cfg(feature="async-std")]
    #[doc(cfg(feature="async-std"))]
    pub async fn read_file<P>(path: P, limit: u64) -> Result<Vec<u8>> where P: AsRef<Path> {
        read_file!(path, limit)
    }

    /// # Reads a whole file to a string
    #[cfg(not(feature="async-std"))]
    #[doc(cfg(not(feature="async-std")))]
    pub fn read_file_to_string<P>(path: P, limit: u64) -> Result<String> where P: AsRef<Path> {
        read_file_to_string!(path, limit)
    }

    /// # Reads a whole file to a string
    #[cfg(feature="async-std")]
    #[doc(cfg(feature="async-std"))]
    pub async fn read_file_to_string<P>(path: P, limit: u64) -> Result<String> where P: AsRef<Path> {
        read_file_to_string!(path, limit)
    }

}

impl Limit<Stdin> {

    /// # Makes new object from standard input
    pub fn stdin(limit: u64) -> Self {
        Self::new(io::stdin(), limit)
    }

}

impl Limit<StdinLock<'_>> {

    /// # Locks standard input and makes new object from it
    pub fn lock_stdin(limit: u64) -> Self {
        Self::new(std::io::stdin().lock(), limit)
    }
}

macro_rules! check { ($self: ident) => {
    if $self.read > $self.limit {
        let err = Err(err!("Limit: {limit}, already read: {read}", limit=$self.limit, read=$self.read));
        #[cfg(feature="async-std")]
        let err = Poll::Ready(err);
        return err;
    }
}}

macro_rules! update_read { ($self: ident, $more: ident) => {{
    match u64::try_from($more).map_err(|_| err!("Cannot convert {more} into u64", more=$more))?.checked_add($self.read) {
        None => {
            let err = Err(err!("Failed: {read} + {more}", read=$self.read, more=$more));
            #[cfg(feature="async-std")]
            let err = Poll::Ready(err);
            return err;
        },
        Some(read) => $self.read = read,
    };
}}}

#[cfg(not(feature="async-std"))]
#[doc(cfg(not(feature="async-std")))]
impl<R> Read for Limit<R> where R: Read {

    fn read(&mut self, buf: &mut [u8]) -> Result<usize> {
        check!(self);
        let result = self.inner.read(buf)?;

        update_read!(self, result);
        check!(self);

        Ok(result)
    }

}

#[cfg(feature="async-std")]
#[doc(cfg(feature="async-std"))]
impl<R> Read for Limit<R> where R: Unpin + Read {

    fn poll_read(self: Pin<&mut Self>, context: &mut Context<'_>, buf: &mut [u8]) -> Poll<Result<usize>> {
        let limit = self.get_mut();

        check!(limit);
        let result = match Read::poll_read(Pin::new(&mut limit.inner), context, buf) {
            Poll::Ready(Ok(count)) => count,
            Poll::Ready(Err(err)) => return Poll::Ready(Err(err)),
            Poll::Pending => return Poll::Pending,
        };

        update_read!(limit, result);
        check!(limit);

        Poll::Ready(Ok(result))
    }

}
