cataclysm 0.3.0-beta.3

use tokio::{
    sync::Semaphore,
    net::{TcpListener, TcpStream}
    //io::AsyncWriteExt
};
use bytes::Buf;
use crate::metafunctions::callback::PipelineKind;
use crate::{
    Stream,
    Branch, Shared, Additional, Cors, branch::PureBranch, Pipeline, Error, session::SessionCreator,
    http::{Request, Response, Method}
};
use std::sync::{Arc};

// Default max connections for the server
const MAX_CONNECTIONS: usize = 2_000;
const RESPONSE_CHUNK_SIZE: usize = 4_096;
const READ_CHUNK_SIZE: usize = 8_192;

/// Builder pattern for the server structure
///
/// It is the main method for building a server and configuring certain behaviour
pub struct ServerBuilder<T> {
    branch: Branch<T>,
    shared: Option<Shared<T>>,
    session_creator: Option<Arc<Box<dyn SessionCreator>>>,
    log_string: Option<String>,
    cors: Option<Cors>,
    max_connections: usize,
    timeout: std::time::Duration
}

impl<T: Sync + Send> ServerBuilder<T> {
    /// Creates a new server from a given branch
    ///
    /// ```rust,no_run
    /// # use cataclysm::{ServerBuilder, Branch, http::{Method, Response}};
    /// let branch: Branch<()> = Branch::new("/").with(Method::Get.to(|| async {Response::ok().body("Ok!")}));
    /// let mut server_builder = ServerBuilder::new(branch);
    /// // ...
    /// ```
    pub fn new(branch: Branch<T>) -> ServerBuilder<T> {
        ServerBuilder {
            branch,
            shared: None,
            session_creator: None,
            log_string: None,
            cors: None,
            max_connections: MAX_CONNECTIONS,
            timeout: std::time::Duration::from_millis(15_000)
        }
    }

    /// Declare some information to be shared with the [Shared](crate::Shared) extractor
    ///
    /// ```rust,no_run
    /// use cataclysm::{Server, Branch, Shared, http::{Response, Method, Path}};
    /// 
    /// // Receives a string, and concatenates the shared suffix
    /// async fn index(path: Path<(String,)>, shared: Shared<String>) -> Response {
    ///     let (prefix,) = path.into_inner();
    ///     Response::ok().body(format!("{}{}", prefix, *shared))
    /// }
    /// 
    /// #[tokio::main]
    /// async fn main() {
    ///     // We create our tree structure
    ///     let branch = Branch::new("/{:prefix}").with(Method::Get.to(index));
    ///     // We create a server with the given tree structure
    ///     let server = Server::builder(branch).share("!!!".into()).build().unwrap();
    ///     // And we launch it on the following address
    ///     server.run("127.0.0.1:8000").await.unwrap();
    /// }
    /// ```
    ///
    /// If you intend to share a mutable variable, consider using rust's [Mutex](https://doc.rust-lang.org/std/sync/struct.Mutex.html), as the shared value is already inside an [Arc](https://doc.rust-lang.org/std/sync/struct.Arc.html).
    pub fn share(mut self, shared: T) -> ServerBuilder<T> {
        self.shared = Some(Shared::new(shared));
        self
    }

    /// Sets a session creator for the Session extractor to work
    ///
    /// ```rust,no_run
    /// use cataclysm::{Server, session::{Session, CookieSession}, Branch, Shared, http::{Response, Method, Path}};
    /// 
    /// async fn index(mut session: Session) -> Response {
    ///     // the session will be empty if the signature was invalid
    ///     // ... do something with the session
    ///     // apply changes to response
    ///     session.apply(Response::ok())
    /// }
    /// 
    /// #[tokio::main]
    /// async fn main() {
    ///     // We create our tree structure
    ///     let branch: Branch<()> = Branch::new("/").with(Method::Get.to(index));
    ///     // We create a server with the given tree structure
    ///     let server = Server::builder(branch).session_creator(
    ///         CookieSession::new() // Default cookie session implementation
    ///     ).build().unwrap();
    ///     // And we launch it on the following address
    ///     server.run("127.0.0.1:8000").await.unwrap();
    /// }
    /// ```
    ///
    /// If no secret is provided, a random key will be used (generated by ring).
    pub fn session_creator<A: 'static + SessionCreator>(mut self, session_creator: A) -> Self {
        self.session_creator = Some(Arc::new(Box::new(session_creator)));
        self
    }

    /// Sets a log string, to log information per call
    ///
    /// ```rust,no_run
    /// # use cataclysm::{Server, Branch, Shared, http::{Response, Method, Path}};
    /// // Tree structure
    /// let branch: Branch<()> = Branch::new("/").with(Method::Get.to(|| async {Response::ok()}));
    /// // Now we configure the server
    /// let server = Server::builder(branch).log_format("[%M %P] %S, from %A").build().unwrap();
    /// ```
    ///
    /// The list of available format elements are the following
    /// 
    /// * `%M`: Method from the request
    /// * `%P`: Path from the request
    /// * `%S`: Status from the response
    /// * `%A`: Socket address and port from the connection
    /// (more data to be added soon)
    pub fn log_format<A: Into<String>>(mut self, log_string: A) -> Self {
        self.log_string = Some(log_string.into());
        self
    }

    /// Adds the cors "middleware"
    ///
    /// ```rust,no_run
    /// # use cataclysm::{Server, Branch, CorsBuilder, http::{Response, Method}};
    /// // Tree structure
    /// let branch: Branch<()> = Branch::new("/").with(Method::Get.to(|| async {Response::ok()}));
    /// // Now we configure the server
    /// let server = Server::builder(branch)
    ///     .cors(CorsBuilder::new()
    ///         .origin("https://fake.domain")
    ///         .max_age(600)
    ///         .build().unwrap()
    ///     ).build().unwrap();
    /// ```
    pub fn cors(mut self, cors: Cors) -> Self {
        self.cors = Some(cors);
        self
    }

    /// Sets up a maximum number of connections for the server to be dealt with
    ///
    /// ```rust,no_run
    /// # use cataclysm::{Server, Branch, http::{Response, Method}};
    /// // Tree structure
    /// let branch: Branch<()> = Branch::new("/").with(Method::Get.to(|| async {Response::ok()}));
    /// // Now we configure the server
    /// let server = Server::builder(branch).max_connections(10_000).build().unwrap();
    /// ```
    pub fn max_connections(mut self, n: usize) -> Self {
        self.max_connections = n;
        self
    }

    /// Sets up a custom timeout for http requests to be finished
    ///
    /// ```rust,no_run
    /// # use cataclysm::{Server, Branch, http::{Response, Method}};
    /// use std::time::Duration;
    /// // Tree structure
    /// let branch: Branch<()> = Branch::new("/").with(Method::Get.to(|| async {Response::ok()}));
    /// // Now we configure the server
    /// let server = Server::builder(branch).timeout(Duration::from_millis(5_000)).build().unwrap();
    /// ```
    pub fn timeout(mut self, duration: std::time::Duration) -> Self {
        self.timeout = duration;
        self
    }

    /// Builds the server
    ///
    /// ```rust,no_run
    /// use cataclysm::{Server, Branch, Shared, http::{Response, Method, Path}};
    /// 
    /// // Receives a string, and concatenates the shared suffix
    /// async fn index() -> Response {
    ///     Response::ok().body("Hello")
    /// }
    /// 
    /// #[tokio::main]
    /// async fn main() {
    ///     // We create our tree structure
    ///     let branch: Branch<()> = Branch::new("/").with(Method::Get.to(index));
    ///     // We create a server with the given tree structure
    ///     let server = Server::builder(branch).build().unwrap();
    ///     // And we launch it on the following address
    ///     server.run("127.0.0.1:8000").await.unwrap();
    /// }
    /// ```
    pub fn build(self) -> Result<Arc<Server<T>>, Error> {
        Ok(Arc::new(Server {
            pure_branch: Arc::new(self.branch.purify()),
            additional: Arc::new(Additional {
                shared: self.shared,
                session_creator: self.session_creator
            }),
            log_string: Arc::new(self.log_string),
            cors: Arc::new(self.cors),
            max_connections: Arc::new(Semaphore::new(self.max_connections)),
            timeout: Arc::new(self.timeout)
        }))
    }
}

/// Http Server instance
///
/// The Server structure hosts all the information to successfully process each call
pub struct Server<T> {
    pure_branch: Arc<PureBranch<T>>,
    additional: Arc<Additional<T>>,
    log_string: Arc<Option<String>>,
    cors: Arc<Option<Cors>>,
    max_connections: Arc<Semaphore>,
    timeout: Arc<std::time::Duration>
}

impl<T: 'static + Sync + Send> Server<T> {
    // Short for ServerBuilder's `new` function.
    pub fn builder(branch: Branch<T>) -> ServerBuilder<T> {
        ServerBuilder::new(branch)
    }

    pub async fn run<S: AsRef<str>>(self: &Arc<Self>, socket: S) -> Result<(), Error> {
        let listener = TcpListener::bind(socket.as_ref()).await.map_err(|e| Error::Io(e))?;

        log::info!("Cataclysm ongoing \u{26c8}");
        #[cfg(feature = "full_log")]
        log::warn!("using the `full_log` feature might impact performance");
        // We need a fused future for the select macro
        tokio::select! {
            _ = async {
                loop {
                    // We lock the loop until one permit becomes available
                    self.max_connections.acquire().await.unwrap().forget();
                    
                    match listener.accept().await {
                        Ok((socket, addr)) => {
                            #[cfg(feature = "full_log")]
                            log::trace!("socket connection accepted");
                            let server = Arc::clone(self);
                            
                            tokio::spawn(async move {
                                tokio::select! {
                                    res = server.dispatch(socket, addr) => match res {
                                        Ok(_) => (),
                                        Err(e) => {
                                            log::error!("{}", e);
                                        }
                                    },
                                    _ = tokio::time::sleep(*server.timeout) => {
                                        log::debug!("timeout for http response");
                                    }
                                }
                                // We set up back the permits
                                server.max_connections.add_permits(1);
                            });
                        },
                        Err(e) => {
                            log::error!("{}", e);
                        }
                    }
                }
            } => (),
            _ = tokio::signal::ctrl_c() => {
                log::info!("Shutting down server");
            }
        };
        Ok(())
    }

    /// Deals with the read part of the socket stream
    async fn dispatch_read(socket: &TcpStream, addr: std::net::SocketAddr) -> Result<Option<Vec<u8>>, Error> {
        let mut request_bytes = Vec::with_capacity(READ_CHUNK_SIZE);
        let mut expected_length = None;
        let mut header_size = 0;
        let mut request = None;
        // First we read
        loop {
            socket.readable().await.map_err(|e| Error::Io(e))?;
            
            // being stored in the async task.
            let mut buf = [0; READ_CHUNK_SIZE];

            // Try to read data, this may still fail with `WouldBlock`
            // if the readiness event is a false positive.
            match socket.try_read(&mut buf) {
                Ok(0) => {
                    break
                },
                Ok(n) => {
                    request_bytes.extend_from_slice(&buf[0..n]);

                    if request.is_none() {
                        request = match Request::parse(request_bytes.clone(), addr) {
                            Ok(r) => {
                                // We check if we need to give a continue 100
                                if r.headers.get("Expect").map(|h| h.get(0).map(|ih| ih == "100-continue")).flatten().unwrap_or(false) {
                                    // We send it
                                    Server::<T>::dispatch_write(&socket, Response::r#continue()).await?;
                                }

                                // We check now if there is a content size hint
                                expected_length = r.headers.get("Content-Length").or_else(|| r.headers.get("content-length")).map(|cl| cl.get(0).map(|v| v.parse::<usize>().ok())).flatten().flatten();
                                #[cfg(feature = "full_log")]
                                log::trace!("expecting to read {:?} bytes in request", expected_length);
                                header_size = r.header_size;
                                Some(r)
                            },
                            Err(_e) => {
                                #[cfg(feature = "full_log")]
                                log::debug!("{}", _e);
                                Server::<T>::dispatch_write(&socket, Response::bad_request()).await?;
                                return Ok(None)
                            }
                        };
                    }

                    // And now we check if, given the hint, we need to act upon.
                    if let Some(expected_length) = &expected_length {
                        if *expected_length > request_bytes.len() - header_size {
                            continue;
                        } else {
                            break;
                        }
                    } else {
                        break;
                    }
                },
                Err(ref e) if e.kind() == std::io::ErrorKind::WouldBlock => {
                    continue
                }
                Err(e) => return Err(Error::Io(e))
            }
        }
        Ok(Some(request_bytes))
    }

    async fn dispatch_write(socket: &TcpStream, mut response: Response) -> Result<(), Error> {
        let serialized_response = response.serialize();
        let mut chunks_iter = serialized_response.chunks(RESPONSE_CHUNK_SIZE);
        #[cfg(feature = "full_log")]
        log::trace!("writting {} chunks of maximum {} bytes each", chunks_iter.len(), RESPONSE_CHUNK_SIZE);
        // We check the first chunk
        let mut current_chunk = match chunks_iter.next() {
            Some(v) => v,
            None => return Ok(()) // Zero length response
        };
        loop {
            // Wait for the socket to be writable
            socket.writable().await.map_err(|e| Error::Io(e))?;
    
            // Try to write data, this may still fail with `WouldBlock`
            // if the readiness event is a false positive.        
            match socket.try_write(&current_chunk) {
                Ok(n) => {
                    if n != current_chunk.remaining() {
                        // There are some bytes still to be written in this chunk
                        #[cfg(feature = "full_log")]
                        log::debug!("incomplete chunk, trying to serve remaining bytes ({}/{})", current_chunk.len(), RESPONSE_CHUNK_SIZE);
                        current_chunk.advance(n);
                        continue;
                    } else {
                        current_chunk = match chunks_iter.next() {
                            Some(v) => v,
                            None => return Ok(())
                        }
                    }
                }
                Err(ref e) if e.kind() == tokio::io::ErrorKind::WouldBlock => {
                    continue;
                }
                Err(e) => break Err(Error::Io(e))
            }
        }
    }

    async fn dispatch(self: &Arc<Self>, socket: TcpStream, addr: std::net::SocketAddr) -> Result<(), Error> {
        // let mut second_part = false;
        let request_bytes = match Server::<T>::dispatch_read(&socket, addr).await? {
            Some(b) => b,
            None => return Ok(())
        };

        let stream = Stream::new(socket);

        let mut request =match Request::parse(request_bytes.clone(), addr) {
            Ok(r) => r,
            Err(_e) => {
                #[cfg(feature = "full_log")]
                log::debug!("{}", _e);
                stream.response(Response::bad_request()).await?;
                return Ok(())
            }
        };

        if let Some(cors) = &*self.cors {
            if request.method == Method::Options {
                if let Some(supported_methods) = self.pure_branch.supported_methods(request.url().path()) {
                    stream.response(cors.preflight(&request, &supported_methods)).await?;
                } // If the method is not options, it will anyways return a not-found
            }
        }

        request.addr = addr;

        // The method will take the request, and modify particularly the "variable count" variable
        let mut response = match self.pure_branch.pipeline(&mut request) {
            Some(pipeline_kind) => {
                match pipeline_kind {
                    PipelineKind::NormalPipeline(pipeline) => {
                        #[cfg(feature = "full_log")]
                        log::trace!("found path {} with method {}", request.url, request.method);
                        match pipeline {
                            Pipeline::Layer(func, pipeline_layer) => func(request.clone(), pipeline_layer, self.additional.clone()),
                            Pipeline::Core(core_fn) => core_fn(request.clone(), self.additional.clone())
                        }.await
                    },
                    #[cfg(feature = "stream")]
                    PipelineKind::StreamPipeline(pipeline) => {
                        #[cfg(feature = "full_log")]
                        log::trace!("found stream path {}", request.url);
                        pipeline(request.clone(), self.additional.clone(), stream).await;
                        return Ok(())
                    }
                }
            },
            None => {
                #[cfg(feature = "full_log")]
                log::trace!("path {} not found, with method {}", request.url, request.method);
                Response::not_found()
            }
        };

        // Cors validation, not as an actual pipeline layer
        if let Some(cors) = &*self.cors {
            cors.apply(&request, &mut response);
        }

        if let Some(log_string) = &*self.log_string {
            log::info!("{}", log_string.replace("%M", request.method.to_str()).replace("%P", &request.url().path()).replace("%A", &format!("{}", addr)).replace("%S", &format!("{}", response.status.0)));
        }

        stream.response(response).await?;
        //Server::<T>::dispatch_write(&stream.into(), response).await?;
        Ok(())
    }
}