mcstatusface/src/http.rs

use std::{collections::HashMap, io::{BufRead, BufReader, Result, Write}, net::{SocketAddr, TcpListener, TcpStream}};

use chrono::Local;

use crate::ThreadPool;

#[derive(Clone, Copy)]
pub enum StatusCode {
    OK,
    BadRequest,
    NotFound,
    InternalServerError,
    // ImATeapot,
}

impl StatusCode {
    pub fn code(&self) -> u16 {
        match self {
            StatusCode::OK => 200,
            StatusCode::BadRequest => 400,
            StatusCode::NotFound => 404,
            // StatusCode::ImATeapot => 418,
            StatusCode::InternalServerError => 500,
        }
    }

    pub fn reason(&self) -> &str {
        match self {
            StatusCode::OK => "OK",
            StatusCode::BadRequest => "Bad Request",
            StatusCode::NotFound => "Not Found",
            // StatusCode::ImATeapot => "I'm a teapot",
            StatusCode::InternalServerError => "Internal Server Error",
        }
    }
}

type HttpHandlerFunc = fn(&Request, Response) -> Result<StatusCode>;

pub struct Request<'a> {
    stream: &'a TcpStream,
    path: &'a str,
    method: &'a str,
    version: &'a str,
    headers: HashMap<&'a str, &'a str>,
    query: HashMap<&'a str, &'a str>,
    body: Option<String>,
}

impl<'a> Request<'a> {
    pub fn new(stream: &'a TcpStream, lines: &'a Vec<String>) -> Result<Request<'a>> {
        let request_line = lines[0].as_str();
        let request_line_split: Vec<&str> = request_line.split(" ").collect();
        let method = request_line_split[0];
        let path = request_line_split[1];
        let version = request_line_split[2];

        let mut query: HashMap<&'a str, &'a str> = HashMap::new();
        match path.split_once("?") {
            Some((_, query_string)) => {
                let query_splits: Vec<&'a str> = query_string.split("&").collect();
                for pair in query_splits {
                    match pair.split_once("=") {
                        Some((name, value)) => {
                            query.insert(name, value);
                        }
                        None => { continue; }
                    }
                }
            }
            None => {},
        }

        let mut headers: HashMap<&'a str, &'a str> = HashMap::new();
        if lines.len() > 1 {
            let mut i: usize = 1;
            loop {
                if i >= lines.len() { break; }
                let line = &lines[i];
                if line.len() == 0 || !line.contains(":") { break; }
                let (name, value) = line.split_once(":").unwrap();
                headers.insert(name, value.trim());
                i += 1;
            }
        }

        let mut body: Option<String> = None;
        if lines.len() > headers.len() + 2 &&
            (method == "POST" || method == "PUT") &&
            lines[headers.len() + 1] == "\r\n"
        {
            body = Some(lines[headers.len() + 2..].join("\n"));
        }

        Ok(Request {
            stream,
            path,
            method,
            version,
            headers,
            query,
            body,
        })
    }

    pub fn address(&self) -> Result<SocketAddr> {
        self.stream.peer_addr()
    }
    pub fn path(&self) -> &'a str {
        self.path
    }
    pub fn method(&self) -> &'a str {
        self.method
    }
    pub fn version(&self) -> &'a str {
        self.version
    }
    pub fn body(&self) -> &Option<String> {
        &self.body
    }
    pub fn headers(&self) -> &HashMap<&'a str, &'a str> {
        &self.headers
    }
    pub fn query(&self) -> &HashMap<&'a str, &'a str> {
        &self.query
    }
}

pub struct Response<'a> {
    stream: &'a TcpStream,
    status: StatusCode,
    headers: HashMap<&'a str, String>,
    body: Option<String>,
}

impl<'a> Response<'a> {
    pub fn new(stream: &'a TcpStream) -> Response<'a> {
        Response {
            stream,
            status: StatusCode::OK,
            headers: HashMap::from([
                         ("Server", "mcstatusface".to_string()),
                         ("Content-Type", "text/plain".to_string()),
            ]),
            body: None,
        }
    }

    pub fn status(&mut self, status: StatusCode) {
        self.status = status;
    }
    pub fn headers(&self) -> &HashMap<&'a str, String> {
        &self.headers
    }
    pub fn set_header(&mut self, name: &'a str, value: String) {
        self.headers.insert(name, value);
    }
    pub fn body(&mut self, body: String) {
        self.body = Some(body);
    }

    pub fn send(&mut self) -> Result<StatusCode> {
        // let mut len: usize = 0;
        self.stream.write(format!("HTTP/1.1 {} {}\r\n", self.status.code(), self.status.reason()).as_bytes()).unwrap();

        let mut content_length: usize = 0;
        if self.body.is_some() {
            content_length = self.body.as_ref().unwrap().len();
        }
        self.set_header("Content-Length", content_length.to_string());
        for (name, value) in &self.headers {
            self.stream.write(format!("{name}: {value}\r\n").as_bytes()).unwrap();
        }

        if self.body.is_some() {
            self.stream.write("\r\n".as_bytes()).unwrap();
            self.stream.write(self.body.as_ref().unwrap().as_bytes()).unwrap();
        }
        Ok(self.status)
    }
}

pub struct HttpServer {
    address: String,
    port: u16,
    max_connections: usize,
}

impl HttpServer {
    pub fn new(address: String, max_connections: usize) -> HttpServer {
        let mut _address = address.clone();
        let mut _port: u16 = 8080;
        match address.split_once(":") {
            Some((ip, port)) => {
                _address = ip.to_string();
                _port = port.parse::<u16>().expect(format!("Invalid port {}", port).as_str());
            }
            None => {}
        }
        HttpServer {
            address: _address,
            port: _port,
            max_connections,
        }
    }

    pub fn start(&self, handler: HttpHandlerFunc) -> Result<()> {
        let pool = ThreadPool::new(self.max_connections);
        let listener = TcpListener::bind(format!("{}:{}", self.address, self.port)).expect("Failed to bind to port");

        println!("Now listening on {}:{}", self.address, self.port);

        for stream in listener.incoming() {
            match stream {
                Ok(stream) => {
                    pool.execute(move || {
                        HttpServer::handle_client(&stream, handler);
                    });
                }
                Err(e) => {
                    eprintln!("Failed to handle incoming connection: {e}");
                }
            }
        }

        Ok(())
    }

    fn handle_client(stream: &TcpStream, handler: HttpHandlerFunc) {
        let buf_reader = BufReader::new(stream);
        let http_request: Vec<String> = buf_reader
            .lines()
            .map(|result| result.unwrap())
            .take_while(|line| !line.is_empty())
            .collect();

        let request = Request::new(stream, &http_request);
        if request.is_err() {
            eprintln!("Failed to process request: {}", request.err().unwrap());
            return;
        }
        let request = request.unwrap();

        let response = Response::new(stream);

        let start_date = Local::now();
        match handler(&request, response) {
            Ok(status) => {
                let end_date = Local::now();
                println!(
                    "[{}] {} {} {} - {} {} - {}ms - {} ({})",
                    start_date.format("%Y-%m-%d %H:%M:%S"),

                    request.method(),
                    request.path(),
                    request.version(),

                    status.code(),
                    status.reason(),

                    (end_date - start_date).num_milliseconds(),

                    request.headers().get("User-Agent").map_or("[]", |v| v),
                    request.address().unwrap().ip(),
                );
            }
            Err(e) => {
                eprintln!("Failed to handle request: {e}")
            }
        }
    }
}