router.rs

use crate::prelude::*;

#[derive(Debug, Clone)]
pub struct Route {
    pub methods: Vec<String>,
    pub path: String,
    pub handler: RouteHandler,
    pub redirect: Option<String>,
    pub directory_listing: bool,
    pub default_file: Option<String>,
    pub cgi_extension: Option<String>,
}

#[derive(Debug, Clone)]
pub enum RouteHandler {
    StaticFile(PathBuf),
    Directory(PathBuf),
    Function(fn(&HttpRequest, &SessionManager) -> HttpResponse),
    CGI(fn(&HttpRequest, &ErrorPage, Duration) -> Result<CGIProcess, HttpResponse>),
}

pub enum RouteResponse {
    Complete(HttpResponse),
    Pending(CGIProcess),
}

#[derive(Debug, Clone)]
pub struct VirtualServerRoutes {
    pub server_name: String,
    pub routes: Vec<Route>,
    pub default_file: String,
    pub error_pages: ErrorPage,
    pub max_body_size: usize,
}

#[derive(Debug, Clone)]
pub struct Router {
    pub virtual_servers: Vec<(String, u16, VirtualServerRoutes)>,
}

impl Router {
    pub fn new(config: &ServerConfig) -> Self {
        let mut virtual_servers = Vec::new();
        for server_config in &config.servers {
            let mut routes = Vec::new();
            if let Some(config_routes) = &server_config.routes {
                for cfg_route in config_routes {
                    let methods = cfg_route.methods.iter().map(|m| m.to_uppercase()).collect();
                    let handler = match &cfg_route.handler {
                        ConfigRouteHandler::File { file } => {
                            let path = PathBuf::from(file);
                            // Check if it's a directory (ends with / or actually exists as directory)
                            if file.ends_with('/') || (path.exists() && path.is_dir()) {
                                RouteHandler::Directory(path)
                            } else {
                                RouteHandler::StaticFile(path)
                            }
                        }
                        ConfigRouteHandler::Directory { directory } => {
                            RouteHandler::Directory(PathBuf::from(directory))
                        }
                        ConfigRouteHandler::Function => match cfg_route.cgi_extension {
                            Some(_) => RouteHandler::CGI(execute),
                            None => RouteHandler::Function(simple_function_handler),
                        },
                    };
                    routes.push(Route {
                        methods,
                        path: cfg_route.path.clone(),
                        handler,
                        redirect: cfg_route.redirect.clone(),
                        directory_listing: cfg_route.directory_listing.unwrap_or(false),
                        default_file: cfg_route.default_file.clone(),
                        cgi_extension: cfg_route.cgi_extension.clone(),
                    });
                }
            }
            virtual_servers.push((
                server_config.host.clone(),
                server_config.port,
                VirtualServerRoutes {
                    server_name: server_config.server_name.clone(),
                    routes,
                    default_file: server_config
                        .default_file
                        .clone()
                        .unwrap_or_else(|| "templates/index.html".to_string()),
                    error_pages: ErrorPage::new(&server_config.error_pages),
                    max_body_size: server_config.max_body_size,
                },
            ));
        }
        Self { virtual_servers }
    }

    pub fn route(
        &self,
        request: &HttpRequest,
        listener_host: &str,
        listener_port: u16,
        session_manager: &SessionManager,
        cgi_timeout: Duration,
    ) -> RouteResponse {
        let host_header = request
            .headers
            .get("host")
            .map(|h| h.as_str())
            .unwrap_or("");
        let host_header = host_header.split(':').next().unwrap_or(host_header);
        debug_log!(
            "Routing request with Host: {}, listener:                                                                                                                                                                                                            {}:{}",
            host_header,
            listener_host,
            listener_port
        );

        let virtual_server = self
            .virtual_servers
            .iter()
            .filter(|(host, port, _)| host == listener_host && *port == listener_port)
            .find(|(_, _, server)| server.server_name == host_header)
            .or_else(|| {
                debug_log!(
                    "Falling back to default server for {}:{}",
                    listener_host,
                    listener_port,
                );
                self.virtual_servers
                    .iter()
                    .find(|(host, port, _)| host == listener_host && *port == listener_port)
            });

        match virtual_server {
            Some((_, _, server)) => {
                debug_log!("Selected server: {}", server.server_name);
                self.route_for_server(
                    request,
                    server,
                    listener_host,
                    listener_port,
                    session_manager,
                    cgi_timeout,
                )
            }
            None => {
                debug_log!(
                    "No server configured for {}:{}",
                    listener_host,
                    listener_port
                );
                RouteResponse::Complete(HttpResponse::new(
                    400,
                    format!(
                        "No server configured for {}:{}",
                        listener_host, listener_port
                    ),
                    "text/plain".to_string(),
                ))
            }
        }
    }
    fn find_virtual_server(
        &self,
        request: &HttpRequest,
        listener_host: &str,
        listener_port: u16,
    ) -> Option<&VirtualServerRoutes> {
        let host_header = request
            .headers
            .get("host")
            .map(|h| h.as_str())
            .unwrap_or("");
        let host_header = host_header.split(':').next().unwrap_or(host_header);

        self.virtual_servers
            .iter()
            .filter(|(host, port, _)| host == listener_host && *port == listener_port)
            .find(|(_, _, server)| server.server_name == host_header)
            .or_else(|| {
                self.virtual_servers
                    .iter()
                    .find(|(host, port, _)| host == listener_host && *port == listener_port)
            })
            .map(|(_, _, server)| server)
    }

    pub fn get_max_body_size(
        &self,
        request: &HttpRequest,
        listener_host: &str,
        listener_port: u16,
    ) -> usize {
        self.find_virtual_server(request, listener_host, listener_port)
            .map(|server| server.max_body_size)
            .unwrap_or(default_max_body_size())
    }

    pub fn get_error_pages(
        &self,
        request: &HttpRequest,
        listener_host: &str,
        listener_port: u16,
    ) -> ErrorPage {
        self.find_virtual_server(request, listener_host, listener_port)
            .map(|server| server.error_pages.clone())
            .unwrap_or_else(|| ErrorPage::new(&HashMap::new()))
    }
    fn route_for_server(
        &self,
        request: &HttpRequest,
        server: &VirtualServerRoutes,
        host_header: &str,
        listener_port: u16,
        session_manager: &SessionManager,
        cgi_timeout: Duration,
    ) -> RouteResponse {
        for route in &server.routes {
            if self.matches_route(route, request) {
                debug_log!(
                    "Matched route: {} for method {}",
                    route.path,
                    request.method
                );
                if !route.methods.contains(&request.method) {
                    return RouteResponse::Complete(
                        server.error_pages.error_response(
                            405,
                            &format!("Method not allowed: {}", request.method),
                        ),
                    );
                }
                if let Some(redirect_url) = &route.redirect {
                    let mut response = HttpResponse::new(301, "", "text/plain".to_string());

                    let absolute_url = if redirect_url.starts_with("http://")
                        || redirect_url.starts_with("https://")
                    {
                        redirect_url.clone()
                    } else {
                        format!("http://{}:{}{}", host_header, listener_port, redirect_url)
                    };
                    response.headers.push(("Location", absolute_url));

                    return RouteResponse::Complete(response);
                }
                return self.handle_route(route, request, server, session_manager, cgi_timeout);
            }
        }
        RouteResponse::Complete(self.match_static_path(request, server))
    }

    fn normalize_path(&self, uri: &str) -> String {
        let request_path = uri.split('?').next().unwrap_or(uri);
        let mut normalized_path = self.url_decode(request_path);
        while normalized_path.contains("//") {
            normalized_path = normalized_path.replace("//", "/");
        }
        normalized_path
    }

    fn matches_route(&self, route: &Route, request: &HttpRequest) -> bool {
        let normalized_path = self.normalize_path(&request.uri);
        if route.path.ends_with("/*") {
            let prefix = route.path.trim_end_matches("/*");
            normalized_path.starts_with(prefix)
        } else {
            normalized_path == route.path
        }
    }

    fn handle_delete_request(
        &self,
        file_path: &Path,
        server: &VirtualServerRoutes,
    ) -> RouteResponse {
        match fs::remove_file(file_path) {
            Ok(_) => RouteResponse::Complete(HttpResponse::new(
                200,
                "File deleted",
                "text/plain".to_string(),
            )),
            Err(e) => {
                debug_log!("Error deleting file {:?}: {}", file_path, e);
                match e.kind() {
                    std::io::ErrorKind::NotFound => RouteResponse::Complete(
                        server.error_pages.error_response(404, "File Not Found"),
                    ),
                    std::io::ErrorKind::PermissionDenied => RouteResponse::Complete(
                        server.error_pages.error_response(403, "Access Forbidden"),
                    ),
                    _ => RouteResponse::Complete(
                        server
                            .error_pages
                            .error_response(500, "Internal Server Error"),
                    ),
                }
            }
        }
    }

    fn handle_route(
        &self,
        route: &Route,
        request: &HttpRequest,
        server: &VirtualServerRoutes,
        session_manager: &SessionManager,
        cgi_timeout: Duration,
    ) -> RouteResponse {
        match &route.handler {
            RouteHandler::StaticFile(file_path) => {
                if request.method == "DELETE" {
                    return self.handle_delete_request(file_path, server);
                }
                RouteResponse::Complete(self.serve_file(file_path, &server.error_pages, request))
            }
            RouteHandler::Directory(dir_path) => {
                if let Err(e) = fs::create_dir_all(dir_path) {
                    debug_log!("Failed to create directory {:?}: {}", dir_path, e);
                    return RouteResponse::Complete(
                        server
                            .error_pages
                            .error_response(500, "Failed to create upload directory"),
                    );
                }
                let route_prefix = route.path.trim_end_matches("/*");
                let normalized_path = self.normalize_path(&request.uri);

                let relative_path = normalized_path
                    .strip_prefix(route_prefix)
                    .unwrap_or("")
                    .trim_start_matches('/');
                let sanitized_path = self.sanitize_path(relative_path);
                let file_path = dir_path.join(sanitized_path);

                if request.method == "DELETE" {
                    return self.handle_delete_request(&file_path, server);
                }

                if request.method == "POST" {
                    if let Some(content_type) = request.headers.get("content-type")
                        && content_type.starts_with("multipart/form-data")
                    {
                        // Parse multipart data HERE from request.body
                        if let Some(body) = &request.body {
                            let parts = match parse_multipart_body(body, content_type) {
                                Ok(parts) => parts,
                                Err(e) => {
                                    debug_log!("Failed to parse multipart/form-data: {}", e);
                                    return RouteResponse::Complete(
                                        server.error_pages.error_response(
                                            400,
                                            &format!("Invalid multipart/form-data: {}", e),
                                        ),
                                    );
                                }
                            };

                            // Process the parsed parts
                            let mut response_body = String::new();
                            for part in parts {
                                if let Some(filename) = &part.filename {
                                    let timestamp = Utc::now().timestamp_millis();
                                    let safe_filename = self.sanitize_filename(filename);
                                    let upload_path =
                                        dir_path.join(format!("{}_{}", timestamp, safe_filename));
                                    debug_log!("Attempting to write file to {:?}", upload_path);
                                    match fs::write(&upload_path, &part.data) {
                                        Ok(_) => {
                                            debug_log!(
                                                "Successfully wrote file to {:?}",
                                                upload_path
                                            );
                                            response_body.push_str(&format!(
                                                "Uploaded file: {}\n",
                                                safe_filename
                                            ));
                                        }
                                        Err(e) => {
                                            debug_log!(
                                                "Error saving file to {:?}: {}",
                                                upload_path,
                                                e
                                            );
                                            return RouteResponse::Complete(
                                                server
                                                    .error_pages
                                                    .error_response(500, "Failed to save upload"),
                                            );
                                        }
                                    }
                                } else {
                                    response_body.push_str(&format!(
                                        "Field {}: {}\n",
                                        part.name,
                                        String::from_utf8_lossy(&part.data)
                                    ));
                                }
                            }
                            return RouteResponse::Complete(HttpResponse::new(
                                200,
                                response_body,
                                "text/plain".to_string(),
                            ));
                        } else {
                            return RouteResponse::Complete(
                                server
                                    .error_pages
                                    .error_response(400, "No body in POST request"),
                            );
                        }
                    }

                    // Handle non-multipart POST uploads
                    if let Some(body) = &request.body {
                        let timestamp = Utc::now().timestamp_millis();
                        // Map Content-Type to extension
                        let extension = get_file_extension_from_content_type(
                            request.headers.get("content-type"),
                        );
                        let upload_path =
                            dir_path.join(format!("uploaded_file_{}.{}", timestamp, extension));
                        match fs::write(&upload_path, body) {
                            Ok(_) => {
                                debug_log!("File uploaded to {:?}", upload_path);
                                return RouteResponse::Complete(HttpResponse::new(
                                    200,
                                    format!(
                                        "File uploaded successfully as uploaded_file_{}.{}",
                                        timestamp, extension
                                    ),
                                    "text/plain".to_string(),
                                ));
                            }
                            Err(e) => {
                                debug_log!("Error saving uploaded file: {}", e);
                                return RouteResponse::Complete(
                                    server
                                        .error_pages
                                        .error_response(500, "Failed to save upload"),
                                );
                            }
                        }
                    } else {
                        return RouteResponse::Complete(
                            server
                                .error_pages
                                .error_response(400, "No body in POST request"),
                        );
                    }
                }

                // Prefer route-specific `default_file`, fall back to server default
                let default_file = route
                    .default_file
                    .as_deref()
                    .or(Some(server.default_file.as_str()));
                RouteResponse::Complete(self.serve_file_or_directory(
                    &file_path,
                    route.directory_listing,
                    default_file,
                    request,
                    &server.error_pages,
                ))
            }
            RouteHandler::Function(handler) => {
                RouteResponse::Complete(handler(request, session_manager))
            }
            RouteHandler::CGI(handler) => {
                // For multipart uploads to CGI, extract the file content
                let mut modified_request = request.clone();

                if request.method == "POST"
                    && let Some(content_type) = request.headers.get("content-type")
                    && content_type.starts_with("multipart/form-data")
                    && let Some(body) = &request.body
                {
                    match parse_multipart_body(body, content_type) {
                        Ok(parts) => {
                            // Find the first file part and use its data
                            if let Some(file_part) = parts.iter().find(|p| p.filename.is_some()) {
                                debug_log!(
                                    "Extracting file '{}' ({} bytes) for CGI",
                                    file_part.filename.as_ref().unwrap(),
                                    file_part.data.len()
                                );
                                // Replace body with just the file content
                                modified_request.body = Some(file_part.data.clone());

                                // Update Content-Type to the file's type
                                if let Some(file_content_type) = &file_part.content_type {
                                    modified_request.headers.insert(
                                        "content-type".to_string(),
                                        file_content_type.clone(),
                                    );
                                }

                                // Update Content-Length
                                modified_request.headers.insert(
                                    "content-length".to_string(),
                                    file_part.data.len().to_string(),
                                );

                                // Optionally pass filename as a custom header
                                if let Some(filename) = &file_part.filename {
                                    modified_request
                                        .headers
                                        .insert("x-file-name".to_string(), filename.clone());
                                }
                            }
                        }
                        Err(e) => {
                            debug_log!("Failed to parse multipart for CGI: {}", e);
                            return RouteResponse::Complete(
                                server
                                    .error_pages
                                    .error_response(400, &format!("Invalid multipart: {}", e)),
                            );
                        }
                    }
                }
                match handler(&modified_request, &server.error_pages, cgi_timeout) {
                    Ok(process) => RouteResponse::Pending(process),
                    Err(response) => RouteResponse::Complete(response),
                }
            }
        }
    }

    fn get_content_type(&self, file_path: &Path) -> String {
        match file_path.extension().and_then(|ext| ext.to_str()) {
            Some("html") | Some("htm") => "text/html",
            Some("css") => "text/css",
            Some("js") => "application/javascript",
            Some("json") => "application/json",
            Some("png") => "image/png",
            Some("jpg") | Some("jpeg") => "image/jpeg",
            Some("gif") => "image/gif",
            Some("svg") => "image/svg+xml",
            Some("ico") => "image/x-icon",
            Some("txt") => "text/plain",
            Some("pdf") => "application/pdf",
            Some("zip") => "application/zip",
            _ => "application/octet-stream",
        }
        .to_string()
    }

    fn serve_file(
        &self,
        file_path: &Path,
        error_pages: &ErrorPage,
        request: &HttpRequest,
    ) -> HttpResponse {
        // First get file metadata to check size and pre-allocate buffer
        match fs::metadata(file_path) {
            Ok(metadata) => {
                let file_size = metadata.len() as usize;
                let content_type = self.get_content_type(file_path);
                // For files < 1MB, direct read is efficient
                // For larger files, use BufReader with capacity
                if !is_acceptable(request, &content_type) {
                    return error_pages.error_response(406, "Not Acceptable");
                }
                if file_size < 1024 * 1024 {
                    match fs::read(file_path) {
                        Ok(content) => HttpResponse::new(200, content, content_type),
                        Err(e) => self.handle_file_error(e, error_pages),
                    }
                } else {
                    // Large file: use BufReader with pre-allocated capacity
                    match fs::File::open(file_path) {
                        Ok(file) => {
                            let mut content = Vec::with_capacity(file_size);
                            let mut reader = io::BufReader::new(file);
                            match reader.read_to_end(&mut content) {
                                Ok(_) => HttpResponse::new(200, content, content_type),
                                Err(e) => self.handle_file_error(e, error_pages),
                            }
                        }
                        Err(e) => self.handle_file_error(e, error_pages),
                    }
                }
            }
            Err(e) => {
                debug_log!("Error getting metadata for {:?}: {}", file_path, e);
                self.handle_file_error(e, error_pages)
            }
        }
    }

    fn handle_file_error(&self, e: std::io::Error, error_pages: &ErrorPage) -> HttpResponse {
        debug_log!("Error reading file: {}", e);
        match e.kind() {
            std::io::ErrorKind::NotFound => error_pages.error_response(404, "File Not Found"),
            std::io::ErrorKind::PermissionDenied => {
                error_pages.error_response(403, "Access Forbidden")
            }
            _ => error_pages.error_response(500, "Internal Server Error"),
        }
    }

    fn serve_file_or_directory(
        &self,
        path: &Path,
        allow_listing: bool,
        default_file: Option<&str>,
        request: &HttpRequest,
        error_pages: &ErrorPage,
    ) -> HttpResponse {
        if path.is_file() {
            self.serve_file(path, error_pages, request)
        } else if path.is_dir() {
            if let Some(default_file) = default_file {
                let df = default_file.trim_start_matches('/');
                let index_path = path.join(df);
                if index_path.is_file() {
                    return self.serve_file(&index_path, error_pages, request);
                }
            }
            if allow_listing {
                self.serve_directory_listing(path, request, error_pages)
            } else {
                error_pages.error_response(403, "Directory Listing Forbidden")
            }
        } else {
            error_pages.error_response(404, "Path Not Found")
        }
    }

    fn serve_directory_listing(
        &self,
        dir_path: &Path,
        request: &HttpRequest,
        error_pages: &ErrorPage,
    ) -> HttpResponse {
        match fs::read_dir(dir_path) {
            Ok(entries) => {
                let mut html = String::from(
                    "<!DOCTYPE html><html><head><title>Directory Listing</title></head><body><h1>Directory Listing</h1><ul>",
                );
                let base_path = request.uri.split('?').next().unwrap_or(&request.uri);
                let base_path = if base_path.ends_with('/') {
                    base_path.to_string()
                } else {
                    format!("{}/", base_path)
                };
                html.push_str("<li><a href=\"../\">..</a></li>");
                let mut entries: Vec<_> = entries.collect();
                entries
                    .sort_by_key(|entry| entry.as_ref().map(|e| e.file_name()).unwrap_or_default());
                for entry in entries.into_iter().flatten() {
                    let name = entry.file_name();
                    let name_str = name.to_string_lossy();

                    let is_dir = entry.file_type().map(|t| t.is_dir()).unwrap_or(false);
                    let display_name = if is_dir {
                        format!("{}/", name_str)
                    } else {
                        name_str.to_string()
                    };
                    let link_path = if is_dir {
                        format!("{}{}/", base_path, self.url_encode(&name_str))
                    } else {
                        format!("{}{}", base_path, self.url_encode(&name_str))
                    };
                    html.push_str(&format!(
                        "<li><a href=\"{}\">{}</a></li>",
                        link_path,
                        html_escape(&display_name)
                    ));
                }
                html.push_str("</ul></body></html>");
                HttpResponse::new(200, html.into_bytes(), "text/html".to_string())
            }
            Err(_) => error_pages.error_response(500, "Unable to read directory"),
        }
    }

    fn url_encode(&self, input: &str) -> String {
        utf8_percent_encode(input, NON_ALPHANUMERIC).to_string()
    }

    fn url_decode(&self, input: &str) -> String {
        percent_decode_str(input).decode_utf8_lossy().to_string()
    }

    fn match_static_path(
        &self,
        request: &HttpRequest,
        server: &VirtualServerRoutes,
    ) -> HttpResponse {
        let request_path = request.uri.split('?').next().unwrap_or(&request.uri);
        let path = self.url_decode(request_path);
        let path = self.sanitize_path(&path);
        let is_valid_route = server.routes.iter().any(|route| {
            if route.path.ends_with("/*") {
                let prefix = route.path.trim_end_matches("/*");
                request_path.starts_with(prefix)
            } else {
                request_path == route.path
            }
        });

        if !is_valid_route {
            debug_log!("No route matches path: {}, returning 404", request_path);
            return server.error_pages.error_response(404, "Path Not Found");
        }

        self.serve_file_or_directory(
            Path::new(&path),
            false,
            Some(&server.default_file),
            request,
            &server.error_pages,
        )
    }

    fn sanitize_path(&self, path: &str) -> String {
        let path = path.trim_start_matches('/');
        let parts: Vec<&str> = path.split('/').collect();
        let mut clean_parts = Vec::new();
        for part in parts {
            match part {
                "" | "." => {
                    continue;
                }
                ".." => {
                    clean_parts.pop();
                }
                _ => clean_parts.push(part),
            }
        }
        clean_parts.join("/")
    }
    fn sanitize_filename(&self, filename: &str) -> String {
        let invalid_chars = ['/', '\\', ':', '*', '?', '"', '<', '>', '|'];
        let mut safe = filename.to_string();
        for c in invalid_chars.iter() {
            safe = safe.replace(*c, "_");
        }
        if safe.is_empty() {
            "unnamed".to_string()
        } else {
            safe
        }
    }
}

/// Returns true if the given content_type matches any of the Accept header values.
fn is_acceptable(request: &HttpRequest, content_type: &str) -> bool {
    let accept = request.headers.get("accept");
    if accept.is_none() {
        return true;
    }
    let accept = accept.unwrap();
    // Split on commas, trim, ignore q values for now
    for val in accept.split(',') {
        let val = val.trim().split(';').next().unwrap_or("").trim();
        if val == "*/*" || val == content_type {
            return true;
        }
        // Handle type/*
        if let Some(idx) = val.find("/*")
            && &val[..idx] == content_type.split('/').next().unwrap_or("")
        {
            return true;
        }
    }
    false
}
fn get_file_extension_from_content_type(content_type: Option<&String>) -> &'static str {
    content_type
        .and_then(|ct| match ct.as_str() {
            "image/jpeg" => Some("jpg"),
            "image/png" => Some("png"),
            "image/gif" => Some("gif"),
            "application/pdf" => Some("pdf"),
            "application/zip" => Some("zip"),
            "application/x-tar" => Some("tar"),
            "application/x-rar-compressed" => Some("rar"),
            "application/x-7z-compressed" => Some("7z"),
            "text/plain" => Some("txt"),
            "text/html" => Some("html"),
            "text/css" => Some("css"),
            "application/javascript" => Some("js"),
            "application/json" => Some("json"),
            "application/xml" => Some("xml"),
            _ => None,
        })
        .unwrap_or("bin")
}
fn html_escape(input: &str) -> String {
    let mut escaped = String::with_capacity(input.len());
    for c in input.chars() {
        match c {
            '&' => escaped.push_str("&amp;"),
            '<' => escaped.push_str("&lt;"),
            '>' => escaped.push_str("&gt;"),
            '"' => escaped.push_str("&quot;"),
            '\'' => escaped.push_str("&#39;"),
            _ => escaped.push(c),
        }
    }
    escaped
}